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

lsp/handlers.go

@@ -98,85 +98,124 @@ 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.
-	f, errs := parser.Parse(uriPath(p.TextDocument.URI), text)
-	if f == nil || len(errs) > 0 {
-		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 {
-						return []Location{{
-							URI: p.TextDocument.URI,
-							Range: Range{
-								Start: Position{Line: lbl.Name.Pos.Line - 1, Character: lbl.Name.Pos.Column - 1},
-								End:   Position{Line: lbl.Name.Pos.Line - 1, Character: lbl.Name.Pos.Column - 1 + len(word)},
-							},
-						}}
+	// The local label definition.
+	f, _ := parser.Parse(uriPath(p.TextDocument.URI), text)
+	if f != nil {
+		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 == name || lbl.Name.Text == word {
+							return []Location{{
+								URI: p.TextDocument.URI,
+								Range: Range{
+									Start: Position{Line: lbl.Name.Pos.Line - 1, Character: lbl.Name.Pos.Column - 1},
+									End:   Position{Line: lbl.Name.Pos.Line - 1, Character: lbl.Name.Pos.Column - 1 + len(word)},
+								},
+							}}
+						}
 					}
 				}
 			}
 		}
 	}
+
+	// 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)
-	if f == nil || len(errs) > 0 {
-		return nil
-	}
+	name := strings.TrimPrefix(word, "\u00B7")
 	uri := p.TextDocument.URI
+
 	var out []Location
-	for _, d := range f.Decls {
-		t, ok := d.(*ast.Text)
-		if !ok {
-			continue
-		}
-		// Include the definition if requested.
-		if p.Context.IncludeDeclaration {
-			if t.Name.Name == word {
-				out = append(out, Location{
-					URI:   uri,
-					Range: symRange(t.Name),
-				})
+	for _, of := range s.openASTs() {
+		sameDoc := of.uri == uri
+		for _, d := range of.file.Decls {
+			t, ok := d.(*ast.Text)
+			if !ok || t.Name == nil {
+				continue
 			}
-		}
-		for _, stmt := range t.Body {
-			switch st := stmt.(type) {
-			case *ast.Label:
-				if st.Name.Text == word {
-					out = append(out, Location{
-						URI:   uri,
-						Range: tokenRange(st.Name),
-					})
-				}
-			case *ast.Instr:
-				for _, op := range st.Operands {
-					if op.Addr.Sym != nil && op.Addr.Sym.Name == word {
+			// Include the definition if requested.
+			if p.Context.IncludeDeclaration && t.Name.Name == name {
+				out = append(out, Location{URI: of.uri, Range: symRange(t.Name)})
+			}
+			for _, stmt := range t.Body {
+				switch st := stmt.(type) {
+				case *ast.Label:
+					if sameDoc && st.Name.Text == name {
+						out = append(out, Location{URI: of.uri, Range: tokenRange(st.Name)})
+					}
+				case *ast.Instr:
+					for _, op := range st.Operands {
+						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)},
 							},
 						})
 					}