go: allow use of go_asm.h assembly header in assembly files (pantsb…

…uild#17611)

As described in pantsbuild#12943, Go assembly files expect an "assembly header" called `go_asm.h` to exist and be available for inclusion via `#include "go_asm.h"`. The assembly header contains macros generated based on constants and type metadata from the package's Go code. The Go backend currently does not request that the compiler generate this header.

This PR splits the `symabis` generation step and assembly step into two different phases:
- `symabis` generation continues to happen _before_ compilation of the package's Go code since it contains API metadata needed by the Go compiler about the assembly code.
- The Go compiler will now generate the `go_asm.h` assembly header and make it available to Go assembly files. The assembly files are now assembled _after_ the package's Go code has been compiled so they can consume `go_asm.h` if need be.

Fixes pantsbuild#12943.

Note: `symabis` and `go_asm.h` serve similar purposes; they differ mainly around which part of a Go package they are metadata for.

src/python/pants/backend/go/util_rules/assembly.py

@@ -9,58 +9,83 @@
 
 from pants.backend.go.util_rules.goroot import GoRoot
 from pants.backend.go.util_rules.sdk import GoSdkProcess, GoSdkToolIDRequest, GoSdkToolIDResult
-from pants.engine.fs import CreateDigest, Digest, FileContent, MergeDigests
+from pants.engine.fs import CreateDigest, Digest, Directory, FileContent, MergeDigests
 from pants.engine.process import FallibleProcessResult
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
 
 
 @dataclass(frozen=True)
-class FallibleAssemblyCompilationResult:
-    result: AssemblyCompilationResult | None
-    exit_code: int = 0
-    stdout: str | None = None
-    stderr: str | None = None
+class GenerateAssemblySymabisRequest:
+    """Generate a `symabis` file with metadata about the assemnbly files for consumption by Go
+    compiler.
+
+    See https://github.com/bazelbuild/rules_go/issues/1893.
+    """
+
+    compilation_input: Digest
+    s_files: tuple[str, ...]
+    dir_path: str
 
 
 @dataclass(frozen=True)
-class AssemblyCompilationResult:
+class GenerateAssemblySymabisResult:
     symabis_digest: Digest
     symabis_path: str
-    assembly_outputs: tuple[tuple[str, Digest], ...]
 
 
 @dataclass(frozen=True)
-class AssemblyCompilationRequest:
-    """Add a `symabis` file for consumption by Go compiler and assemble all `.s` files.
+class FallibleGenerateAssemblySymabisResult:
+    result: GenerateAssemblySymabisResult | None
+    exit_code: int = 0
+    stdout: str | None = None
+    stderr: str | None = None
 
-    See https://github.com/bazelbuild/rules_go/issues/1893.
-    """
 
-    compilation_input: Digest
+@dataclass(frozen=True)
+class AssembleGoAssemblyFilesRequest:
+    """Assemble Go assembly files to object files."""
+
+    input_digest: Digest
     s_files: tuple[str, ...]
     dir_path: str
+    asm_header_path: str | None
     import_path: str
 
 
+@dataclass(frozen=True)
+class AssembleGoAssemblyFilesResult:
+    assembly_outputs: tuple[tuple[str, Digest], ...]
+
+
+@dataclass(frozen=True)
+class FallibleAssembleGoAssemblyFilesResult:
+    result: AssembleGoAssemblyFilesResult | None
+    exit_code: int = 0
+    stdout: str | None = None
+    stderr: str | None = None
+
+
 @rule
-async def setup_assembly_pre_compilation(
-    request: AssemblyCompilationRequest,
+async def generate_go_assembly_symabisfile(
+    request: GenerateAssemblySymabisRequest,
     goroot: GoRoot,
-) -> FallibleAssemblyCompilationResult:
+) -> FallibleGenerateAssemblySymabisResult:
     # From Go tooling comments:
     #
     #   Supply an empty go_asm.h as if the compiler had been run. -symabis parsing is lax enough
     #   that we don't need the actual definitions that would appear in go_asm.h.
     #
     # See https://go-review.googlesource.com/c/go/+/146999/8/src/cmd/go/internal/work/gc.go
-    go_asm_h_digest, asm_tool_id = await MultiGet(
+    obj_dir_path = PurePath(".", request.dir_path, "__obj__")
+    symabis_path = str(obj_dir_path / "symabis")
+    go_asm_h_digest, asm_tool_id, obj_dir_digest = await MultiGet(
         Get(Digest, CreateDigest([FileContent("go_asm.h", b"")])),
         Get(GoSdkToolIDResult, GoSdkToolIDRequest("asm")),
+        Get(Digest, CreateDigest([Directory(str(obj_dir_path))])),
     )
     symabis_input_digest = await Get(
-        Digest, MergeDigests([request.compilation_input, go_asm_h_digest])
+        Digest, MergeDigests([request.compilation_input, go_asm_h_digest, obj_dir_digest])
     )
-    symabis_path = "symabis"
     symabis_result = await Get(
         FallibleProcessResult,
         GoSdkProcess(
@@ -74,7 +99,7 @@ async def setup_assembly_pre_compilation(
                 "-o",
                 symabis_path,
                 "--",
-                *(f"./{request.dir_path}/{name}" for name in request.s_files),
+                *(str(PurePath(".", request.dir_path, s_file)) for s_file in request.s_files),
             ),
             env={
                 "__PANTS_GO_ASM_TOOL_ID": asm_tool_id.tool_id,
@@ -84,37 +109,66 @@ async def setup_assembly_pre_compilation(
         ),
     )
     if symabis_result.exit_code != 0:
-        return FallibleAssemblyCompilationResult(
+        return FallibleGenerateAssemblySymabisResult(
             None, symabis_result.exit_code, symabis_result.stderr.decode("utf-8")
         )
 
+    return FallibleGenerateAssemblySymabisResult(
+        result=GenerateAssemblySymabisResult(
+            symabis_digest=symabis_result.output_digest,
+            symabis_path=symabis_path,
+        ),
+    )
+
+
+@rule
+async def assemble_go_assembly_files(
+    request: AssembleGoAssemblyFilesRequest,
+    goroot: GoRoot,
+) -> FallibleAssembleGoAssemblyFilesResult:
     # On Go 1.19+, the import path must be supplied via the `-p` option to `go tool asm`.
     # See https://go.dev/doc/go1.19#assembler and
     # https://github.com/bazelbuild/rules_go/commit/cde7d7bc27a34547c014369790ddaa95b932d08d (Bazel rules_go).
-    maybe_package_path_args = (
+    maybe_package_import_path_args = (
         ["-p", request.import_path] if goroot.is_compatible_version("1.19") else []
     )
 
+    obj_dir_path = PurePath(".", request.dir_path, "__obj__")
+    asm_tool_id, obj_dir_digest = await MultiGet(
+        Get(GoSdkToolIDResult, GoSdkToolIDRequest("asm")),
+        Get(Digest, CreateDigest([Directory(str(obj_dir_path))])),
+    )
+
+    input_digest = await Get(Digest, MergeDigests([request.input_digest, obj_dir_digest]))
+
+    maybe_asm_header_path_args = (
+        ["-I", str(PurePath(request.asm_header_path).parent)] if request.asm_header_path else []
+    )
+
+    def obj_output_path(s_file: str) -> str:
+        return str(obj_dir_path / PurePath(s_file).with_suffix(".o"))
+
     assembly_results = await MultiGet(
         Get(
             FallibleProcessResult,
             GoSdkProcess(
-                input_digest=request.compilation_input,
+                input_digest=input_digest,
                 command=(
                     "tool",
                     "asm",
                     "-I",
                     os.path.join(goroot.path, "pkg", "include"),
-                    *maybe_package_path_args,
+                    *maybe_asm_header_path_args,
+                    *maybe_package_import_path_args,
                     "-o",
-                    f"./{request.dir_path}/{PurePath(s_file).with_suffix('.o')}",
-                    f"./{request.dir_path}/{s_file}",
+                    obj_output_path(s_file),
+                    str(os.path.normpath(PurePath(".", request.dir_path, s_file))),
                 ),
                 env={
                     "__PANTS_GO_ASM_TOOL_ID": asm_tool_id.tool_id,
                 },
                 description=f"Assemble {s_file} with Go",
-                output_files=(f"./{request.dir_path}/{PurePath(s_file).with_suffix('.o')}",),
+                output_files=(obj_output_path(s_file),),
             ),
         )
         for s_file in request.s_files
@@ -127,17 +181,15 @@ async def setup_assembly_pre_compilation(
         stderr = "\n\n".join(
             result.stderr.decode("utf-8") for result in assembly_results if result.stderr
         )
-        return FallibleAssemblyCompilationResult(None, exit_code, stdout, stderr)
+        return FallibleAssembleGoAssemblyFilesResult(None, exit_code, stdout, stderr)
 
     assembly_outputs = tuple(
-        (f"./{request.dir_path}/{PurePath(s_file).with_suffix('.o')}", result.output_digest)
+        (obj_output_path(s_file), result.output_digest)
         for s_file, result in zip(request.s_files, assembly_results)
     )
 
-    return FallibleAssemblyCompilationResult(
-        AssemblyCompilationResult(
-            symabis_digest=symabis_result.output_digest,
-            symabis_path=symabis_path,
+    return FallibleAssembleGoAssemblyFilesResult(
+        AssembleGoAssemblyFilesResult(
             assembly_outputs=assembly_outputs,
         )
     )

src/python/pants/backend/go/util_rules/assembly_integration_test.py

@@ -156,10 +156,7 @@ def test_build_invalid_package(rule_runner: RuleRunner) -> None:
     result = rule_runner.request(FallibleBuiltGoPackage, [request])
     assert result.output is None
     assert result.exit_code == 1
-    assert (
-        result.stdout
-        == ".//add_amd64.s:1: unexpected EOF\nasm: assembly of .//add_amd64.s failed\n"
-    )
+    assert result.stdout == "add_amd64.s:1: unexpected EOF\nasm: assembly of add_amd64.s failed\n"
 
 
 def test_build_package_with_prebuilt_object_files(rule_runner: RuleRunner) -> None:
@@ -255,3 +252,73 @@ def test_build_package_with_prebuilt_object_files(rule_runner: RuleRunner) -> No
     result = subprocess.run([os.path.join(rule_runner.build_root, "bin")], stdout=subprocess.PIPE)
     assert result.returncode == 0
     assert result.stdout == b"42\n"
+
+
+def test_build_package_using_api_metdata(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            "go.mod": dedent(
+                """\
+                module example.com/assembly
+                go 1.17
+                """
+            ),
+            "main.go": dedent(
+                """\
+                package main
+
+                import "fmt"
+
+                const MagicValueToBeUsedByAssembly int = 42
+
+                func main() {
+                    fmt.Println(add_magic(10))
+                }
+                """
+            ),
+            "add_amd64.go": "package main\nfunc add_magic(x int64) int64",
+            "add_arm64.go": "package main\nfunc add_magic(x int64) int64",
+            "add_amd64.s": dedent(
+                """\
+                #include "textflag.h"  // for NOSPLIT
+                #include "go_asm.h"  // for const_MagicValueToBeUsedByAssembly
+                TEXT ·add_magic(SB),NOSPLIT,$0
+                    MOVQ  x+0(FP), BX
+                    MOVQ  $const_MagicValueToBeUsedByAssembly, BP
+
+                    ADDQ  BP, BX
+                    MOVQ  BX, ret+8(FP)
+                    RET
+                """
+            ),
+            "add_arm64.s": dedent(
+                """\
+                #include "textflag.h"  // for NOSPLIT
+                #include "go_asm.h"  // for const_MagicValueToBeUsedByAssembly
+                TEXT ·add_magic(SB),NOSPLIT,$0
+                    MOVD  x+0(FP), R0
+                    MOVD  $const_MagicValueToBeUsedByAssembly, R1
+
+                    ADD   R1, R0, R0
+                    MOVD  R0, ret+8(FP)
+                    RET
+                """
+            ),
+            "BUILD": dedent(
+                """\
+                go_mod(name="mod")
+                go_package(name="pkg", sources=["*.go", "*.s"])
+                go_binary(name="bin")
+                """
+            ),
+        }
+    )
+
+    binary_tgt = rule_runner.get_target(Address("", target_name="bin"))
+    built_package = build_package(rule_runner, binary_tgt)
+    assert len(built_package.artifacts) == 1
+    assert built_package.artifacts[0].relpath == "bin"
+
+    result = subprocess.run([os.path.join(rule_runner.build_root, "bin")], stdout=subprocess.PIPE)
+    assert result.returncode == 0
+    assert result.stdout == b"52\n"  # should be 10 + the 42 "magic" value