Implement workspace symbol query

README.md

@@ -38,8 +38,8 @@ Features in progress:
     - [x] Resolved import paths
   - [x] Rename symbols
   - [x] Multi-workspace support
-  - [x] Workspace symbol index/search
-    - [x] Editor Breadcrumbs
+  - [x] Document symbols
+  - [x] Workspace symbol query with fuzzy matching
   - [ ] Completion:
     - [x] Message and enum types
     - [x] Extendee types
@@ -48,7 +48,6 @@ Features in progress:
     - [x] Package names
     - [x] Message and field literals
     - [ ] Field literal values
-  - [ ] Type hierarchy
 - [x] Import resolution
   - [x] Local/relative paths
   - [x] Go module path lookup with inline sources

go.mod

@@ -8,8 +8,8 @@ require (
 	github.com/bufbuild/protovalidate-go v0.5.1
 	github.com/google/cel-go v0.19.0
 	github.com/kralicky/gpkg v0.0.0-20240119195700-64f32830b14f
-	github.com/kralicky/protocompile v0.0.0-20240209041722-c2e867015504
-	github.com/kralicky/tools-lite v0.0.0-20240104191314-c259ddd5a342
+	github.com/kralicky/protocompile v0.0.0-20240210012852-09c34c606fa5
+	github.com/kralicky/tools-lite v0.0.0-20240209234032-93b7eedbea2e
 	github.com/mattn/go-tty v0.0.5
 	github.com/spf13/cobra v1.8.0
 	golang.org/x/mod v0.15.0

go.sum

@@ -32,10 +32,10 @@ github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 h1:Z9n2FFNU
 github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51/go.mod h1:CzGEWj7cYgsdH8dAjBGEr58BoE7ScuLd+fwFZ44+/x8=
 github.com/kralicky/gpkg v0.0.0-20240119195700-64f32830b14f h1:MsNe8A51V+7Fu5OMXSl8SK02erPJ40vFs2zDHn89w1g=
 github.com/kralicky/gpkg v0.0.0-20240119195700-64f32830b14f/go.mod h1:vOkwMjs49XmP/7Xfo9ZL6eg2ei51lmtD/4U/Az5GTq8=
-github.com/kralicky/protocompile v0.0.0-20240209041722-c2e867015504 h1:W4mo/TNzDUS3rDdmOKBOfpFKM+PPPgVZMwV6XQoSox4=
-github.com/kralicky/protocompile v0.0.0-20240209041722-c2e867015504/go.mod h1:wlbdmEq73w8fMZJOgh0sEgjIFYaKkfrmX/DCA/JfNlc=
-github.com/kralicky/tools-lite v0.0.0-20240104191314-c259ddd5a342 h1:lZLWHXKHmOhTrs3oSZoCRtb8Y9a0mqUwCsaKut+Y1eU=
-github.com/kralicky/tools-lite v0.0.0-20240104191314-c259ddd5a342/go.mod h1:NKsdxFI6awifvNvxDwtCU1YCaKRoSSPpbHXkKOMuq24=
+github.com/kralicky/protocompile v0.0.0-20240210012852-09c34c606fa5 h1:iUY6ncakDlHsfGGs81PZphQMnQGgTrs2SseZ7BCdAcw=
+github.com/kralicky/protocompile v0.0.0-20240210012852-09c34c606fa5/go.mod h1:wlbdmEq73w8fMZJOgh0sEgjIFYaKkfrmX/DCA/JfNlc=
+github.com/kralicky/tools-lite v0.0.0-20240209234032-93b7eedbea2e h1:Mic4oZbKrGxJ6l3FmVLax0+RNOxE/J6KnnlN4BH2d58=
+github.com/kralicky/tools-lite v0.0.0-20240209234032-93b7eedbea2e/go.mod h1:NKsdxFI6awifvNvxDwtCU1YCaKRoSSPpbHXkKOMuq24=
 github.com/mattn/go-colorable v0.1.2/go.mod h1:U0ppj6V5qS13XJ6of8GYAs25YV2eR4EVcfRqFIhoBtE=
 github.com/mattn/go-colorable v0.1.4 h1:snbPLB8fVfU9iwbbo30TPtbLRzwWu6aJS6Xh4eaaviA=
 github.com/mattn/go-colorable v0.1.4/go.mod h1:U0ppj6V5qS13XJ6of8GYAs25YV2eR4EVcfRqFIhoBtE=

pkg/lsp/cache.go

@@ -306,44 +306,90 @@ func (c *Cache) FindDefinitionForTypeDescriptor(desc protoreflect.Descriptor) (p
 	}, nil
 }
 
-func (c *Cache) FindAllDescriptorsByPrefix(ctx context.Context, prefix string, filter ...func(protoreflect.Descriptor) bool) []protoreflect.Descriptor {
+type WorkspaceDescriptors interface {
+	Len() int
+	All() []protoreflect.Descriptor
+	Range(func(path string, descriptors []protoreflect.Descriptor))
+}
+
+type workspaceDescriptors struct {
+	results   [][]protoreflect.Descriptor
+	pathIndex []string
+	len       int
+}
+
+func (wd workspaceDescriptors) All() []protoreflect.Descriptor {
+	all := make([]protoreflect.Descriptor, 0, wd.len)
+	for _, item := range wd.results {
+		all = append(all, item...)
+	}
+	return all
+}
+
+func (wd workspaceDescriptors) Range(fn func(path string, descriptors []protoreflect.Descriptor)) {
+	for i, path := range wd.pathIndex {
+		fn(path, wd.results[i])
+	}
+}
+
+func (wd workspaceDescriptors) Len() int {
+	return wd.len
+}
+
+func (c *Cache) FindAllDescriptorsByPrefix(ctx context.Context, prefix string, filter ...func(protoreflect.Descriptor) bool) WorkspaceDescriptors {
 	c.resultsMu.RLock()
 	defer c.resultsMu.RUnlock()
+	return c.findAllDescriptorsByPrefixLocked(ctx, prefix, filter...)
+}
+
+func (c *Cache) findAllDescriptorsByPrefixLocked(ctx context.Context, prefix string, filter ...func(protoreflect.Descriptor) bool) WorkspaceDescriptors {
 	eg, ctx := errgroup.WithContext(ctx)
-	resultsByPackage := make([][]protoreflect.Descriptor, len(c.results))
+	eg.SetLimit(runtime.NumCPU())
+	descriptorsByFile := make([][]protoreflect.Descriptor, len(c.results))
+	pathIndex := make([]string, len(c.results))
 	for i, res := range c.results {
-		i, res := i, res
+		pathIndex[i] = res.Path()
 		pkg := res.Package()
 		if pkg == "" {
 			// skip searching in files that do not have a package name
 			continue
 		}
 		eg.Go(func() (err error) {
-			resultsByPackage[i], err = res.(linker.Result).FindDescriptorsByPrefix(ctx, string(pkg.Append(protoreflect.Name(prefix))), filter...)
+			descriptorsByFile[i], err = res.(linker.Result).FindDescriptorsByPrefix(ctx, string(pkg.Append(protoreflect.Name(prefix))), filter...)
 			return
 		})
 	}
 	eg.Wait()
-	combined := make([]protoreflect.Descriptor, 0, len(c.results))
-	for _, results := range resultsByPackage {
-		combined = append(combined, results...)
+	totalLen := 0
+	for _, item := range descriptorsByFile {
+		totalLen += len(item)
+	}
+	return workspaceDescriptors{
+		results:   descriptorsByFile,
+		pathIndex: pathIndex,
+		len:       totalLen,
 	}
-	return combined
 }
 
 // Like FindAllDescriptorsByPrefix, but assumes a fully qualified prefix with
 // package name.
-func (c *Cache) FindAllDescriptorsByQualifiedPrefix(ctx context.Context, prefix string, filter ...func(protoreflect.Descriptor) bool) []protoreflect.Descriptor {
+func (c *Cache) FindAllDescriptorsByQualifiedPrefix(ctx context.Context, prefix string, filter ...func(protoreflect.Descriptor) bool) WorkspaceDescriptors {
 	c.resultsMu.RLock()
 	defer c.resultsMu.RUnlock()
+	return c.findAllDescriptorsByQualifiedPrefixLocked(ctx, prefix, filter...)
+}
+
+func (c *Cache) findAllDescriptorsByQualifiedPrefixLocked(ctx context.Context, prefix string, filter ...func(protoreflect.Descriptor) bool) WorkspaceDescriptors {
 	eg, ctx := errgroup.WithContext(ctx)
+	eg.SetLimit(runtime.NumCPU())
 	resultsByPackage := make([][]protoreflect.Descriptor, len(c.results))
+	pathIndex := make([]string, len(c.results))
 	isFullyQualified := strings.HasPrefix(prefix, ".")
 	if isFullyQualified {
 		prefix = prefix[1:]
 	}
 	for i, res := range c.results {
-		i, res := i, res
+		pathIndex[i] = res.Path()
 		pkg := res.Package()
 		if pkg == "" {
 			// skip searching in files that do not have a package name
@@ -359,11 +405,34 @@ func (c *Cache) FindAllDescriptorsByQualifiedPrefix(ctx context.Context, prefix
 		})
 	}
 	eg.Wait()
-	combined := make([]protoreflect.Descriptor, 0, len(c.results))
-	for _, results := range resultsByPackage {
-		combined = append(combined, results...)
+	totalLen := 0
+	for _, item := range resultsByPackage {
+		totalLen += len(item)
+	}
+	return workspaceDescriptors{
+		results:   resultsByPackage,
+		pathIndex: pathIndex,
+		len:       totalLen,
+	}
+}
+
+// RangeAllDescriptors calls fn for each descriptor in the cache, possibly in
+// a separate goroutine for each file.
+func (c *Cache) RangeAllDescriptors(ctx context.Context, fn func(protoreflect.Descriptor) bool) error {
+	c.resultsMu.RLock()
+	defer c.resultsMu.RUnlock()
+	return c.rangeAllDescriptorsLocked(ctx, fn)
+}
+
+func (c *Cache) rangeAllDescriptorsLocked(ctx context.Context, fn func(protoreflect.Descriptor) bool) error {
+	eg, ctx := errgroup.WithContext(ctx)
+	eg.SetLimit(runtime.NumCPU())
+	for _, res := range c.results {
+		eg.Go(func() (err error) {
+			return res.(linker.Result).RangeDescriptors(ctx, fn)
+		})
 	}
-	return combined
+	return eg.Wait()
 }
 
 func (c *Cache) FindImportPathsByPrefix(ctx context.Context, prefix string) map[protocol.DocumentURI]string {

pkg/lsp/completion.go

@@ -965,9 +965,9 @@ func completeTypeNames(cache *Cache, partialName, partialNameSuffix string, link
 		return false
 	}
 	if strings.Contains(partialName, ".") {
-		candidates = cache.FindAllDescriptorsByQualifiedPrefix(context.TODO(), partialName, filter)
+		candidates = cache.FindAllDescriptorsByQualifiedPrefix(context.TODO(), partialName, filter).All()
 	} else {
-		candidates = cache.FindAllDescriptorsByPrefix(context.TODO(), partialName, filter)
+		candidates = cache.FindAllDescriptorsByPrefix(context.TODO(), partialName, filter).All()
 	}
 	return completeTypeNamesFromList(candidates, partialName, partialNameSuffix, linkRes, scope, pos)
 }
@@ -983,9 +983,9 @@ func completeExtendeeTypeNames(cache *Cache, partialName, partialNameSuffix stri
 			return false
 		}
 		if strings.Contains(partialName, ".") {
-			candidates = append(candidates, cache.FindAllDescriptorsByQualifiedPrefix(context.TODO(), partialName, filter)...)
+			candidates = append(candidates, cache.FindAllDescriptorsByQualifiedPrefix(context.TODO(), partialName, filter).All()...)
 		} else {
-			candidates = append(candidates, cache.FindAllDescriptorsByPrefix(context.TODO(), partialName, filter)...)
+			candidates = append(candidates, cache.FindAllDescriptorsByPrefix(context.TODO(), partialName, filter).All()...)
 		}
 	}
 	return completeTypeNamesFromList(candidates, partialName, partialNameSuffix, linkRes, scope, pos)

pkg/lsp/refactor.go

@@ -59,9 +59,9 @@ func RefactorUndeclaredName(ctx context.Context, cache *Cache, uri protocol.Docu
 	}
 	var candidates []protoreflect.Descriptor
 	if strings.Contains(name, ".") {
-		candidates = cache.FindAllDescriptorsByQualifiedPrefix(ctx, name, filter)
+		candidates = cache.FindAllDescriptorsByQualifiedPrefix(ctx, name, filter).All()
 	} else {
-		candidates = cache.FindAllDescriptorsByQualifiedPrefix(ctx, string(linkRes.Package().Append(protoreflect.Name(name))), filter)
+		candidates = cache.FindAllDescriptorsByQualifiedPrefix(ctx, string(linkRes.Package().Append(protoreflect.Name(name))), filter).All()
 	}
 
 	var matches []protoreflect.Descriptor

pkg/lsp/server.go

@@ -203,9 +203,9 @@ func (s *Server) Initialize(ctx context.Context, params *protocol.ParamInitializ
 			},
 			// DeclarationProvider: &protocol.Or_ServerCapabilities_declarationProvider{Value: true},
 			// TypeDefinitionProvider: true,
-			ReferencesProvider: &protocol.Or_ServerCapabilities_referencesProvider{Value: true},
-			// WorkspaceSymbolProvider: &protocol.Or_ServerCapabilities_workspaceSymbolProvider{Value: true},
-			DefinitionProvider: &protocol.Or_ServerCapabilities_definitionProvider{Value: true},
+			ReferencesProvider:      &protocol.Or_ServerCapabilities_referencesProvider{Value: true},
+			WorkspaceSymbolProvider: &protocol.Or_ServerCapabilities_workspaceSymbolProvider{Value: true},
+			DefinitionProvider:      &protocol.Or_ServerCapabilities_definitionProvider{Value: true},
 
 			SemanticTokensProvider: &protocol.SemanticTokensOptions{
 				Legend: protocol.SemanticTokensLegend{
@@ -562,7 +562,7 @@ func (s *Server) DocumentSymbol(ctx context.Context, params *protocol.DocumentSy
 	if err != nil {
 		return nil, err
 	}
-	symbols, err := c.DocumentSymbolsForFile(params.TextDocument)
+	symbols, err := c.DocumentSymbolsForFile(params.TextDocument.URI)
 	if err != nil {
 		return nil, err
 	}
@@ -809,6 +809,15 @@ func (s *Server) CodeLens(ctx context.Context, params *protocol.CodeLensParams)
 	return c.ComputeCodeLens(params.TextDocument.URI)
 }
 
+// Symbol implements protocol.Server.
+func (s *Server) Symbol(ctx context.Context, params *protocol.WorkspaceSymbolParams) ([]protocol.SymbolInformation, error) {
+	var symbolInfos []protocol.SymbolInformation
+	for _, c := range s.caches {
+		symbolInfos = append(symbolInfos, c.QueryWorkspaceSymbols(ctx, params.Query)...)
+	}
+	return symbolInfos, nil
+}
+
 // =====================
 // Unimplemented Methods
 // =====================
@@ -837,11 +846,6 @@ func (*Server) Supertypes(context.Context, *protocol.TypeHierarchySupertypesPara
 	return nil, notImplemented("Supertypes")
 }
 
-// Symbol implements protocol.Server.
-func (*Server) Symbol(context.Context, *protocol.WorkspaceSymbolParams) ([]protocol.SymbolInformation, error) {
-	return nil, notImplemented("Symbol")
-}
-
 // TypeDefinition implements protocol.Server.
 func (*Server) TypeDefinition(context.Context, *protocol.TypeDefinitionParams) ([]protocol.Location, error) {
 	return nil, notImplemented("TypeDefinition")