fix(langserver): support go-to-definition for plugin-defined rules (#3491)

This is a follow up MR from
#3485 . As suggested by
@toastwaffle I split the MR into two smaller pieces. This one is about
go-to-definition in plugin rules.

I applied every single requested change. Furthermore, I removed all
formatting changes my IDE did, to keep the patch short.

---------

Co-authored-by: Andrzej J Skalski <gitstuff@s5i.ch>

Author: njskalski

src/cmap/cerrmap.go

@@ -78,3 +78,14 @@ func (m *ErrMap[K, V]) GetOrSet(key K, f func() (V, error)) (V, error) {
 	}
 	return v.Val, v.Err
 }
+
+// Range calls f for each key-value pair in the map.
+// No particular consistency guarantees are made during iteration.
+func (m *ErrMap[K, V]) Range(f func(key K, val V)) {
+	m.m.Range(func(key K, val errV[V]) {
+		if val.Err != nil {
+			return // skip errors
+		}
+		f(key, val.Val)
+	})
+}

src/cmap/cmap.go

@@ -94,6 +94,14 @@ func (m *Map[K, V]) Values() []V {
 	return ret
 }
 
+// Range calls f for each key-value pair in the map.
+// No particular consistency guarantees are made during iteration.
+func (m *Map[K, V]) Range(f func(key K, val V)) {
+	for i := 0; i < len(m.shards); i++ {
+		m.shards[i].Range(f)
+	}
+}
+
 // An awaitableValue represents a value in the map & an awaitable channel for it to exist.
 type awaitableValue[V any] struct {
 	Val  V
@@ -195,3 +203,14 @@ func (s *shard[K, V]) Contains(key K) bool {
 	_, ok := s.m[key]
 	return ok
 }
+
+// Range calls f for each key-value pair in this shard.
+func (s *shard[K, V]) Range(f func(key K, val V)) {
+	s.l.RLock()
+	defer s.l.RUnlock()
+	for k, v := range s.m {
+		if v.Wait == nil { // Only include completed values
+			f(k, v.Val)
+		}
+	}
+}

src/parse/asp/parser.go

@@ -257,6 +257,24 @@ func (p *Parser) optimiseBuiltinCalls(stmts []*Statement) {
 	}
 }
 
+// AllFunctionsByFile returns all function definitions grouped by filename.
+// This includes functions from builtins, plugins, and subincludes.
+// It iterates over the ASTs stored by the interpreter.
+func (p *Parser) AllFunctionsByFile() map[string][]*Statement {
+	if p.interpreter == nil || p.interpreter.asts == nil {
+		return nil
+	}
+	result := make(map[string][]*Statement)
+	p.interpreter.asts.Range(func(filename string, stmts []*Statement) {
+		for _, stmt := range stmts {
+			if stmt.FuncDef != nil {
+				result[filename] = append(result[filename], stmt)
+			}
+		}
+	})
+	return result
+}
+
 // whitelistedKwargs returns true if the given built-in function name is allowed to
 // be called as non-kwargs.
 // TODO(peterebden): Come up with a syntax that exposes this directly in the file.

src/parse/init.go

@@ -25,6 +25,19 @@ func InitParser(state *core.BuildState) *core.BuildState {
 	return state
 }
 
+// GetAspParser returns the underlying asp.Parser from the state's parser.
+// This is useful for tools like the language server that need direct access to AST information.
+// Returns nil if the state's parser is not set or is not an aspParser.
+func GetAspParser(state *core.BuildState) *asp.Parser {
+	if state.Parser == nil {
+		return nil
+	}
+	if ap, ok := state.Parser.(*aspParser); ok {
+		return ap.parser
+	}
+	return nil
+}
+
 // aspParser implements the core.Parser interface around our parser package.
 type aspParser struct {
 	parser *asp.Parser

tools/build_langserver/lsp/BUILD

@@ -17,6 +17,7 @@ go_library(
         "//rules",
         "//src/core",
         "//src/fs",
+        "//src/parse",
         "//src/parse/asp",
         "//src/plz",
         "//tools/build_langserver/lsp/astutils",

tools/build_langserver/lsp/completion.go

@@ -109,10 +109,10 @@ func (h *Handler) completeString(doc *doc, s string, line, col int) (*lsp.Comple
 // completeIdent completes an arbitrary identifier
 func (h *Handler) completeIdent(doc *doc, s string, line, col int) (*lsp.CompletionList, error) {
 	list := &lsp.CompletionList{}
-	for name, f := range h.builtins {
-		if strings.HasPrefix(name, s) {
+	for name, builtins := range h.builtins {
+		if strings.HasPrefix(name, s) && len(builtins) > 0 {
 			item := completionItem(name, "", line, col)
-			item.Documentation = f.Stmt.FuncDef.Docstring
+			item.Documentation = builtins[0].Stmt.FuncDef.Docstring
 			item.Kind = lsp.CIKFunction
 			list.Items = append(list.Items, item)
 		}

tools/build_langserver/lsp/definition.go

@@ -18,41 +18,40 @@ func (h *Handler) definition(params *lsp.TextDocumentPositionParams) ([]lsp.Loca
 	ast := h.parseIfNeeded(doc)
 	f := doc.AspFile()
 
-	var locs []lsp.Location
+	locs := []lsp.Location{}
 	pos := aspPos(params.Position)
 	asp.WalkAST(ast, func(expr *asp.Expression) bool {
-		if !asp.WithinRange(pos, f.Pos(expr.Pos), f.Pos(expr.EndPos)) {
+		exprStart := f.Pos(expr.Pos)
+		exprEnd := f.Pos(expr.EndPos)
+		if !asp.WithinRange(pos, exprStart, exprEnd) {
 			return false
 		}
-
 		if expr.Val.Ident != nil {
 			if loc := h.findGlobal(expr.Val.Ident.Name); loc.URI != "" {
 				locs = append(locs, loc)
 			}
 			return false
 		}
-
 		if expr.Val.String != "" {
 			label := astutils.TrimStrLit(expr.Val.String)
 			if loc := h.findLabel(doc.PkgName, label); loc.URI != "" {
 				locs = append(locs, loc)
 			}
 			return false
 		}
-
 		return true
 	})
-	// It might also be a statement.
+	// It might also be a statement (e.g. a function call like go_library(...))
 	asp.WalkAST(ast, func(stmt *asp.Statement) bool {
 		if stmt.Ident != nil {
-			endPos := f.Pos(stmt.Pos)
+			stmtStart := f.Pos(stmt.Pos)
+			endPos := stmtStart
 			// TODO(jpoole): The AST should probably just have this information
 			endPos.Column += len(stmt.Ident.Name)
 
-			if !asp.WithinRange(pos, f.Pos(stmt.Pos), endPos) {
-				return false
+			if !asp.WithinRange(pos, stmtStart, endPos) {
+				return true // continue to other statements
 			}
-
 			if loc := h.findGlobal(stmt.Ident.Name); loc.URI != "" {
 				locs = append(locs, loc)
 			}
@@ -78,6 +77,9 @@ func (h *Handler) findLabel(currentPath, label string) lsp.Location {
 	}
 
 	pkg := h.state.Graph.PackageByLabel(l)
+	if pkg == nil {
+		return lsp.Location{}
+	}
 	uri := lsp.DocumentURI("file://" + filepath.Join(h.root, pkg.Filename))
 	loc := lsp.Location{URI: uri}
 	doc, err := h.maybeOpenDoc(uri)
@@ -137,9 +139,18 @@ func findName(args []asp.CallArgument) string {
 
 // findGlobal returns the location of a global of the given name.
 func (h *Handler) findGlobal(name string) lsp.Location {
-	if f, present := h.builtins[name]; present {
+	h.mutex.Lock()
+	builtins := h.builtins[name]
+	h.mutex.Unlock()
+	if len(builtins) > 0 {
+		f := builtins[0]
+		filename := f.Pos.Filename
+		// Make path absolute if it's relative
+		if !filepath.IsAbs(filename) {
+			filename = filepath.Join(h.root, filename)
+		}
 		return lsp.Location{
-			URI:   lsp.DocumentURI("file://" + f.Pos.Filename),
+			URI:   lsp.DocumentURI("file://" + filename),
 			Range: rng(f.Pos, f.EndPos),
 		}
 	}

tools/build_langserver/lsp/lsp.go

@@ -20,6 +20,7 @@ import (
 	"github.com/thought-machine/please/rules"
 	"github.com/thought-machine/please/src/core"
 	"github.com/thought-machine/please/src/fs"
+	"github.com/thought-machine/please/src/parse"
 	"github.com/thought-machine/please/src/parse/asp"
 	"github.com/thought-machine/please/src/plz"
 )
@@ -33,7 +34,7 @@ type Handler struct {
 	mutex    sync.Mutex // guards docs
 	state    *core.BuildState
 	parser   *asp.Parser
-	builtins map[string]builtin
+	builtins map[string][]builtin
 	pkgs     *pkg
 	root     string
 }
@@ -55,7 +56,7 @@ func NewHandler() *Handler {
 	return &Handler{
 		docs:     map[string]*doc{},
 		pkgs:     &pkg{},
-		builtins: map[string]builtin{},
+		builtins: map[string][]builtin{},
 	}
 }
 
@@ -195,13 +196,35 @@ func (h *Handler) initialize(params *lsp.InitializeParams) (*lsp.InitializeResul
 	}
 	h.state = core.NewBuildState(config)
 	h.state.NeedBuild = false
-	// We need an unwrapped parser instance as well for raw access.
-	h.parser = asp.NewParser(h.state)
+	// Initialize the parser on state first, so that plz.RunHost uses the same parser.
+	// This ensures plugin subincludes are stored in the same AST cache we use.
+	parse.InitParser(h.state)
+	h.parser = parse.GetAspParser(h.state)
+	if h.parser == nil {
+		return nil, fmt.Errorf("failed to get asp parser from state")
+	}
 	// Parse everything in the repo up front.
 	// This is a lot easier than trying to do clever partial parses later on, although
 	// eventually we may want that if we start dealing with truly large repos.
 	go func() {
+		// Start a goroutine to periodically load parser functions as they become available.
+		// This allows go-to-definition to work progressively while the full parse runs.
+		done := make(chan struct{})
+		go func() {
+			ticker := time.NewTicker(2 * time.Second)
+			defer ticker.Stop()
+			for {
+				select {
+				case <-done:
+					h.loadParserFunctions()
+					return
+				case <-ticker.C:
+					h.loadParserFunctions()
+				}
+			}
+		}()
 		plz.RunHost(core.WholeGraph, h.state)
+		close(done)
 		log.Debug("initial parse complete")
 		h.buildPackageTree()
 		log.Debug("built completion package tree")
@@ -256,18 +279,46 @@ func (h *Handler) loadBuiltins() error {
 		f := asp.NewFile(dest, data)
 		for _, stmt := range stmts {
 			if stmt.FuncDef != nil {
-				h.builtins[stmt.FuncDef.Name] = builtin{
+				h.builtins[stmt.FuncDef.Name] = append(h.builtins[stmt.FuncDef.Name], builtin{
 					Stmt:   stmt,
 					Pos:    f.Pos(stmt.Pos),
 					EndPos: f.Pos(stmt.EndPos),
-				}
+				})
 			}
 		}
 	}
 	log.Debug("loaded builtin function information")
 	return nil
 }
 
+// loadParserFunctions loads function definitions from the parser's ASTs.
+// This includes plugin-defined functions like go_library, python_library, etc.
+func (h *Handler) loadParserFunctions() {
+	funcsByFile := h.parser.AllFunctionsByFile()
+	if funcsByFile == nil {
+		return
+	}
+	h.mutex.Lock()
+	defer h.mutex.Unlock()
+	for filename, stmts := range funcsByFile {
+		// Read the file to create a File object for position conversion
+		data, err := os.ReadFile(filename)
+		if err != nil {
+			log.Warning("failed to read file %s: %v", filename, err)
+			continue
+		}
+		file := asp.NewFile(filename, data)
+		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),
+			})
+		}
+	}
+}
+
 // fromURI converts a DocumentURI to a path.
 func fromURI(uri lsp.DocumentURI) string {
 	if !strings.HasPrefix(string(uri), "file://") {

tools/build_langserver/lsp/lsp_test.go

@@ -458,7 +458,7 @@ func TestCompletionFunction(t *testing.T) {
 			Kind:             lsp.CIKFunction,
 			InsertTextFormat: lsp.ITFPlainText,
 			TextEdit:         textEdit("plugin_repo", 0, 4, 0),
-			Documentation:    h.builtins["plugin_repo"].Stmt.FuncDef.Docstring,
+			Documentation:    h.builtins["plugin_repo"][0].Stmt.FuncDef.Docstring,
 		}},
 	}, completions)
 }
@@ -492,7 +492,7 @@ func TestCompletionPartialFunction(t *testing.T) {
 			Kind:             lsp.CIKFunction,
 			InsertTextFormat: lsp.ITFPlainText,
 			TextEdit:         textEdit("plugin_repo", 0, 9, 0),
-			Documentation:    h.builtins["plugin_repo"].Stmt.FuncDef.Docstring,
+			Documentation:    h.builtins["plugin_repo"][0].Stmt.FuncDef.Docstring,
 		}},
 	}, completions)
 }