lsp: add inlay hint support (#621)

This PR adds support for inlay hints of function arguments.
Similar to the hover feature, only built-in functions are
supported at this point.

Also did some minor refactoring in the hover function to
have it use the same method for collecting built-in functions
from a module as the inlay hint system does.

Signed-off-by: Anders Eknert <anders@styra.com>

internal/lsp/hover.go

@@ -150,23 +150,16 @@ func updateBuiltinPositions(cache *Cache, uri string) error {
 
 	builtinsOnLine := map[uint][]BuiltinPosition{}
 
-	ast.WalkTerms(module, func(t *ast.Term) bool {
-		if call, ok := t.Value.(ast.Call); ok {
-			name := call[0].Value.String()
-			line := uint(call[0].Location.Row)
-
-			if b, ok := builtins[name]; ok {
-				builtinsOnLine[line] = append(builtinsOnLine[line], BuiltinPosition{
-					Builtin: b,
-					Line:    line,
-					Start:   uint(call[0].Location.Col),
-					End:     uint(call[0].Location.Col + len(name)),
-				})
-			}
-		}
-
-		return false
-	})
+	for _, call := range AllBuiltinCalls(module) {
+		line := uint(call.Location.Row)
+
+		builtinsOnLine[line] = append(builtinsOnLine[line], BuiltinPosition{
+			Builtin: call.Builtin,
+			Line:    line,
+			Start:   uint(call.Location.Col),
+			End:     uint(call.Location.Col + len(call.Builtin.Name)),
+		})
+	}
 
 	cache.SetBuiltinPositions(uri, builtinsOnLine)
 

internal/lsp/inlayhint.go

@@ -0,0 +1,46 @@
+package lsp
+
+import (
+	"fmt"
+
+	"github.com/open-policy-agent/opa/ast"
+	"github.com/open-policy-agent/opa/types"
+)
+
+func createInlayTooltip(named *types.NamedType) string {
+	if named.Descr == "" {
+		return fmt.Sprintf("Type: `%s`", named.Type.String())
+	}
+
+	return fmt.Sprintf("%s\n\nType: `%s`", named.Descr, named.Type.String())
+}
+
+func getInlayHints(module *ast.Module) []InlayHint {
+	inlayHints := make([]InlayHint, 0)
+
+	for _, call := range AllBuiltinCalls(module) {
+		for i, arg := range call.Builtin.Decl.NamedFuncArgs().Args {
+			if len(call.Args) <= i {
+				// avoid panic if provided a builtin function where the args
+				// have yet to be provided, like if the user types `split()`
+				continue
+			}
+
+			if named, ok := arg.(*types.NamedType); ok {
+				inlayHints = append(inlayHints, InlayHint{
+					Position:     positionFromLocation(call.Args[i].Location),
+					Label:        named.Name + ":",
+					Kind:         2,
+					PaddingLeft:  false,
+					PaddingRight: true,
+					Tooltip: MarkupContent{
+						Kind:  "markdown",
+						Value: createInlayTooltip(named),
+					},
+				})
+			}
+		}
+	}
+
+	return inlayHints
+}

internal/lsp/inlayhint_test.go

@@ -0,0 +1,86 @@
+package lsp
+
+import (
+	"testing"
+
+	"github.com/open-policy-agent/opa/ast"
+)
+
+// A function call may either be represented as an ast.Call.
+func TestGetInlayHintsAstCall(t *testing.T) {
+	t.Parallel()
+
+	policy := `package p
+
+	r := json.filter({}, [])`
+
+	module := ast.MustParseModule(policy)
+
+	inlayHints := getInlayHints(module)
+
+	if len(inlayHints) != 2 {
+		t.Errorf("Expected 2 inlay hints, got %d", len(inlayHints))
+	}
+
+	if inlayHints[0].Label != "object:" {
+		t.Errorf("Expected label to be 'object:', got %s", inlayHints[0].Label)
+	}
+
+	if inlayHints[0].Position.Line != 2 && inlayHints[0].Position.Character != 18 {
+		t.Errorf("Expected line 2, character 18, got %d, %d",
+			inlayHints[0].Position.Line, inlayHints[0].Position.Character)
+	}
+
+	if inlayHints[0].Tooltip.Value != "Type: `object[any: any]`" {
+		t.Errorf("Expected tooltip to be 'Type: `object[any: any]`, got %s", inlayHints[0].Tooltip.Value)
+	}
+
+	if inlayHints[1].Label != "paths:" {
+		t.Errorf("Expected label to be 'paths:', got %s", inlayHints[1].Label)
+	}
+
+	if inlayHints[1].Position.Line != 2 && inlayHints[1].Position.Character != 22 {
+		t.Errorf("Expected line 2, character 22, got %d, %d",
+			inlayHints[1].Position.Line, inlayHints[1].Position.Character)
+	}
+
+	if inlayHints[1].Tooltip.Value != "JSON string paths\n\nType: `any<array[any<string, array[any]>],"+
+		" set[any<string, array[any]>]>`" {
+		t.Errorf("Expected tooltip to be 'JSON string paths\n\nType: `any<array[any<string, array[any]>], "+
+			"set[any<string, array[any]>]>`, got %s", inlayHints[1].Tooltip.Value)
+	}
+}
+
+// Or a function call may be represented as the terms of an ast.Expr.
+func TestGetInlayHintsAstTerms(t *testing.T) {
+	t.Parallel()
+
+	policy := `package p
+
+	import rego.v1
+
+	allow if {
+		is_string("yes")
+	}`
+
+	module := ast.MustParseModule(policy)
+
+	inlayHints := getInlayHints(module)
+
+	if len(inlayHints) != 1 {
+		t.Errorf("Expected 1 inlay hints, got %d", len(inlayHints))
+	}
+
+	if inlayHints[0].Label != "x:" {
+		t.Errorf("Expected label to be 'x:', got %s", inlayHints[0].Label)
+	}
+
+	if inlayHints[0].Position.Line != 5 && inlayHints[0].Position.Character != 12 {
+		t.Errorf("Expected line 5, character 12, got %d, %d",
+			inlayHints[0].Position.Line, inlayHints[0].Position.Character)
+	}
+
+	if inlayHints[0].Tooltip.Value != "Type: `any`" {
+		t.Errorf("Expected tooltip to be 'Type: `any`, got %s", inlayHints[0].Tooltip.Value)
+	}
+}

internal/lsp/messages.go

@@ -80,6 +80,7 @@ type ServerCapabilities struct {
 	TextDocumentSyncOptions TextDocumentSyncOptions `json:"textDocumentSync"`
 	DiagnosticProvider      DiagnosticOptions       `json:"diagnosticProvider"`
 	Workspace               WorkspaceOptions        `json:"workspace"`
+	InlayHintProvider       InlayHintOptions        `json:"inlayHintProvider"`
 	HoverProvider           bool                    `json:"hoverProvider"`
 }
 
@@ -111,6 +112,24 @@ type DiagnosticOptions struct {
 	WorkspaceDiagnostics  bool   `json:"workspaceDiagnostics"`
 }
 
+type InlayHintOptions struct {
+	ResolveProvider bool `json:"resolveProvider"`
+}
+
+type InlayHint struct {
+	Position     Position      `json:"position"`
+	Label        string        `json:"label"`
+	Kind         uint          `json:"kind"`
+	PaddingLeft  bool          `json:"paddingLeft"`
+	PaddingRight bool          `json:"paddingRight"`
+	Tooltip      MarkupContent `json:"tooltip"`
+}
+
+type TextDocumentInlayHintParams struct {
+	TextDocument TextDocumentIdentifier `json:"textDocument"`
+	Range        Range                  `json:"range"`
+}
+
 type TextDocumentSyncOptions struct {
 	OpenClose bool `json:"openClose"`
 	Change    uint `json:"change"`

internal/lsp/rego.go

@@ -0,0 +1,60 @@
+package lsp
+
+import "github.com/open-policy-agent/opa/ast"
+
+type BuiltInCall struct {
+	Builtin  *ast.Builtin
+	Location *ast.Location
+	Args     []*ast.Term
+}
+
+func positionFromLocation(loc *ast.Location) Position {
+	return Position{
+		Line:      uint(loc.Row - 1),
+		Character: uint(loc.Col - 1),
+	}
+}
+
+// AllBuiltinCalls returns all built-in calls in the module, excluding operators
+// and any other function identified by an infix.
+func AllBuiltinCalls(module *ast.Module) []BuiltInCall {
+	builtinCalls := make([]BuiltInCall, 0)
+
+	callVisitor := ast.NewGenericVisitor(func(x interface{}) bool {
+		var terms []*ast.Term
+
+		switch node := x.(type) {
+		case ast.Call:
+			terms = node
+		case *ast.Expr:
+			if call, ok := node.Terms.([]*ast.Term); ok {
+				terms = call
+			}
+		default:
+			return false
+		}
+
+		if len(terms) == 0 {
+			return false
+		}
+
+		if b, ok := builtins[terms[0].Value.String()]; ok {
+			// Exclude operators and similar builtins
+			if b.Infix != "" {
+				return false
+			}
+
+			builtinCalls = append(builtinCalls, BuiltInCall{
+				Builtin:  b,
+				Location: terms[0].Location,
+				Args:     terms[1:],
+			})
+		}
+
+		return false
+	})
+
+	callVisitor.Walk(module)
+
+	return builtinCalls
+}

internal/lsp/server.go

@@ -100,6 +100,8 @@ func (l *LanguageServer) Handle(
 		return l.handleTextDocumentDidChange(ctx, conn, req)
 	case "textDocument/hover":
 		return l.handleHover(ctx, conn, req)
+	case "textDocument/inlayHint":
+		return l.handleInlayHint(ctx, conn, req)
 	case "workspace/didChangeWatchedFiles":
 		return l.handleWorkspaceDidChangeWatchedFiles(ctx, conn, req)
 	case "workspace/diagnostic":
@@ -355,6 +357,26 @@ func (l *LanguageServer) handleHover(
 	return struct{}{}, nil
 }
 
+func (l *LanguageServer) handleInlayHint(
+	_ context.Context,
+	_ *jsonrpc2.Conn,
+	req *jsonrpc2.Request,
+) (result any, err error) {
+	var params TextDocumentInlayHintParams
+	if err := json.Unmarshal(*req.Params, &params); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal params: %w", err)
+	}
+
+	module, ok := l.cache.GetModule(params.TextDocument.URI)
+	if !ok {
+		l.log(fmt.Sprintf("failed to get module for uri %q", params.TextDocument.URI))
+
+		return []InlayHint{}, nil
+	}
+
+	return getInlayHints(module), nil
+}
+
 func (l *LanguageServer) handleTextDocumentDidOpen(
 	_ context.Context,
 	_ *jsonrpc2.Conn,
@@ -570,6 +592,9 @@ func (l *LanguageServer) handleInitialize(
 					},
 				},
 			},
+			InlayHintProvider: InlayHintOptions{
+				ResolveProvider: false,
+			},
 			HoverProvider: true,
 		},
 	}