From e5c7ac2815317961752c5656214e95eb3ea0dbc9 Mon Sep 17 00:00:00 2001
From: groot-guo <genyuan.guo@gmail.com>
Date: Sun, 13 Apr 2025 11:28:15 +0800
Subject: [PATCH 1/3] gopls: add CompiledAsmFiles in cache.Package

The CompiledAsmFiles field supports storing assembly files.

Fixes golang/go#71754
---
 gopls/internal/cache/check.go             | 25 +++++++++++++++++------
 gopls/internal/cache/load.go              |  9 ++++++++
 gopls/internal/cache/metadata/metadata.go |  1 +
 gopls/internal/cache/package.go           |  2 ++
 gopls/internal/cache/parse.go             | 21 +++++++++++++++++++
 5 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/gopls/internal/cache/check.go b/gopls/internal/cache/check.go
index 909003288bc..93518886c7e 100644
--- a/gopls/internal/cache/check.go
+++ b/gopls/internal/cache/check.go
@@ -1456,12 +1456,12 @@ type typeCheckInputs struct {
 	id PackageID
 
 	// Used for type checking:
-	pkgPath                  PackagePath
-	name                     PackageName
-	goFiles, compiledGoFiles []file.Handle
-	sizes                    types.Sizes
-	depsByImpPath            map[ImportPath]PackageID
-	goVersion                string // packages.Module.GoVersion, e.g. "1.18"
+	pkgPath                            PackagePath
+	name                               PackageName
+	goFiles, compiledGoFiles, asmFiles []file.Handle
+	sizes                              types.Sizes
+	depsByImpPath                      map[ImportPath]PackageID
+	goVersion                          string // packages.Module.GoVersion, e.g. "1.18"
 
 	// Used for type check diagnostics:
 	// TODO(rfindley): consider storing less data in gobDiagnostics, and
@@ -1491,6 +1491,10 @@ func (s *Snapshot) typeCheckInputs(ctx context.Context, mp *metadata.Package) (*
 	if err != nil {
 		return nil, err
 	}
+	asmFiles, err := readFiles(ctx, s, mp.AsmFiles)
+	if err != nil {
+		return nil, err
+	}
 
 	goVersion := ""
 	if mp.Module != nil && mp.Module.GoVersion != "" {
@@ -1503,6 +1507,7 @@ func (s *Snapshot) typeCheckInputs(ctx context.Context, mp *metadata.Package) (*
 		name:            mp.Name,
 		goFiles:         goFiles,
 		compiledGoFiles: compiledGoFiles,
+		asmFiles:        asmFiles,
 		sizes:           mp.TypesSizes,
 		depsByImpPath:   mp.DepsByImpPath,
 		goVersion:       goVersion,
@@ -1555,6 +1560,10 @@ func localPackageKey(inputs *typeCheckInputs) file.Hash {
 	for _, fh := range inputs.goFiles {
 		fmt.Fprintln(hasher, fh.Identity())
 	}
+	fmt.Fprintf(hasher, "asmFiles:%d\n", len(inputs.asmFiles))
+	for _, fh := range inputs.asmFiles {
+		fmt.Fprintln(hasher, fh.Identity())
+	}
 
 	// types sizes
 	wordSize := inputs.sizes.Sizeof(types.Typ[types.Int])
@@ -1611,6 +1620,10 @@ func (b *typeCheckBatch) checkPackage(ctx context.Context, fset *token.FileSet,
 			pkg.parseErrors = append(pkg.parseErrors, pgf.ParseErr)
 		}
 	}
+	pkg.asmFiles, err = parseAsmFiles(ctx, inputs.asmFiles...)
+	if err != nil {
+		return nil, err
+	}
 
 	// Use the default type information for the unsafe package.
 	if inputs.pkgPath == "unsafe" {
diff --git a/gopls/internal/cache/load.go b/gopls/internal/cache/load.go
index e15e0cef0b6..956ba6d396f 100644
--- a/gopls/internal/cache/load.go
+++ b/gopls/internal/cache/load.go
@@ -454,6 +454,15 @@ func buildMetadata(updates map[PackageID]*metadata.Package, loadDir string, stan
 			*dst = append(*dst, protocol.URIFromPath(filename))
 		}
 	}
+	copyAsmURIs := func(dst *[]protocol.DocumentURI, src []string) {
+		for _, filename := range src {
+			if !strings.HasSuffix(filename, ".s") {
+				continue
+			}
+			*dst = append(*dst, protocol.URIFromPath(filename))
+		}
+	}
+	copyAsmURIs(&mp.AsmFiles, pkg.OtherFiles)
 	copyURIs(&mp.CompiledGoFiles, pkg.CompiledGoFiles)
 	copyURIs(&mp.GoFiles, pkg.GoFiles)
 	copyURIs(&mp.IgnoredFiles, pkg.IgnoredFiles)
diff --git a/gopls/internal/cache/metadata/metadata.go b/gopls/internal/cache/metadata/metadata.go
index 81b6dc57e1f..5408b5a1105 100644
--- a/gopls/internal/cache/metadata/metadata.go
+++ b/gopls/internal/cache/metadata/metadata.go
@@ -48,6 +48,7 @@ type Package struct {
 	// These fields are as defined by go/packages.Package
 	GoFiles         []protocol.DocumentURI
 	CompiledGoFiles []protocol.DocumentURI
+	AsmFiles        []protocol.DocumentURI
 	IgnoredFiles    []protocol.DocumentURI
 	OtherFiles      []protocol.DocumentURI
 
diff --git a/gopls/internal/cache/package.go b/gopls/internal/cache/package.go
index 3477d522cee..370f0959389 100644
--- a/gopls/internal/cache/package.go
+++ b/gopls/internal/cache/package.go
@@ -19,6 +19,7 @@ import (
 	"golang.org/x/tools/gopls/internal/cache/testfuncs"
 	"golang.org/x/tools/gopls/internal/cache/xrefs"
 	"golang.org/x/tools/gopls/internal/protocol"
+	"golang.org/x/tools/gopls/internal/util/asm"
 )
 
 // Convenient aliases for very heavily used types.
@@ -49,6 +50,7 @@ type syntaxPackage struct {
 	fset            *token.FileSet // for now, same as the snapshot's FileSet
 	goFiles         []*parsego.File
 	compiledGoFiles []*parsego.File
+	asmFiles        []*asm.File
 	diagnostics     []*Diagnostic
 	parseErrors     []scanner.ErrorList
 	typeErrors      []types.Error
diff --git a/gopls/internal/cache/parse.go b/gopls/internal/cache/parse.go
index d733ca76799..0f41824d7ea 100644
--- a/gopls/internal/cache/parse.go
+++ b/gopls/internal/cache/parse.go
@@ -13,6 +13,7 @@ import (
 
 	"golang.org/x/tools/gopls/internal/cache/parsego"
 	"golang.org/x/tools/gopls/internal/file"
+	"golang.org/x/tools/gopls/internal/util/asm"
 )
 
 // ParseGo parses the file whose contents are provided by fh.
@@ -43,3 +44,23 @@ func parseGoImpl(ctx context.Context, fset *token.FileSet, fh file.Handle, mode
 	pgf, _ := parsego.Parse(ctx, fset, fh.URI(), content, mode, purgeFuncBodies) // ignore 'fixes'
 	return pgf, nil
 }
+
+// parseAsmFiles pases the assembly files whose contents are provided by fhs.
+func parseAsmFiles(ctx context.Context, fhs ...file.Handle) ([]*asm.File, error) {
+	pgfs := make([]*asm.File, len(fhs))
+
+	for i, fh := range fhs {
+		var err error
+		content, err := fh.Content()
+		if err != nil {
+			return nil, err
+		}
+		// Check for context cancellation before actually doing the parse.
+		if ctx.Err() != nil {
+			return nil, ctx.Err()
+		}
+		pgfs[i] = asm.Parse(content)
+	}
+	return pgfs, nil
+
+}

From 48076194a3a1c2fda562e384f81089b5b3bcaf11 Mon Sep 17 00:00:00 2001
From: groot-guo <genyuan.guo@gmail.com>
Date: Mon, 14 Apr 2025 08:05:12 +0800
Subject: [PATCH 2/3] gopls: add asmFiles in cache.Package

The asmFiles field supports storing assembly files.

Fixes golang/go#71754
---
 gopls/internal/cache/load.go | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/gopls/internal/cache/load.go b/gopls/internal/cache/load.go
index 956ba6d396f..b7d8d671168 100644
--- a/gopls/internal/cache/load.go
+++ b/gopls/internal/cache/load.go
@@ -454,15 +454,13 @@ func buildMetadata(updates map[PackageID]*metadata.Package, loadDir string, stan
 			*dst = append(*dst, protocol.URIFromPath(filename))
 		}
 	}
-	copyAsmURIs := func(dst *[]protocol.DocumentURI, src []string) {
-		for _, filename := range src {
-			if !strings.HasSuffix(filename, ".s") {
-				continue
-			}
-			*dst = append(*dst, protocol.URIFromPath(filename))
+	// Copy SFiles to AsmFiles.
+	for _, filename := range pkg.OtherFiles {
+		if !strings.HasSuffix(filename, ".s") {
+			continue
 		}
+		mp.AsmFiles = append(mp.AsmFiles, protocol.URIFromPath(filename))
 	}
-	copyAsmURIs(&mp.AsmFiles, pkg.OtherFiles)
 	copyURIs(&mp.CompiledGoFiles, pkg.CompiledGoFiles)
 	copyURIs(&mp.GoFiles, pkg.GoFiles)
 	copyURIs(&mp.IgnoredFiles, pkg.IgnoredFiles)

From f8deb4d6267c5d0b095f86a13a8de202b096ff61 Mon Sep 17 00:00:00 2001
From: groot-guo <genyuan.guo@gmail.com>
Date: Thu, 24 Apr 2025 23:19:05 +0800
Subject: [PATCH 3/3] gopls: augmentation of the xrefs index to include
 cross-package references from assembly files.

This commit enhances the xrefs index to include references from assembly (.s) files.
It processes assembly files, extracts cross-package references, and updates the index.

Updates golang/go#71754
---
 gopls/internal/cache/metadata/metadata.go |  4 +-
 gopls/internal/cache/package.go           |  2 +-
 gopls/internal/cache/parse.go             | 10 ++---
 gopls/internal/cache/xrefs/xrefs.go       | 45 +++++++++++++++++++++--
 gopls/internal/goasm/definition.go        |  2 +-
 gopls/internal/util/asm/parse.go          | 13 ++++++-
 gopls/internal/util/asm/parse_test.go     |  2 +-
 7 files changed, 62 insertions(+), 16 deletions(-)

diff --git a/gopls/internal/cache/metadata/metadata.go b/gopls/internal/cache/metadata/metadata.go
index 5408b5a1105..629d25852f6 100644
--- a/gopls/internal/cache/metadata/metadata.go
+++ b/gopls/internal/cache/metadata/metadata.go
@@ -48,10 +48,12 @@ type Package struct {
 	// These fields are as defined by go/packages.Package
 	GoFiles         []protocol.DocumentURI
 	CompiledGoFiles []protocol.DocumentURI
-	AsmFiles        []protocol.DocumentURI
 	IgnoredFiles    []protocol.DocumentURI
 	OtherFiles      []protocol.DocumentURI
 
+	// These fields ares as defined by asm.File
+	AsmFiles []protocol.DocumentURI
+
 	ForTest       PackagePath // q in a "p [q.test]" package, else ""
 	TypesSizes    types.Sizes
 	Errors        []packages.Error          // must be set for packages in import cycles
diff --git a/gopls/internal/cache/package.go b/gopls/internal/cache/package.go
index 370f0959389..71b2a07e3f0 100644
--- a/gopls/internal/cache/package.go
+++ b/gopls/internal/cache/package.go
@@ -71,7 +71,7 @@ type syntaxPackage struct {
 
 func (p *syntaxPackage) xrefs() []byte {
 	p.xrefsOnce.Do(func() {
-		p._xrefs = xrefs.Index(p.compiledGoFiles, p.types, p.typesInfo)
+		p._xrefs = xrefs.Index(p.compiledGoFiles, p.types, p.typesInfo, p.asmFiles)
 	})
 	return p._xrefs
 }
diff --git a/gopls/internal/cache/parse.go b/gopls/internal/cache/parse.go
index 0f41824d7ea..ec94b182997 100644
--- a/gopls/internal/cache/parse.go
+++ b/gopls/internal/cache/parse.go
@@ -45,10 +45,9 @@ func parseGoImpl(ctx context.Context, fset *token.FileSet, fh file.Handle, mode
 	return pgf, nil
 }
 
-// parseAsmFiles pases the assembly files whose contents are provided by fhs.
+// parseAsmFiles parses the assembly files whose contents are provided by fhs.
 func parseAsmFiles(ctx context.Context, fhs ...file.Handle) ([]*asm.File, error) {
-	pgfs := make([]*asm.File, len(fhs))
-
+	pafs := make([]*asm.File, len(fhs))
 	for i, fh := range fhs {
 		var err error
 		content, err := fh.Content()
@@ -59,8 +58,7 @@ func parseAsmFiles(ctx context.Context, fhs ...file.Handle) ([]*asm.File, error)
 		if ctx.Err() != nil {
 			return nil, ctx.Err()
 		}
-		pgfs[i] = asm.Parse(content)
+		pafs[i] = asm.Parse(fh.URI(), content)
 	}
-	return pgfs, nil
-
+	return pafs, nil
 }
diff --git a/gopls/internal/cache/xrefs/xrefs.go b/gopls/internal/cache/xrefs/xrefs.go
index d9b7051737a..33fda837af3 100644
--- a/gopls/internal/cache/xrefs/xrefs.go
+++ b/gopls/internal/cache/xrefs/xrefs.go
@@ -17,13 +17,15 @@ import (
 	"golang.org/x/tools/gopls/internal/cache/metadata"
 	"golang.org/x/tools/gopls/internal/cache/parsego"
 	"golang.org/x/tools/gopls/internal/protocol"
+	"golang.org/x/tools/gopls/internal/util/asm"
 	"golang.org/x/tools/gopls/internal/util/bug"
 	"golang.org/x/tools/gopls/internal/util/frob"
+	"golang.org/x/tools/gopls/internal/util/morestrings"
 )
 
 // Index constructs a serializable index of outbound cross-references
 // for the specified type-checked package.
-func Index(files []*parsego.File, pkg *types.Package, info *types.Info) []byte {
+func Index(files []*parsego.File, pkg *types.Package, info *types.Info, asmFiles []*asm.File) []byte {
 	// pkgObjects maps each referenced package Q to a mapping:
 	// from each referenced symbol in Q to the ordered list
 	// of references to that symbol from this package.
@@ -70,7 +72,7 @@ func Index(files []*parsego.File, pkg *types.Package, info *types.Info) []byte {
 								// (e.g. local const/var/type).
 								continue
 							}
-							gobObj = &gobObject{Path: path}
+							gobObj = &gobObject{Path: path, isAsm: false}
 							objects[obj] = gobObj
 						}
 
@@ -112,6 +114,37 @@ func Index(files []*parsego.File, pkg *types.Package, info *types.Info) []byte {
 		}
 	}
 
+	for fileIndex, af := range asmFiles {
+		for _, id := range af.Idents {
+			_, name, ok := morestrings.CutLast(id.Name, ".")
+			if id.Kind != asm.Text {
+				continue
+			}
+			obj := pkg.Scope().Lookup(name)
+			if obj == nil {
+				continue
+			}
+			objects := getObjects(pkg)
+			gobObj, ok := objects[obj]
+			if !ok {
+				path, err := objectpathFor(obj)
+				if err != nil {
+					// Capitalized but not exported
+					// (e.g. local const/var/type).
+					continue
+				}
+				gobObj = &gobObject{Path: path, isAsm: true}
+				objects[obj] = gobObj
+			}
+			if rng, err := af.NodeRange(id); err == nil {
+				gobObj.Refs = append(gobObj.Refs, gobRef{
+					FileIndex: fileIndex,
+					Range:     rng,
+				})
+			}
+		}
+	}
+
 	// Flatten the maps into slices, and sort for determinism.
 	var packages []*gobPackage
 	for p := range pkgObjects {
@@ -148,6 +181,9 @@ func Lookup(mp *metadata.Package, data []byte, targets map[metadata.PackagePath]
 				if _, ok := objectSet[gobObj.Path]; ok {
 					for _, ref := range gobObj.Refs {
 						uri := mp.CompiledGoFiles[ref.FileIndex]
+						if gobObj.isAsm {
+							uri = mp.AsmFiles[ref.FileIndex]
+						}
 						locs = append(locs, protocol.Location{
 							URI:   uri,
 							Range: ref.Range,
@@ -184,8 +220,9 @@ type gobPackage struct {
 
 // A gobObject records all references to a particular symbol.
 type gobObject struct {
-	Path objectpath.Path // symbol name within package; "" => import of package itself
-	Refs []gobRef        // locations of references within P, in lexical order
+	Path  objectpath.Path // symbol name within package; "" => import of package itself
+	Refs  []gobRef        // locations of references within P, in lexical order
+	isAsm bool            // true if this is an assembly object
 }
 
 type gobRef struct {
diff --git a/gopls/internal/goasm/definition.go b/gopls/internal/goasm/definition.go
index 903916d265d..3eeb0bd7b9c 100644
--- a/gopls/internal/goasm/definition.go
+++ b/gopls/internal/goasm/definition.go
@@ -45,7 +45,7 @@ func Definition(ctx context.Context, snapshot *cache.Snapshot, fh file.Handle, p
 	//
 	// TODO(adonovan): make this just another
 	// attribute of the type-checked cache.Package.
-	file := asm.Parse(content)
+	file := asm.Parse(fh.URI(), content)
 
 	// Figure out the selected symbol.
 	// For now, just find the identifier around the cursor.
diff --git a/gopls/internal/util/asm/parse.go b/gopls/internal/util/asm/parse.go
index 11c59a7cc3d..448a26bc6d2 100644
--- a/gopls/internal/util/asm/parse.go
+++ b/gopls/internal/util/asm/parse.go
@@ -11,6 +11,8 @@ import (
 	"fmt"
 	"strings"
 	"unicode"
+
+	"golang.org/x/tools/gopls/internal/protocol"
 )
 
 // Kind describes the nature of an identifier in an assembly file.
@@ -43,12 +45,19 @@ var kindString = [...]string{
 
 // A file represents a parsed file of Go assembly language.
 type File struct {
+	URI    protocol.DocumentURI
 	Idents []Ident
 
+	Mapper *protocol.Mapper // may map fixed Src, not file content
+
 	// TODO(adonovan): use token.File? This may be important in a
 	// future in which analyzers can report diagnostics in .s files.
 }
 
+func (f *File) NodeRange(ident Ident) (protocol.Range, error) {
+	return f.Mapper.OffsetRange(ident.Offset, ident.End())
+}
+
 // Ident represents an identifier in an assembly file.
 type Ident struct {
 	Name   string // symbol name (after correcting [·∕]); Name[0]='.' => current package
@@ -61,7 +70,7 @@ func (id Ident) End() int { return id.Offset + len(id.Name) }
 
 // Parse extracts identifiers from Go assembly files.
 // Since it is a best-effort parser, it never returns an error.
-func Parse(content []byte) *File {
+func Parse(uri protocol.DocumentURI, content []byte) *File {
 	var idents []Ident
 	offset := 0 // byte offset of start of current line
 
@@ -192,7 +201,7 @@ func Parse(content []byte) *File {
 
 	_ = scan.Err() // ignore scan errors
 
-	return &File{Idents: idents}
+	return &File{Idents: idents, Mapper: protocol.NewMapper(uri, content), URI: uri}
 }
 
 // isIdent reports whether s is a valid Go assembly identifier.
diff --git a/gopls/internal/util/asm/parse_test.go b/gopls/internal/util/asm/parse_test.go
index 67a1286d28b..6d9bd7b8b93 100644
--- a/gopls/internal/util/asm/parse_test.go
+++ b/gopls/internal/util/asm/parse_test.go
@@ -39,7 +39,7 @@ TEXT ·g(SB),NOSPLIT,$0
 `[1:])
 	const filename = "asm.s"
 	m := protocol.NewMapper(protocol.URIFromPath(filename), src)
-	file := asm.Parse(src)
+	file := asm.Parse("", src)
 
 	want := `
 asm.s:5:6-11:	data "hello"