Fix LSP references for dependencies

CHANGELOG.md

@@ -5,6 +5,8 @@
 - Fix `buf format` dropping comments next to commas or semicolons in message literals.
 - Fix compilation failing to resolve symbols re-exported through `import public` when the
   re-exporting file also reaches those symbols through a non-public import.
+- Fix LSP finding only a subset of references to symbols declared in
+  dependencies, including well-known types.
 
 ## [v1.72.0] - 2026-07-17
 

private/buf/buflsp/buflsp.go

@@ -100,6 +100,7 @@ func Serve(
 	}
 	lsp.fileManager = newFileManager(lsp)
 	lsp.workspaceManager = newWorkspaceManager(lsp)
+	lsp.referenceIndex = newReferenceIndex()
 	lsp.bufYAMLManager = newBufYAMLManager(lsp)
 	lsp.bufGenYAMLManager = newBufGenYAMLManager(lsp)
 	lsp.bufPolicyYAMLManager = newBufPolicyYAMLManager()
@@ -137,6 +138,7 @@ type lsp struct {
 	wasmRuntime          wasm.Runtime
 	fileManager          *fileManager
 	workspaceManager     *workspaceManager
+	referenceIndex       *referenceIndex
 	bufYAMLManager       *bufYAMLManager
 	bufGenYAMLManager    *bufGenYAMLManager
 	bufPolicyYAMLManager *bufPolicyYAMLManager

private/buf/buflsp/file.go

@@ -63,7 +63,6 @@ type file struct {
 
 	ir                   *ir.File
 	referenceableSymbols map[ir.FullName]*symbol
-	referenceSymbols     []*symbol
 	symbols              []*symbol
 	irReport             *report.Report        // IR diagnostic report for code actions
 	diagnostics          []protocol.Diagnostic // Converted LSP diagnostics
@@ -98,6 +97,8 @@ func (f *file) Reset(ctx context.Context) {
 		f.workspace.Release()
 		f.workspace = nil
 	}
+	// Drop this file's references from the index before the file is zeroed.
+	f.lsp.referenceIndex.RemoveFile(f.uri)
 	// Evict the query key if there is a query cached on the file. We cache the [queries.File]
 	// query since this allows the executor to evict all dependent queries, e.g. AST and IR.
 	f.lsp.queryExecutor.Evict(f.queryFileKeys()...)
@@ -354,7 +355,6 @@ func (f *file) IndexSymbols(ctx context.Context) {
 	// Throw away all the old symbols and rebuild symbols unconditionally. This is because if
 	// this file depends on a file that has since been modified, we may need to update references.
 	f.symbols = nil
-	f.referenceSymbols = nil
 	f.referenceableSymbols = make(map[ir.FullName]*symbol)
 
 	// Process all imports as symbols
@@ -365,7 +365,6 @@ func (f *file) IndexSymbols(ctx context.Context) {
 	resolved, unresolved := f.indexSymbols()
 	f.symbols = append(f.symbols, resolved...)
 	f.symbols = append(f.symbols, unresolved...)
-	f.referenceSymbols = append(f.referenceSymbols, unresolved...)
 
 	// Index all referenceable symbols
 	for _, sym := range resolved {
@@ -376,47 +375,29 @@ func (f *file) IndexSymbols(ctx context.Context) {
 		f.referenceableSymbols[sym.ir.FullName()] = sym
 	}
 
-	// TODO: this could use a refactor, probably.
-	//
-	// Resolve all unresolved symbols from this file
+	// Resolve all unresolved symbols from this file, and record the references they make in
+	// the reference index. References are keyed by definition site, so recording does not
+	// depend on the order files are indexed in.
+	var references map[referenceKey][]*symbol
+	addReference := func(def ast.DeclDef, fullName ir.FullName, sym *symbol) {
+		key, ok := newReferenceKey(def, fullName)
+		if !ok {
+			return
+		}
+		if references == nil {
+			references = make(map[referenceKey][]*symbol)
+		}
+		references[key] = append(references[key], sym)
+	}
 	for _, sym := range unresolved {
 		switch kind := sym.kind.(type) {
 		case *reference:
-			def := f.resolveASTDefinition(kind.def, kind.fullName)
-			sym.def = def
-			if def == nil {
-				// In the case where the symbol is not resolved, we continue
-				continue
-			}
-			referenceable, ok := def.kind.(*referenceable)
-			if !ok {
-				// This shouldn't happen, logging a warning
-				f.lsp.logger.Warn(
-					"found non-referenceable symbol in index",
-					slog.String("file", f.uri.Filename()),
-					slog.Any("symbol", def),
-				)
-				continue
-			}
-			referenceable.references = append(referenceable.references, sym)
+			sym.def = f.resolveASTDefinition(kind.def, kind.fullName)
+			addReference(kind.def, kind.fullName, sym)
 		case *option:
-			def := f.resolveASTDefinition(kind.def, kind.defFullName)
-			sym.def = def
-			if def != nil {
-				referenceable, ok := def.kind.(*referenceable)
-				if !ok {
-					// This shouldn't happen, logging a warning
-					f.lsp.logger.Warn(
-						"found non-referenceable symbol in index",
-						slog.String("file", f.uri.Filename()),
-						slog.Any("symbol", def),
-					)
-				} else {
-					referenceable.references = append(referenceable.references, sym)
-				}
-			}
-			typeDef := f.resolveASTDefinition(kind.typeDef, kind.typeDefFullName)
-			sym.typeDef = typeDef
+			sym.def = f.resolveASTDefinition(kind.def, kind.defFullName)
+			addReference(kind.def, kind.defFullName, sym)
+			sym.typeDef = f.resolveASTDefinition(kind.typeDef, kind.typeDefFullName)
 		default:
 			// This shouldn't happen, logging a warning
 			f.lsp.logger.Warn(
@@ -426,58 +407,7 @@ func (f *file) IndexSymbols(ctx context.Context) {
 			)
 		}
 	}
-
-	// Resolve all references outside of this file to symbols in this file
-	for _, file := range f.workspace.PathToFile() {
-		if f == file {
-			continue // ignore self
-		}
-		for _, sym := range file.referenceSymbols {
-			var fullName ir.FullName
-			switch kind := sym.kind.(type) {
-			case *reference:
-				if kind.def.Span().Path() != f.objectInfo.LocalPath() {
-					continue
-				}
-				fullName = kind.fullName
-			case *option:
-				if kind.def.Span().Path() != f.objectInfo.LocalPath() {
-					continue
-				}
-				fullName = kind.defFullName
-			default:
-				// This shouldn't happen, logging a warning
-				f.lsp.logger.Warn(
-					"found unresolved non-reference and non-option symbol",
-					slog.String("file", f.uri.Filename()),
-					slog.Any("symbol", sym),
-				)
-				continue
-			}
-			def, ok := f.referenceableSymbols[fullName]
-			if !ok {
-				// This shouldn't happen, if a symbol is pointing at this file, all definitions
-				// should be resolved, logging a warning
-				f.lsp.logger.Warn(
-					"found reference to unknown symbol",
-					slog.String("file", f.uri.Filename()),
-					slog.Any("reference", sym),
-				)
-				continue
-			}
-			referenceable, ok := def.kind.(*referenceable)
-			if !ok {
-				// This shouldn't happen, logging a warning
-				f.lsp.logger.Warn(
-					"found non-referenceable symbol in index",
-					slog.String("file", f.uri.Filename()),
-					slog.Any("symbol", def),
-				)
-				continue
-			}
-			referenceable.references = append(referenceable.references, sym)
-		}
-	}
+	f.lsp.referenceIndex.SetFile(f.uri, references)
 
 	// Finally, sort the symbols in position order, with shorter symbols sorting smaller.
 	slices.SortFunc(f.symbols, func(s1, s2 *symbol) int {

private/buf/buflsp/reference_index.go

@@ -0,0 +1,141 @@
+// Copyright 2020-2026 Buf Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// This file defines the reverse index from symbol definitions to their references.
+
+package buflsp
+
+import (
+	"cmp"
+	"slices"
+
+	"github.com/bufbuild/protocompile/experimental/ast"
+	"github.com/bufbuild/protocompile/experimental/ir"
+	"go.lsp.dev/protocol"
+)
+
+// referenceKey identifies a symbol definition site: the path of the declaring file as it
+// appears in source spans, plus the symbol's fully-qualified name.
+type referenceKey struct {
+	path     string
+	fullName ir.FullName
+}
+
+// newReferenceKey builds the key for the definition an unresolved symbol names. Returns
+// false if the definition has no span to take a path from, such as a zero [ast.DeclDef].
+func newReferenceKey(def ast.DeclDef, fullName ir.FullName) (referenceKey, bool) {
+	path := def.Span().Path()
+	if path == "" || fullName == "" {
+		return referenceKey{}, false
+	}
+	return referenceKey{path: path, fullName: fullName}, true
+}
+
+// newReferenceKeyForDeclaration builds the key for a declaration symbol. This is the only
+// other way to construct a key, and derives the path from the same source as spans do, so
+// lookups cannot diverge from what [newReferenceKey] indexed.
+func newReferenceKeyForDeclaration(declaration *symbol) (referenceKey, bool) {
+	if declaration.file == nil || declaration.file.file == nil || declaration.ir.IsZero() {
+		return referenceKey{}, false
+	}
+	return referenceKey{
+		path:     declaration.file.file.Path(),
+		fullName: declaration.ir.FullName(),
+	}, true
+}
+
+// referenceIndex is a reverse index from a symbol definition to the symbols referencing it,
+// across all workspaces. References are grouped by the file containing them, so re-indexing
+// a file replaces only that file's contribution.
+//
+// The index is not safe for concurrent use; it is protected by the [lsp] lock.
+type referenceIndex struct {
+	keyToFileRefs map[referenceKey]map[protocol.URI][]*symbol
+	// uriToKeys records which definitions each file contributes references to, so that a
+	// file's contribution can be removed without scanning the whole index.
+	uriToKeys map[protocol.URI][]referenceKey
+}
+
+// newReferenceIndex creates a new reference index.
+func newReferenceIndex() *referenceIndex {
+	return &referenceIndex{
+		keyToFileRefs: make(map[referenceKey]map[protocol.URI][]*symbol),
+		uriToKeys:     make(map[protocol.URI][]referenceKey),
+	}
+}
+
+// SetFile replaces all references contributed by the file at the given URI.
+func (i *referenceIndex) SetFile(uri protocol.URI, references map[referenceKey][]*symbol) {
+	i.RemoveFile(uri)
+	if len(references) == 0 {
+		return
+	}
+	keys := make([]referenceKey, 0, len(references))
+	for key, symbols := range references {
+		fileRefs, ok := i.keyToFileRefs[key]
+		if !ok {
+			fileRefs = make(map[protocol.URI][]*symbol, 1)
+			i.keyToFileRefs[key] = fileRefs
+		}
+		fileRefs[uri] = symbols
+		keys = append(keys, key)
+	}
+	i.uriToKeys[uri] = keys
+}
+
+// RemoveFile drops all references contributed by the file at the given URI.
+//
+// This must be called when a file is evicted. The index holds symbols pointing back at
+// their file, and an evicted file is zeroed, so entries left behind would resolve to
+// locations with an empty URI.
+func (i *referenceIndex) RemoveFile(uri protocol.URI) {
+	for _, key := range i.uriToKeys[uri] {
+		fileRefs, ok := i.keyToFileRefs[key]
+		if !ok {
+			continue
+		}
+		delete(fileRefs, uri)
+		if len(fileRefs) == 0 {
+			delete(i.keyToFileRefs, key)
+		}
+	}
+	delete(i.uriToKeys, uri)
+}
+
+// References returns all symbols referencing the given definition, ordered by file URI and
+// then by position so that results are stable across calls.
+func (i *referenceIndex) References(key referenceKey) []*symbol {
+	fileRefs := i.keyToFileRefs[key]
+	if len(fileRefs) == 0 {
+		return nil
+	}
+	var symbols []*symbol
+	for _, fileSymbols := range fileRefs {
+		symbols = append(symbols, fileSymbols...)
+	}
+	slices.SortFunc(symbols, func(symbol1, symbol2 *symbol) int {
+		return cmp.Or(
+			cmp.Compare(symbol1.file.uri, symbol2.file.uri),
+			cmp.Compare(symbol1.span.Start, symbol2.span.Start),
+			cmp.Compare(symbol1.span.End, symbol2.span.End),
+		)
+	})
+	return symbols
+}
+
+// FileReferences returns the symbols in the file at the given URI that reference the given
+// definition. This is a direct lookup, cheap enough for per-cursor-move requests.
+func (i *referenceIndex) FileReferences(key referenceKey, uri protocol.URI) []*symbol {
+	return i.keyToFileRefs[key][uri]
+}