feat(langserver): find references and tests

tools/build_langserver/lsp/BUILD

@@ -5,6 +5,7 @@ go_library(
         "definition.go",
         "diagnostics.go",
         "lsp.go",
+        "references.go",
         "symbols.go",
         "text.go",
     ],
@@ -20,6 +21,7 @@ go_library(
         "//src/parse",
         "//src/parse/asp",
         "//src/plz",
+        "//src/query",
         "//tools/build_langserver/lsp/astutils",
     ],
 )
@@ -30,15 +32,20 @@ go_test(
     srcs = [
         "definition_test.go",
         "lsp_test.go",
+        "references_test.go",
         "symbols_test.go",
     ],
-    data = ["test_data"],
+    data = [
+        "test_data",
+        "test_data_find_references",
+    ],
     deps = [
         ":lsp",
         "///third_party/go/github.com_please-build_buildtools//build",
         "///third_party/go/github.com_sourcegraph_go-lsp//:go-lsp",
         "///third_party/go/github.com_sourcegraph_jsonrpc2//:jsonrpc2",
         "///third_party/go/github.com_stretchr_testify//assert",
+        "///third_party/go/github.com_stretchr_testify//require",
         "//src/cli",
         "//src/core",
     ],

tools/build_langserver/lsp/lsp.go

@@ -42,6 +42,7 @@ type Handler struct {
 type builtin struct {
 	Stmt        *asp.Statement
 	Pos, EndPos asp.FilePosition
+	Labels      []core.BuildLabel // which subinclude targets can provide this (empty for core builtins)
 }
 
 // A Conn is a minimal set of the jsonrpc2.Conn that we need.
@@ -174,6 +175,12 @@ func (h *Handler) handle(method string, params *json.RawMessage) (res interface{
 			return nil, &jsonrpc2.Error{Code: jsonrpc2.CodeInvalidParams}
 		}
 		return h.definition(positionParams)
+	case "textDocument/references":
+		referenceParams := &lsp.ReferenceParams{}
+		if err := json.Unmarshal(*params, referenceParams); err != nil {
+			return nil, &jsonrpc2.Error{Code: jsonrpc2.CodeInvalidParams}
+		}
+		return h.references(referenceParams)
 	default:
 		return nil, &jsonrpc2.Error{Code: jsonrpc2.CodeMethodNotFound}
 	}
@@ -244,6 +251,7 @@ func (h *Handler) initialize(params *lsp.InitializeParams) (*lsp.InitializeResul
 			DocumentFormattingProvider: true,
 			DocumentSymbolProvider:     true,
 			DefinitionProvider:         true,
+			ReferencesProvider:         true,
 			CompletionProvider: &lsp.CompletionOptions{
 				TriggerCharacters: []string{"/", ":"},
 			},
@@ -308,17 +316,43 @@ func (h *Handler) loadParserFunctions() {
 			continue
 		}
 		file := asp.NewFile(filename, data)
+		labels := h.findLabelsForFile(filename)
 		for _, stmt := range stmts {
 			name := stmt.FuncDef.Name
 			h.builtins[name] = append(h.builtins[name], builtin{
 				Stmt:   stmt,
 				Pos:    file.Pos(stmt.Pos),
 				EndPos: file.Pos(stmt.EndPos),
+				Labels: labels,
 			})
 		}
 	}
 }
 
+// findLabelsForFile finds all build labels that produce the given output file.
+// A file can be exposed by multiple targets (e.g., multiple filegroups pointing to the same source).
+func (h *Handler) findLabelsForFile(filename string) []core.BuildLabel {
+	// Make the filename absolute for comparison
+	if !filepath.IsAbs(filename) {
+		filename = filepath.Join(h.root, filename)
+	}
+	var labels []core.BuildLabel
+	for _, pkg := range h.state.Graph.PackageMap() {
+		for _, target := range pkg.AllTargets() {
+			for _, out := range target.FullOutputs() {
+				absOut := out
+				if !filepath.IsAbs(out) {
+					absOut = filepath.Join(h.root, out)
+				}
+				if absOut == filename {
+					labels = append(labels, target.Label)
+				}
+			}
+		}
+	}
+	return labels
+}
+
 // fromURI converts a DocumentURI to a path.
 func fromURI(uri lsp.DocumentURI) string {
 	if !strings.HasPrefix(string(uri), "file://") {

tools/build_langserver/lsp/references.go

@@ -0,0 +1,287 @@
+package lsp
+
+import (
+	"path/filepath"
+	"slices"
+	"strings"
+
+	"github.com/sourcegraph/go-lsp"
+
+	"github.com/thought-machine/please/src/core"
+	"github.com/thought-machine/please/src/parse/asp"
+	"github.com/thought-machine/please/src/query"
+	"github.com/thought-machine/please/tools/build_langserver/lsp/astutils"
+)
+
+// references implements 'find all references' support.
+func (h *Handler) references(params *lsp.ReferenceParams) ([]lsp.Location, error) {
+	doc := h.doc(params.TextDocument.URI)
+	ast := h.parseIfNeeded(doc)
+	f := doc.AspFile()
+	pos := aspPos(params.Position)
+
+	// Check if cursor is on a function definition (def funcname(...))
+	var funcName string
+	asp.WalkAST(ast, func(stmt *asp.Statement) bool {
+		if stmt.FuncDef != nil {
+			stmtStart := f.Pos(stmt.Pos)
+			// Check if cursor is on the function name
+			nameEnd := stmtStart
+			nameEnd.Column += len("def ") + len(stmt.FuncDef.Name)
+			if asp.WithinRange(pos, stmtStart, nameEnd) {
+				funcName = stmt.FuncDef.Name
+				return false
+			}
+		}
+		return true
+	})
+
+	// If we found a function definition, find all calls to it
+	if funcName != "" {
+		// Pass the current filename so we can find the labels for THIS definition
+		return h.findFunctionReferences(funcName, f.Name, params.Context.IncludeDeclaration)
+	}
+
+	// Check if cursor is on a function call (e.g., go_library(...))
+	asp.WalkAST(ast, func(stmt *asp.Statement) bool {
+		if stmt.Ident != nil {
+			stmtStart := f.Pos(stmt.Pos)
+			nameEnd := stmtStart
+			nameEnd.Column += len(stmt.Ident.Name)
+			if asp.WithinRange(pos, stmtStart, nameEnd) {
+				funcName = stmt.Ident.Name
+				return false
+			}
+		}
+		return true
+	})
+
+	if funcName != "" {
+		// For function calls, we don't have a specific definition file - use empty string
+		return h.findFunctionReferences(funcName, "", params.Context.IncludeDeclaration)
+	}
+
+	// Otherwise, look for build label references
+	return h.findLabelReferences(doc, ast, f, pos, params.Context.IncludeDeclaration)
+}
+
+// findFunctionReferences finds all calls to a function across all BUILD files.
+// sourceFile is the file where the definition we're searching from is located (empty if from a call site).
+func (h *Handler) findFunctionReferences(funcName string, sourceFile string, includeDeclaration bool) ([]lsp.Location, error) {
+	locs := []lsp.Location{}
+
+	// Get the labels that can provide this function (a file may be exposed by multiple targets)
+	h.mutex.Lock()
+	builtins := h.builtins[funcName]
+	var defLabels []core.BuildLabel
+	// If we know the source file, find the labels for THAT specific definition
+	// This prevents false positives when multiple files define functions with the same name
+	for _, b := range builtins {
+		if sourceFile != "" && pathsMatch(sourceFile, b.Pos.Filename, h.root) {
+			defLabels = b.Labels
+			break
+		}
+	}
+	// If no source file specified or not found, fall back to first definition's labels
+	if len(defLabels) == 0 && len(builtins) > 0 {
+		defLabels = builtins[0].Labels
+	}
+	h.mutex.Unlock()
+
+	// Search all packages for calls to this function
+	for _, pkg := range h.state.Graph.PackageMap() {
+		// If we know which labels define this function, only search packages that subinclude at least one
+		if len(defLabels) > 0 {
+			allSubincludes := pkg.AllSubincludes(h.state.Graph)
+			hasSubinclude := false
+			for _, defLabel := range defLabels {
+				if slices.Contains(allSubincludes, defLabel) {
+					hasSubinclude = true
+					break
+				}
+			}
+			if !hasSubinclude {
+				continue
+			}
+		}
+
+		uri := lsp.DocumentURI("file://" + filepath.Join(h.root, pkg.Filename))
+		refDoc, err := h.maybeOpenDoc(uri)
+		if err != nil {
+			continue
+		}
+		refAst := h.parseIfNeeded(refDoc)
+		refFile := refDoc.AspFile()
+
+		// Find all statement calls to the function (e.g., go_library(...))
+		asp.WalkAST(refAst, func(stmt *asp.Statement) bool {
+			if stmt.Ident != nil && stmt.Ident.Name == funcName && stmt.Ident.Action != nil && stmt.Ident.Action.Call != nil {
+				start := refFile.Pos(stmt.Pos)
+				end := start
+				end.Column += len(funcName)
+				locs = append(locs, lsp.Location{
+					URI: uri,
+					Range: lsp.Range{
+						Start: lsp.Position{Line: start.Line - 1, Character: start.Column - 1},
+						End:   lsp.Position{Line: end.Line - 1, Character: end.Column - 1},
+					},
+				})
+			}
+			return true
+		})
+
+		// Find expression calls (e.g., x = go_library(...))
+		asp.WalkAST(refAst, func(expr *asp.Expression) bool {
+			if expr.Val.Ident != nil && expr.Val.Ident.Name == funcName && len(expr.Val.Ident.Action) > 0 && expr.Val.Ident.Action[0].Call != nil {
+				start := refFile.Pos(expr.Pos)
+				end := start
+				end.Column += len(funcName)
+				locs = append(locs, lsp.Location{
+					URI: uri,
+					Range: lsp.Range{
+						Start: lsp.Position{Line: start.Line - 1, Character: start.Column - 1},
+						End:   lsp.Position{Line: end.Line - 1, Character: end.Column - 1},
+					},
+				})
+			}
+			return true
+		})
+	}
+
+	// Include the definition itself if requested
+	if includeDeclaration {
+		h.mutex.Lock()
+		if builtins, ok := h.builtins[funcName]; ok && len(builtins) > 0 {
+			b := builtins[0]
+			filename := b.Pos.Filename
+			if !filepath.IsAbs(filename) {
+				filename = filepath.Join(h.root, filename)
+			}
+			locs = append(locs, lsp.Location{
+				URI:   lsp.DocumentURI("file://" + filename),
+				Range: rng(b.Pos, b.EndPos),
+			})
+		}
+		h.mutex.Unlock()
+	}
+
+	return locs, nil
+}
+
+// findLabelReferences finds all references to a build label.
+func (h *Handler) findLabelReferences(doc *doc, ast []*asp.Statement, f *asp.File, pos asp.FilePosition, includeDeclaration bool) ([]lsp.Location, error) {
+	var targetLabel core.BuildLabel
+	var targetName string
+
+	// Check if cursor is on a string (build label)
+	asp.WalkAST(ast, func(expr *asp.Expression) bool {
+		exprStart := f.Pos(expr.Pos)
+		exprEnd := f.Pos(expr.EndPos)
+		if !asp.WithinRange(pos, exprStart, exprEnd) {
+			return false
+		}
+		if expr.Val.String != "" {
+			label := astutils.TrimStrLit(expr.Val.String)
+			if l, err := core.TryParseBuildLabel(label, doc.PkgName, ""); err == nil {
+				targetLabel = l
+			}
+			return false
+		}
+		return true
+	})
+
+	// Check if cursor is on a target definition (name = "...")
+	if targetLabel.IsEmpty() {
+		asp.WalkAST(ast, func(stmt *asp.Statement) bool {
+			if stmt.Ident != nil && stmt.Ident.Action != nil && stmt.Ident.Action.Call != nil {
+				stmtStart := f.Pos(stmt.Pos)
+				stmtEnd := f.Pos(stmt.EndPos)
+				if asp.WithinRange(pos, stmtStart, stmtEnd) {
+					if name := findName(stmt.Ident.Action.Call.Arguments); name != "" {
+						targetLabel = core.BuildLabel{PackageName: doc.PkgName, Name: name}
+						targetName = name
+					}
+				}
+				return false
+			}
+			return true
+		})
+	}
+
+	if targetLabel.IsEmpty() {
+		return []lsp.Location{}, nil
+	}
+
+	// Verify the target exists in the graph before querying revdeps
+	// This prevents panics when the label has an invalid package name
+	if h.state.Graph.Package(targetLabel.PackageName, "") == nil {
+		log.Warning("references: package %q not found in build graph for label %s", targetLabel.PackageName, targetLabel)
+		return []lsp.Location{}, nil
+	}
+
+	// Use query.FindRevdeps to find all reverse dependencies
+	// Parameters: hidden=false, followSubincludes=true, includeSubrepos=true, depth=-1 (unlimited)
+	revdeps := query.FindRevdeps(h.state, core.BuildLabels{targetLabel}, false, true, true, -1)
+
+	locs := []lsp.Location{}
+
+	// For each reverse dependency, find the exact location of the reference in its BUILD file
+	for target := range revdeps {
+		pkg := h.state.Graph.PackageByLabel(target.Label)
+		if pkg == nil {
+			continue
+		}
+
+		uri := lsp.DocumentURI("file://" + filepath.Join(h.root, pkg.Filename))
+		refDoc, err := h.maybeOpenDoc(uri)
+		if err != nil {
+			continue
+		}
+		refAst := h.parseIfNeeded(refDoc)
+		refFile := refDoc.AspFile()
+
+		// Find all string literals that reference our target
+		labelStr := targetLabel.String()
+		shortLabelStr := ":" + targetLabel.Name // For same-package references
+
+		asp.WalkAST(refAst, func(expr *asp.Expression) bool {
+			if expr.Val.String != "" {
+				str := astutils.TrimStrLit(expr.Val.String)
+				// Check if this string matches our target label
+				if str == labelStr || (refDoc.PkgName == targetLabel.PackageName && str == shortLabelStr) {
+					// Also try parsing it as a label to handle relative references
+					if l, err := core.TryParseBuildLabel(str, refDoc.PkgName, ""); err == nil && l == targetLabel {
+						start := refFile.Pos(expr.Pos)
+						end := refFile.Pos(expr.EndPos)
+						locs = append(locs, lsp.Location{
+							URI: uri,
+							Range: lsp.Range{
+								Start: lsp.Position{Line: start.Line - 1, Character: start.Column - 1},
+								End:   lsp.Position{Line: end.Line - 1, Character: end.Column - 1},
+							},
+						})
+					}
+				}
+			}
+			return true
+		})
+	}
+
+	// Optionally include the definition itself if requested
+	if includeDeclaration && targetName != "" {
+		if defLoc := h.findLabel(doc.PkgName, targetLabel.String()); defLoc.URI != "" {
+			locs = append(locs, defLoc)
+		}
+	}
+
+	return locs, nil
+}
+
+// pathsMatch checks if two file paths refer to the same file, handling plz-out/gen paths.
+func pathsMatch(path1, path2, root string) bool {
+	// Strip plz-out/gen/ prefix if present
+	norm1 := strings.TrimPrefix(path1, "plz-out/gen/")
+	norm2 := strings.TrimPrefix(path2, "plz-out/gen/")
+
+	return norm1 == norm2
+}