Fix regressions for well-known types in LSP

Author: doriable

private/buf/buflsp/file.go

@@ -321,7 +321,7 @@ func (f *file) IndexSymbols(ctx context.Context) {
 	for _, sym := range unresolved {
 		switch kind := sym.kind.(type) {
 		case *reference:
-			def := f.resolveASTDefinition(kind.def, kind.fullName)
+			def := f.resolveASTDefinition(ctx, kind.def, kind.fullName)
 			sym.def = def
 			if def == nil {
 				// In the case where the symbol is not resolved, we continue
@@ -339,7 +339,7 @@ func (f *file) IndexSymbols(ctx context.Context) {
 			}
 			referenceable.references = append(referenceable.references, sym)
 		case *option:
-			def := f.resolveASTDefinition(kind.def, kind.defFullName)
+			def := f.resolveASTDefinition(ctx, kind.def, kind.defFullName)
 			sym.def = def
 			if def != nil {
 				referenceable, ok := def.kind.(*referenceable)
@@ -354,7 +354,7 @@ func (f *file) IndexSymbols(ctx context.Context) {
 					referenceable.references = append(referenceable.references, sym)
 				}
 			}
-			typeDef := f.resolveASTDefinition(kind.typeDef, kind.typeDefFullName)
+			typeDef := f.resolveASTDefinition(ctx, kind.typeDef, kind.typeDefFullName)
 			sym.typeDef = typeDef
 		default:
 			// This shouldn't happen, logging a warning
@@ -756,36 +756,38 @@ func (f *file) messageToSymbolsHelper(msg ir.MessageValue, index int, parents []
 
 // resolveASTDefinition is a helper for resolving the [ast.DeclDef] to the *[symbol], if
 // there is a matching indexed *[symbol].
-func (f *file) resolveASTDefinition(def ast.DeclDef, defName ir.FullName) *symbol {
-	// No workspace, we cannot resolve the AST definition
-	if f.workspace == nil {
-		return nil
+func (f *file) resolveASTDefinition(ctx context.Context, def ast.DeclDef, defName ir.FullName) *symbol {
+	// First check if the definition is in the current file.
+	if def.Span().Path() == f.objectInfo.LocalPath() {
+		return f.referenceableSymbols[defName]
 	}
-	// We resolve the import path of the span of the AST definition and search for it in our
-	// workspace.
-	fileInfo, ok := f.workspace.fileNameToFileInfo[def.Span().Path()]
-	if !ok {
-		// Unable to resolve an importable path for the file in our workspace, log a debug
-		// statement and return no definition.
-		f.lsp.logger.Debug(
-			"unable to resolve an importable file path for local path",
-			slog.String("file", f.uri.Filename()),
-			slog.String("local path", def.Span().Path()),
-		)
+	// No workspace, we cannot resolve the AST definition from outside of the file.
+	if f.workspace == nil {
 		return nil
 	}
-	file, ok := f.workspace.PathToFile()[fileInfo.Path()]
-	if !ok {
-		// Check current file
-		if def.Span().Path() != f.objectInfo.Path() {
-			// If there is no file for the [ast.DeclDef] span, which can if it is a predeclared
-			// type, for example, or if the file we are checking has not indexed imports, then
-			// we return nil.
-			return nil
+	for objectInfo := range f.workspace.fileInfos(ctx) {
+		if objectInfo.LocalPath() == def.Span().Path() {
+			file, ok := f.workspace.PathToFile()[objectInfo.Path()]
+			if !ok {
+				// This shouldn't happen, if it does, we log an error and optimistically continue
+				// to walk file infos.
+				f.lsp.logger.Error(
+					"found local path that does not resolve to a file in the workspace",
+					slog.String("file", f.uri.Filename()),
+					slog.String("local path", def.Span().Path()),
+				)
+			}
+			return file.referenceableSymbols[defName]
 		}
-		file = f
 	}
-	return file.referenceableSymbols[defName]
+	// Unable to resolve an importable path for the file in our workspace, log a debug
+	// statement and return no definition.
+	f.lsp.logger.Debug(
+		"unable to resolve an importable file path for local path",
+		slog.String("file", f.uri.Filename()),
+		slog.String("local path", def.Span().Path()),
+	)
+	return nil
 }
 
 // SymbolAt finds a symbol in this file at the given cursor position, if one exists.

private/buf/buflsp/workspace.go

@@ -16,7 +16,9 @@ package buflsp
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"io"
 	"iter"
 	"log/slog"
 
@@ -296,11 +298,11 @@ func (w *workspace) fileInfos(ctx context.Context) iter.Seq[storage.ObjectInfo]
 		if err := w.lsp.wktBucket.Walk(ctx, "", func(objectInfo storage.ObjectInfo) error {
 			if _, ok := seen[objectInfo.Path()]; !ok {
 				if !yield(wktObjectInfo{objectInfo}) {
-					return nil
+					return io.EOF
 				}
 			}
 			return nil
-		}); err != nil {
+		}); err != nil && !errors.Is(err, io.EOF) {
 			w.lsp.logger.Error("wkt bucket failed", xslog.ErrorAttr(err))
 		}
 	}