lsp: implement workspace/symbol (#673)

This was rather trivial with document symbols supported. It's now
made possible to search for symbols across the whole workspace, which
is rather nice!

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

internal/lsp/documentsymbol.go

@@ -267,3 +267,28 @@ func isConstant(rule *ast.Rule) bool {
 		rule.Body.Equal(ast.NewBody(ast.NewExpr(ast.BooleanTerm(true)))) &&
 		rule.Else == nil
 }
+
+func toWorkspaceSymbol(docSym types.DocumentSymbol, docURL string) types.WorkspaceSymbol {
+	return types.WorkspaceSymbol{
+		Name: docSym.Name,
+		Kind: docSym.Kind,
+		Location: types.Location{
+			URI:   docURL,
+			Range: docSym.Range,
+		},
+	}
+}
+
+func toWorkspaceSymbols(docSym []types.DocumentSymbol, docURL string, symbols *[]types.WorkspaceSymbol) {
+	for _, sym := range docSym {
+		// Only include the "main" symbol for incremental rules and functions
+		// as numeric items isn't very useful in the workspace symbol list.
+		if !strings.HasPrefix(sym.Name, "#") {
+			*symbols = append(*symbols, toWorkspaceSymbol(sym, docURL))
+
+			if sym.Children != nil {
+				toWorkspaceSymbols(*sym.Children, docURL, symbols)
+			}
+		}
+	}
+}

internal/lsp/server.go

@@ -135,6 +135,8 @@ func (l *LanguageServer) Handle(
 		return l.handleWorkspaceDidCreateFiles(ctx, conn, req)
 	case "workspace/executeCommand":
 		return l.handleWorkspaceExecuteCommand(ctx, conn, req)
+	case "workspace/symbol":
+		return l.handleWorkspaceSymbol(ctx, conn, req)
 	case "shutdown":
 		err = conn.Close()
 		if err != nil {
@@ -611,6 +613,37 @@ func (l *LanguageServer) handleTextDocumentInlayHint(
 	return inlayHints, nil
 }
 
+func (l *LanguageServer) handleWorkspaceSymbol(
+	_ context.Context,
+	_ *jsonrpc2.Conn,
+	req *jsonrpc2.Request,
+) (result any, err error) {
+	var params types.WorkspaceSymbolParams
+	if err := json.Unmarshal(*req.Params, &params); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal params: %w", err)
+	}
+
+	symbols := make([]types.WorkspaceSymbol, 0)
+	contents := l.cache.GetAllFiles()
+
+	// Note: currently ignoring params.Query, as the client seems to do a good
+	// job of filtering anyway, and that would merely be an optimization here.
+	// But perhaps a good one to do at some point, and I'm not sure all clients
+	// do this filtering.
+
+	for moduleURL, module := range l.cache.GetAllModules() {
+		content := contents[moduleURL]
+		docSyms := documentSymbols(content, module)
+		wrkSyms := make([]types.WorkspaceSymbol, 0)
+
+		toWorkspaceSymbols(docSyms, moduleURL, &wrkSyms)
+
+		symbols = append(symbols, wrkSyms...)
+	}
+
+	return symbols, nil
+}
+
 func (l *LanguageServer) handleTextDocumentDefinition(
 	_ context.Context,
 	_ *jsonrpc2.Conn,
@@ -966,6 +999,7 @@ func (l *LanguageServer) handleInitialize(
 			FoldingRangeProvider:       true,
 			DefinitionProvider:         true,
 			DocumentSymbolProvider:     true,
+			WorkspaceSymbolProvider:    true,
 		},
 	}
 

internal/lsp/types/types.go

@@ -89,6 +89,7 @@ type ServerCapabilities struct {
 	DocumentFormattingProvider bool                    `json:"documentFormattingProvider"`
 	FoldingRangeProvider       bool                    `json:"foldingRangeProvider"`
 	DocumentSymbolProvider     bool                    `json:"documentSymbolProvider"`
+	WorkspaceSymbolProvider    bool                    `json:"workspaceSymbolProvider"`
 	DefinitionProvider         bool                    `json:"definitionProvider"`
 }
 
@@ -173,6 +174,17 @@ type DocumentSymbol struct {
 	Children       *[]DocumentSymbol  `json:"children,omitempty"`
 }
 
+type WorkspaceSymbolParams struct {
+	Query string `json:"query"`
+}
+
+type WorkspaceSymbol struct {
+	Name          string             `json:"name"`
+	Kind          symbols.SymbolKind `json:"kind"`
+	Location      Location           `json:"location"`
+	ContainerName *string            `json:"containerName,omitempty"`
+}
+
 type FoldingRangeParams struct {
 	TextDocument TextDocumentIdentifier `json:"textDocument"`
 }