gopls/internal/goasm: support Definition in Go *.s assembly

This CL provides a minimal implementation of the
Definition query within Go assembly files, plus a test.
For now it only works for references to package-level
symbols in the same package or a dependency.

Details:
- add file.Kind Asm and protocol.LanguageKind "go.s".
- include .s files in metadata.Graph.IDs mapping.
- set LanguageKind correctly in gopls CLI.

Also:
- add String() method to file.Handle.
- add convenient forward deps iterator to Graph.
- internal/extract: extract notes from .s files too.

Updates golang/go#71754

Change-Id: I0c518c3279f825411221ebe23dc04654e129fc56
Reviewed-on: https://go-review.googlesource.com/c/tools/+/649461
Auto-Submit: Alan Donovan <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Alan Donovan <[email protected]>
Reviewed-by: Robert Findley <[email protected]>
Commit-Queue: Alan Donovan <[email protected]>

Author: adonovan

gopls/internal/cache/fs_memoized.go

@@ -41,6 +41,8 @@ type diskFile struct {
 	err     error
 }
 
+func (h *diskFile) String() string { return h.uri.Path() }
+
 func (h *diskFile) URI() protocol.DocumentURI { return h.uri }
 
 func (h *diskFile) Identity() file.Identity {

gopls/internal/cache/fs_overlay.go

@@ -64,6 +64,8 @@ type overlay struct {
 	saved bool
 }
 
+func (o *overlay) String() string { return o.uri.Path() }
+
 func (o *overlay) URI() protocol.DocumentURI { return o.uri }
 
 func (o *overlay) Identity() file.Identity {

gopls/internal/cache/metadata/graph.go

@@ -5,7 +5,9 @@
 package metadata
 
 import (
+	"iter"
 	"sort"
+	"strings"
 
 	"golang.org/x/tools/go/packages"
 	"golang.org/x/tools/gopls/internal/protocol"
@@ -99,6 +101,11 @@ func newGraph(pkgs map[PackageID]*Package) *Graph {
 		for _, uri := range mp.GoFiles {
 			uris[uri] = struct{}{}
 		}
+		for _, uri := range mp.OtherFiles {
+			if strings.HasSuffix(string(uri), ".s") { // assembly
+				uris[uri] = struct{}{}
+			}
+		}
 		for uri := range uris {
 			uriIDs[uri] = append(uriIDs[uri], id)
 		}
@@ -160,6 +167,35 @@ func (g *Graph) ReverseReflexiveTransitiveClosure(ids ...PackageID) map[PackageI
 	return seen
 }
 
+// ForwardReflexiveTransitiveClosure returns an iterator over the
+// specified nodes and all their forward dependencies, in an arbitrary
+// topological (dependencies-first) order. The order may vary.
+func (g *Graph) ForwardReflexiveTransitiveClosure(ids ...PackageID) iter.Seq[*Package] {
+	return func(yield func(*Package) bool) {
+		seen := make(map[PackageID]bool)
+		var visit func(PackageID) bool
+		visit = func(id PackageID) bool {
+			if !seen[id] {
+				seen[id] = true
+				if mp := g.Packages[id]; mp != nil {
+					for _, depID := range mp.DepsByPkgPath {
+						if !visit(depID) {
+							return false
+						}
+					}
+					if !yield(mp) {
+						return false
+					}
+				}
+			}
+			return true
+		}
+		for _, id := range ids {
+			visit(id)
+		}
+	}
+}
+
 // breakImportCycles breaks import cycles in the metadata by deleting
 // Deps* edges. It modifies only metadata present in the 'updates'
 // subset. This function has an internal test.

gopls/internal/cache/parse_cache_test.go

@@ -218,6 +218,10 @@ type fakeFileHandle struct {
 	hash file.Hash
 }
 
+func (h fakeFileHandle) String() string {
+	return h.uri.Path()
+}
+
 func (h fakeFileHandle) URI() protocol.DocumentURI {
 	return h.uri
 }

gopls/internal/cache/session.go

@@ -1084,6 +1084,7 @@ type brokenFile struct {
 	err error
 }
 
+func (b brokenFile) String() string            { return b.uri.Path() }
 func (b brokenFile) URI() protocol.DocumentURI { return b.uri }
 func (b brokenFile) Identity() file.Identity   { return file.Identity{URI: b.uri} }
 func (b brokenFile) SameContentsOnDisk() bool  { return false }

gopls/internal/cache/snapshot.go

@@ -323,6 +323,8 @@ func fileKind(fh file.Handle) file.Kind {
 		return file.Sum
 	case ".work":
 		return file.Work
+	case ".s":
+		return file.Asm
 	}
 	return file.UnknownKind
 }
@@ -645,13 +647,32 @@ func (s *Snapshot) Tests(ctx context.Context, ids ...PackageID) ([]*testfuncs.In
 	return indexes, s.forEachPackage(ctx, ids, pre, post)
 }
 
+// NarrowestMetadataForFile returns metadata for the narrowest package
+// (the one with the fewest files) that encloses the specified file.
+// The result may be a test variant, but never an intermediate test variant.
+func (snapshot *Snapshot) NarrowestMetadataForFile(ctx context.Context, uri protocol.DocumentURI) (*metadata.Package, error) {
+	mps, err := snapshot.MetadataForFile(ctx, uri)
+	if err != nil {
+		return nil, err
+	}
+	metadata.RemoveIntermediateTestVariants(&mps)
+	if len(mps) == 0 {
+		return nil, fmt.Errorf("no package metadata for file %s", uri)
+	}
+	return mps[0], nil
+}
+
 // MetadataForFile returns a new slice containing metadata for each
 // package containing the Go file identified by uri, ordered by the
 // number of CompiledGoFiles (i.e. "narrowest" to "widest" package),
 // and secondarily by IsIntermediateTestVariant (false < true).
 // The result may include tests and intermediate test variants of
 // importable packages.
 // It returns an error if the context was cancelled.
+//
+// TODO(adonovan): in nearly all cases the caller must use
+// [metadata.RemoveIntermediateTestVariants]. Make this a parameter to
+// force the caller to consider it (and reduce code).
 func (s *Snapshot) MetadataForFile(ctx context.Context, uri protocol.DocumentURI) ([]*metadata.Package, error) {
 	if s.view.typ == AdHocView {
 		// As described in golang/go#57209, in ad-hoc workspaces (where we load ./

gopls/internal/cmd/cmd.go

@@ -773,10 +773,25 @@ func (c *connection) openFile(ctx context.Context, uri protocol.DocumentURI) (*c
 		return nil, file.err
 	}
 
+	// Choose language ID from file extension.
+	var langID protocol.LanguageKind // "" eventually maps to file.UnknownKind
+	switch filepath.Ext(uri.Path()) {
+	case ".go":
+		langID = "go"
+	case ".mod":
+		langID = "go.mod"
+	case ".sum":
+		langID = "go.sum"
+	case ".work":
+		langID = "go.work"
+	case ".s":
+		langID = "go.s"
+	}
+
 	p := &protocol.DidOpenTextDocumentParams{
 		TextDocument: protocol.TextDocumentItem{
 			URI:        uri,
-			LanguageID: "go",
+			LanguageID: langID,
 			Version:    1,
 			Text:       string(file.mapper.Content),
 		},

gopls/internal/cmd/definition.go

@@ -96,7 +96,7 @@ func (d *definition) Run(ctx context.Context, args ...string) error {
 	}
 
 	if len(locs) == 0 {
-		return fmt.Errorf("%v: not an identifier", from)
+		return fmt.Errorf("%v: no definition location (not an identifier?)", from)
 	}
 	file, err = conn.openFile(ctx, locs[0].URI)
 	if err != nil {

gopls/internal/file/file.go

@@ -49,6 +49,8 @@ type Handle interface {
 	// Content returns the contents of a file.
 	// If the file is not available, returns a nil slice and an error.
 	Content() ([]byte, error)
+	// String returns the file's path.
+	String() string
 }
 
 // A Source maps URIs to Handles.

gopls/internal/file/kind.go

@@ -28,6 +28,8 @@ const (
 	Tmpl
 	// Work is a go.work file.
 	Work
+	// Asm is a Go assembly (.s) file.
+	Asm
 )
 
 func (k Kind) String() string {
@@ -42,13 +44,15 @@ func (k Kind) String() string {
 		return "tmpl"
 	case Work:
 		return "go.work"
+	case Asm:
+		return "Go assembly"
 	default:
 		return fmt.Sprintf("internal error: unknown file kind %d", k)
 	}
 }
 
 // KindForLang returns the gopls file [Kind] associated with the given LSP
-// LanguageKind string from protocol.TextDocumentItem.LanguageID,
+// LanguageKind string from the LanguageID field of [protocol.TextDocumentItem],
 // or UnknownKind if the language is not one recognized by gopls.
 func KindForLang(langID protocol.LanguageKind) Kind {
 	switch langID {
@@ -62,6 +66,8 @@ func KindForLang(langID protocol.LanguageKind) Kind {
 		return Tmpl
 	case "go.work":
 		return Work
+	case "go.s":
+		return Asm
 	default:
 		return UnknownKind
 	}

gopls/internal/goasm/definition.go

@@ -0,0 +1,125 @@
+// Copyright 2025 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package goasm
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"go/token"
+	"strings"
+	"unicode"
+
+	"golang.org/x/tools/gopls/internal/cache"
+	"golang.org/x/tools/gopls/internal/cache/metadata"
+	"golang.org/x/tools/gopls/internal/file"
+	"golang.org/x/tools/gopls/internal/protocol"
+	"golang.org/x/tools/gopls/internal/util/morestrings"
+	"golang.org/x/tools/internal/event"
+)
+
+// Definition handles the textDocument/definition request for Go assembly files.
+func Definition(ctx context.Context, snapshot *cache.Snapshot, fh file.Handle, position protocol.Position) ([]protocol.Location, error) {
+	ctx, done := event.Start(ctx, "goasm.Definition")
+	defer done()
+
+	mp, err := snapshot.NarrowestMetadataForFile(ctx, fh.URI())
+	if err != nil {
+		return nil, err
+	}
+
+	// Read the file.
+	content, err := fh.Content()
+	if err != nil {
+		return nil, err
+	}
+	mapper := protocol.NewMapper(fh.URI(), content)
+	offset, err := mapper.PositionOffset(position)
+	if err != nil {
+		return nil, err
+	}
+
+	// Figure out the selected symbol.
+	// For now, just find the identifier around the cursor.
+	//
+	// TODO(adonovan): use a real asm parser; see cmd/asm/internal/asm/parse.go.
+	// Ideally this would just be just another attribute of the
+	// type-checked cache.Package.
+	nonIdentRune := func(r rune) bool { return !isIdentRune(r) }
+	i := bytes.LastIndexFunc(content[:offset], nonIdentRune)
+	j := bytes.IndexFunc(content[offset:], nonIdentRune)
+	if j < 0 || j == 0 {
+		return nil, nil // identifier runs to EOF, or not an identifier
+	}
+	sym := string(content[i+1 : offset+j])
+	sym = strings.ReplaceAll(sym, "·", ".") // (U+00B7 MIDDLE DOT)
+	sym = strings.ReplaceAll(sym, "∕", "/") // (U+2215 DIVISION SLASH)
+	if sym != "" && sym[0] == '.' {
+		sym = string(mp.PkgPath) + sym
+	}
+
+	// package-qualified symbol?
+	if pkgpath, name, ok := morestrings.CutLast(sym, "."); ok {
+		// Find declaring package among dependencies.
+		//
+		// TODO(adonovan): assembly may legally reference
+		// non-dependencies. For example, sync/atomic calls
+		// internal/runtime/atomic. Perhaps we should search
+		// the entire metadata graph, but that's path-dependent.
+		var declaring *metadata.Package
+		for pkg := range snapshot.MetadataGraph().ForwardReflexiveTransitiveClosure(mp.ID) {
+			if pkg.PkgPath == metadata.PackagePath(pkgpath) {
+				declaring = pkg
+				break
+			}
+		}
+		if declaring == nil {
+			return nil, fmt.Errorf("package %q is not a dependency", pkgpath)
+		}
+
+		pkgs, err := snapshot.TypeCheck(ctx, declaring.ID)
+		if err != nil {
+			return nil, err
+		}
+		pkg := pkgs[0]
+		def := pkg.Types().Scope().Lookup(name)
+		if def == nil {
+			return nil, fmt.Errorf("no symbol %q in package %q", name, pkgpath)
+		}
+		loc, err := mapPosition(ctx, pkg.FileSet(), snapshot, def.Pos(), def.Pos())
+		if err == nil {
+			return []protocol.Location{loc}, nil
+		}
+	}
+
+	// TODO(adonovan): support jump to var, block label, and other
+	// TEXT, DATA, and GLOBAL symbols in the same file. Needs asm parser.
+
+	return nil, nil
+}
+
+// The assembler allows center dot (· U+00B7) and
+// division slash (∕ U+2215) to work as identifier characters.
+func isIdentRune(r rune) bool {
+	return unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_' || r == '·' || r == '∕'
+}
+
+// TODO(rfindley): avoid the duplicate column mapping here, by associating a
+// column mapper with each file handle.
+// TODO(adonovan): plundered from ../golang; factor.
+func mapPosition(ctx context.Context, fset *token.FileSet, s file.Source, start, end token.Pos) (protocol.Location, error) {
+	file := fset.File(start)
+	uri := protocol.URIFromPath(file.Name())
+	fh, err := s.ReadFile(ctx, uri)
+	if err != nil {
+		return protocol.Location{}, err
+	}
+	content, err := fh.Content()
+	if err != nil {
+		return protocol.Location{}, err
+	}
+	m := protocol.NewMapper(fh.URI(), content)
+	return m.PosLocation(file, start, end)
+}

gopls/internal/golang/snapshot.go

@@ -14,19 +14,9 @@ import (
 	"golang.org/x/tools/gopls/internal/protocol"
 )
 
-// NarrowestMetadataForFile returns metadata for the narrowest package
-// (the one with the fewest files) that encloses the specified file.
-// The result may be a test variant, but never an intermediate test variant.
+//go:fix inline
 func NarrowestMetadataForFile(ctx context.Context, snapshot *cache.Snapshot, uri protocol.DocumentURI) (*metadata.Package, error) {
-	mps, err := snapshot.MetadataForFile(ctx, uri)
-	if err != nil {
-		return nil, err
-	}
-	metadata.RemoveIntermediateTestVariants(&mps)
-	if len(mps) == 0 {
-		return nil, fmt.Errorf("no package metadata for file %s", uri)
-	}
-	return mps[0], nil
+	return snapshot.NarrowestMetadataForFile(ctx, uri)
 }
 
 // NarrowestPackageForFile is a convenience function that selects the narrowest

gopls/internal/server/definition.go

@@ -9,6 +9,7 @@ import (
 	"fmt"
 
 	"golang.org/x/tools/gopls/internal/file"
+	"golang.org/x/tools/gopls/internal/goasm"
 	"golang.org/x/tools/gopls/internal/golang"
 	"golang.org/x/tools/gopls/internal/label"
 	"golang.org/x/tools/gopls/internal/protocol"
@@ -37,6 +38,8 @@ func (s *server) Definition(ctx context.Context, params *protocol.DefinitionPara
 		return template.Definition(snapshot, fh, params.Position)
 	case file.Go:
 		return golang.Definition(ctx, snapshot, fh, params.Position)
+	case file.Asm:
+		return goasm.Definition(ctx, snapshot, fh, params.Position)
 	default:
 		return nil, fmt.Errorf("can't find definitions for file type %s", kind)
 	}

gopls/internal/test/integration/fake/editor.go

@@ -113,11 +113,12 @@ type EditorConfig struct {
 
 	// Map of language ID -> regexp to match, used to set the file type of new
 	// buffers. Applied as an overlay on top of the following defaults:
-	//  "go" -> ".*\.go"
+	//  "go"     -> ".*\.go"
 	//  "go.mod" -> "go\.mod"
 	//  "go.sum" -> "go\.sum"
 	//  "gotmpl" -> ".*tmpl"
-	FileAssociations map[string]string
+	//  "go.s"   -> ".*\.s"
+	FileAssociations map[protocol.LanguageKind]string
 
 	// Settings holds user-provided configuration for the LSP server.
 	Settings map[string]any
@@ -619,27 +620,28 @@ func (e *Editor) sendDidOpen(ctx context.Context, item protocol.TextDocumentItem
 	return nil
 }
 
-var defaultFileAssociations = map[string]*regexp.Regexp{
+var defaultFileAssociations = map[protocol.LanguageKind]*regexp.Regexp{
 	"go":      regexp.MustCompile(`^.*\.go$`), // '$' is important: don't match .gotmpl!
 	"go.mod":  regexp.MustCompile(`^go\.mod$`),
 	"go.sum":  regexp.MustCompile(`^go(\.work)?\.sum$`),
 	"go.work": regexp.MustCompile(`^go\.work$`),
 	"gotmpl":  regexp.MustCompile(`^.*tmpl$`),
+	"go.s":    regexp.MustCompile(`\.s$`),
 }
 
 // languageID returns the language identifier for the path p given the user
 // configured fileAssociations.
-func languageID(p string, fileAssociations map[string]string) protocol.LanguageKind {
+func languageID(p string, fileAssociations map[protocol.LanguageKind]string) protocol.LanguageKind {
 	base := path.Base(p)
 	for lang, re := range fileAssociations {
 		re := regexp.MustCompile(re)
 		if re.MatchString(base) {
-			return protocol.LanguageKind(lang)
+			return lang
 		}
 	}
 	for lang, re := range defaultFileAssociations {
 		if re.MatchString(base) {
-			return protocol.LanguageKind(lang)
+			return lang
 		}
 	}
 	return ""

gopls/internal/test/marker/testdata/definition/asm.txt

@@ -0,0 +1,33 @@
+This test exercises the Definition request in a Go assembly file.
+
+For now we support only references to package-level symbols defined in
+the same package or a dependency.
+
+Repeatedly jumping to Definition on ff ping-pongs between the Go and
+assembly declarations.
+
+-- go.mod --
+module example.com
+go 1.18
+
+-- a/a.go --
+package a
+
+import _ "fmt"
+import _ "example.com/b"
+
+func ff() //@ loc(ffgo, re"()ff"), def("ff", ffasm)
+
+var _ = ff // pacify unusedfunc analyzer
+
+-- a/asm.s --
+// portable assembly
+
+TEXT ·ff(SB), $16                       //@ loc(ffasm, "ff"), def("ff", ffgo)
+        CALL    example·com∕b·B         //@ def("com", bB)
+        JMP     ·ff                     //@ def("ff", ffgo)
+
+-- b/b.go --
+package b
+
+func B() {} //@ loc(bB, re"()B")

internal/expect/extract.go

@@ -8,7 +8,9 @@ import (
 	"fmt"
 	"go/ast"
 	"go/parser"
+	goscanner "go/scanner"
 	"go/token"
+	"os"
 	"path/filepath"
 	"regexp"
 	"strconv"
@@ -32,21 +34,54 @@ type Identifier string
 // See the package documentation for details about the syntax of those
 // notes.
 func Parse(fset *token.FileSet, filename string, content []byte) ([]*Note, error) {
-	var src any
-	if content != nil {
-		src = content
+	if content == nil {
+		data, err := os.ReadFile(filename)
+		if err != nil {
+			return nil, err
+		}
+		content = data
 	}
+
 	switch filepath.Ext(filename) {
+	case ".s":
+		// The assembler uses a custom scanner,
+		// but the go/scanner package is close
+		// enough: we only want the comments.
+		file := fset.AddFile(filename, -1, len(content))
+		var scan goscanner.Scanner
+		scan.Init(file, content, nil, goscanner.ScanComments)
+
+		var notes []*Note
+		for {
+			pos, tok, lit := scan.Scan()
+			if tok == token.EOF {
+				break
+			}
+			if tok == token.COMMENT {
+				text, adjust := getAdjustedNote(lit)
+				if text == "" {
+					continue
+				}
+				parsed, err := parse(fset, pos+token.Pos(adjust), text)
+				if err != nil {
+					return nil, err
+				}
+				notes = append(notes, parsed...)
+			}
+		}
+		return notes, nil
+
 	case ".go":
-		// TODO: We should write this in terms of the scanner.
+		// TODO: We should write this in terms of the scanner, like the .s case above.
 		// there are ways you can break the parser such that it will not add all the
 		// comments to the ast, which may result in files where the tests are silently
 		// not run.
-		file, err := parser.ParseFile(fset, filename, src, parser.ParseComments|parser.AllErrors|parser.SkipObjectResolution)
+		file, err := parser.ParseFile(fset, filename, content, parser.ParseComments|parser.AllErrors|parser.SkipObjectResolution)
 		if file == nil {
 			return nil, err
 		}
 		return ExtractGo(fset, file)
+
 	case ".mod":
 		file, err := modfile.Parse(filename, content, nil)
 		if err != nil {
@@ -64,6 +99,7 @@ func Parse(fset *token.FileSet, filename string, content []byte) ([]*Note, error
 			note.Pos += token.Pos(f.Base())
 		}
 		return notes, nil
+
 	case ".work":
 		file, err := modfile.ParseWork(filename, content, nil)
 		if err != nil {