feat(lsp): resolve definition and references across open documents

lsp/handlers.go

@@ -98,26 +98,50 @@ func (s *Server) hover(p hoverParams) *Hover {
 	}
 }
 
-// definition returns the location of the label definition for a label reference.
+// openAST is one open document with its parsed file.
+type openAST struct {
+	uri  string
+	file *ast.File
+}
+
+// openASTs parses every open document, in URI order for deterministic
+// results.  Parsing is tolerant: a buffer with syntax errors still
+// contributes its usable declarations to the workspace scans.
+func (s *Server) openASTs() []openAST {
+	uris := make([]string, 0, len(s.docs))
+	for uri := range s.docs {
+		uris = append(uris, uri)
+	}
+	sort.Strings(uris)
+	out := make([]openAST, 0, len(uris))
+	for _, uri := range uris {
+		if f, _ := parser.Parse(uriPath(uri), s.docs[uri]); f != nil {
+			out = append(out, openAST{uri: uri, file: f})
+		}
+	}
+	return out
+}
+
+// definition returns the location of the named label or function: a local
+// label in the current document wins, then the TEXT functions of every open
+// document are searched, so a `CALL ·helper(SB)` jumps to its definition in
+// another file.
 func (s *Server) definition(p definitionParams) []Location {
 	text := s.docs[p.TextDocument.URI]
 	word, _ := wordAt(text, p.Position)
 	if word == "" {
 		return nil
 	}
+	name := strings.TrimPrefix(word, "\u00B7")
 
-	// Parse the document to find label definitions.
+	// The local label definition.
-	f, errs := parser.Parse(uriPath(p.TextDocument.URI), text)
+	f, _ := parser.Parse(uriPath(p.TextDocument.URI), text)
-	if f == nil || len(errs) > 0 {
+	if f != nil {
-		return nil
-	}
-
-	// Find the label definition.
 		for _, d := range f.Decls {
 			if t, ok := d.(*ast.Text); ok {
 				for _, stmt := range t.Body {
 					if lbl, ok := stmt.(*ast.Label); ok {
-					if lbl.Name.Text == word {
+						if lbl.Name.Text == name || lbl.Name.Text == word {
 							return []Location{{
 								URI: p.TextDocument.URI,
 								Range: Range{
@@ -130,53 +154,68 @@ func (s *Server) definition(p definitionParams) []Location {
 				}
 			}
 		}
+	}
+
+	// Function definitions across the open workspace.
+	for _, of := range s.openASTs() {
+		for _, d := range of.file.Decls {
+			t, ok := d.(*ast.Text)
+			if !ok || t.Name == nil {
+				continue
+			}
+			if t.Name.Name == name {
+				return []Location{{URI: of.uri, Range: symRange(t.Name)}}
+			}
+		}
+	}
 	return nil
 }
 
-// references returns all locations where the symbol under the cursor appears.
+// references returns all locations where the symbol under the cursor appears
+// across every open document.  The current document matches labels and any
+// operand name, as before; other documents only match SB-qualified operand
+// references and the definition itself, because a bare name is a
+// function-local label whose repeats in other files are unrelated.
 func (s *Server) references(p referenceParams) []Location {
 	text := s.docs[p.TextDocument.URI]
 	word, _ := wordAt(text, p.Position)
 	if word == "" {
 		return nil
 	}
-	f, errs := parser.Parse(uriPath(p.TextDocument.URI), text)
+	name := strings.TrimPrefix(word, "\u00B7")
-	if f == nil || len(errs) > 0 {
-		return nil
-	}
 	uri := p.TextDocument.URI
+
 	var out []Location
-	for _, d := range f.Decls {
+	for _, of := range s.openASTs() {
+		sameDoc := of.uri == uri
+		for _, d := range of.file.Decls {
 			t, ok := d.(*ast.Text)
-		if !ok {
+			if !ok || t.Name == nil {
 				continue
 			}
 			// Include the definition if requested.
-		if p.Context.IncludeDeclaration {
+			if p.Context.IncludeDeclaration && t.Name.Name == name {
-			if t.Name.Name == word {
+				out = append(out, Location{URI: of.uri, Range: symRange(t.Name)})
-				out = append(out, Location{
-					URI:   uri,
-					Range: symRange(t.Name),
-				})
-			}
 			}
 			for _, stmt := range t.Body {
 				switch st := stmt.(type) {
 				case *ast.Label:
-				if st.Name.Text == word {
+					if sameDoc && st.Name.Text == name {
-					out = append(out, Location{
+						out = append(out, Location{URI: of.uri, Range: tokenRange(st.Name)})
-						URI:   uri,
-						Range: tokenRange(st.Name),
-					})
 					}
 				case *ast.Instr:
 					for _, op := range st.Operands {
-					if op.Addr.Sym != nil && op.Addr.Sym.Name == word {
+						if op.Addr.Sym == nil || op.Addr.Sym.Name != name {
+							continue
+						}
+						if !sameDoc && op.Addr.Sym.Pseudo != "SB" {
+							continue
+						}
 						out = append(out, Location{
-							URI: uri,
+							URI: of.uri,
 							Range: Range{
 								Start: Position{Line: op.Pos.Line - 1, Character: op.Pos.Column - 1},
-								End:   Position{Line: op.Pos.Line - 1, Character: op.Pos.Column - 1 + runeLen(word)},
+								End:   Position{Line: op.Pos.Line - 1, Character: op.Pos.Column - 1 + runeLen(name)},
 							},
 						})
 					}

lsp/server_test.go

@@ -691,3 +691,64 @@ func TestCodeActionsTargetsFlaggedFunctionOnly(t *testing.T) {
 		t.Errorf("edit does not insert RET at the end of second:\n%s", actions[0].Edit.Changes["file:///f_amd64.s"][0].NewText)
 	}
 }
+
+// TestCrossFileDefinitionAndReferences opens two documents: docA calls
+// ·helper(SB), docB defines it.  Definition must jump to docB and references
+// must collect the call site in docA plus the definition in docB.
+func TestCrossFileDefinitionAndReferences(t *testing.T) {
+	docA := "TEXT \u00b7caller(SB), NOSPLIT, $0\n" +
+		"\tCALL \u00b7helper(SB)\n" +
+		"\tRET\n"
+	docB := "TEXT \u00b7helper(SB), NOSPLIT, $0\n" +
+		"\tRET\n"
+	in := frame(1, "initialize", map[string]any{}) +
+		frame(nil, "initialized", map[string]any{}) +
+		frame(nil, "textDocument/didOpen", map[string]any{
+			"textDocument": map[string]any{"uri": "file:///a_amd64.s", "languageId": "gasm", "version": 1, "text": docA},
+		}) +
+		frame(nil, "textDocument/didOpen", map[string]any{
+			"textDocument": map[string]any{"uri": "file:///b_amd64.s", "languageId": "gasm", "version": 1, "text": docB},
+		}) +
+		frame(2, "textDocument/definition", map[string]any{
+			"textDocument": map[string]any{"uri": "file:///a_amd64.s"},
+			"position":     map[string]any{"line": 1, "character": 8}, // on helper in CALL ·helper(SB)
+		}) +
+		frame(3, "textDocument/references", map[string]any{
+			"textDocument": map[string]any{"uri": "file:///b_amd64.s"},
+			"position":     map[string]any{"line": 0, "character": 7}, // on helper in TEXT ·helper(SB)
+			"context":      map[string]any{"includeDeclaration": true},
+		}) +
+		frame(nil, "exit", nil)
+	msgs := run(t, in)
+
+	dresp := findByID(msgs, 2)
+	if dresp == nil {
+		t.Fatal("no definition response")
+	}
+	var locs []Location
+	if err := json.Unmarshal(mustResult(t, dresp), &locs); err != nil {
+		t.Fatal(err)
+	}
+	if len(locs) != 1 || locs[0].URI != "file:///b_amd64.s" || locs[0].Range.Start.Line != 0 {
+		t.Fatalf("definition = %+v, want the TEXT in b_amd64.s line 0", locs)
+	}
+
+	rresp := findByID(msgs, 3)
+	if rresp == nil {
+		t.Fatal("no references response")
+	}
+	locs = nil
+	if err := json.Unmarshal(mustResult(t, rresp), &locs); err != nil {
+		t.Fatal(err)
+	}
+	if len(locs) != 2 {
+		t.Fatalf("references = %+v, want the definition in b_amd64.s and the call in a_amd64.s", locs)
+	}
+	byURI := map[string]int{}
+	for _, l := range locs {
+		byURI[l.URI]++
+	}
+	if byURI["file:///a_amd64.s"] != 1 || byURI["file:///b_amd64.s"] != 1 {
+		t.Errorf("references by uri = %v, want one in each file", byURI)
+	}
+}