[internal] Refactor and test Go assembly support (#12945)

This makes it much easier to read `build_go_pkg.py`.

This also adds an integration test, although the test does not properly test that our assembly support works...#12944. (I've exceeded my time box for this change.)

[ci skip-rust]
[ci skip-build-wheels]

src/python/pants/backend/experimental/go/register.py

@@ -7,6 +7,7 @@
 from pants.backend.go.subsystems import golang
 from pants.backend.go.target_types import GoBinary, GoExternalPackageTarget, GoModule, GoPackage
 from pants.backend.go.util_rules import (
+    assembly,
     build_go_pkg,
     external_module,
     go_mod,
@@ -22,6 +23,7 @@ def target_types():
 
 def rules():
     return [
+        *assembly.rules(),
         *build_go_pkg.rules(),
         *external_module.rules(),
         *golang.rules(),

src/python/pants/backend/go/goals/package_binary_integration_test.py

@@ -12,6 +12,7 @@
 from pants.backend.go.goals.package_binary import GoBinaryFieldSet
 from pants.backend.go.target_types import GoBinary, GoModule, GoPackage
 from pants.backend.go.util_rules import (
+    assembly,
     build_go_pkg,
     external_module,
     go_mod,
@@ -33,6 +34,7 @@ def rule_runner() -> RuleRunner:
         target_types=[GoBinary, GoPackage, GoModule],
         rules=[
             *external_tool.rules(),
+            *assembly.rules(),
             *source_files.rules(),
             *import_analysis.rules(),
             *package_binary.rules(),

src/python/pants/backend/go/target_types.py

@@ -16,7 +16,7 @@
 
 
 class GoSources(Sources):
-    expected_file_extensions = (".go",)
+    expected_file_extensions = (".go", ".s")
 
 
 # -----------------------------------------------------------------------------------------------
@@ -25,7 +25,7 @@ class GoSources(Sources):
 
 
 class GoPackageSources(GoSources):
-    default = ("*.go",)
+    default = ("*.go", "*.s")
 
 
 class GoImportPath(StringField):

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

@@ -0,0 +1,143 @@
+# Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import PurePath
+
+from pants.backend.go.util_rules.sdk import GoSdkProcess
+from pants.engine.fs import CreateDigest, Digest, FileContent, MergeDigests
+from pants.engine.process import ProcessResult
+from pants.engine.rules import Get, MultiGet, collect_rules, rule
+
+
+@dataclass(frozen=True)
+class AssemblyPreCompilation:
+    merged_compilation_input_digest: Digest
+    assembly_digests: tuple[Digest, ...]
+    EXTRA_COMPILATION_ARGS = ("-symabis", "./symabis")
+
+
+@dataclass(frozen=True)
+class AssemblyPreCompilationRequest:
+    """Add a `symabis` file for consumption by Go compiler and assemble all `.s` files.
+
+    See https://github.com/bazelbuild/rules_go/issues/1893.
+    """
+
+    compilation_input: Digest
+    s_files: tuple[str, ...]
+    source_files_subpath: str
+
+
+@dataclass(frozen=True)
+class AssemblyPostCompilation:
+    merged_output_digest: Digest
+
+
+@dataclass(frozen=True)
+class AssemblyPostCompilationRequest:
+    """Link the assembly_digests into the compilation_result."""
+
+    compilation_result: Digest
+    assembly_digests: tuple[Digest, ...]
+    s_files: tuple[str, ...]
+    source_files_subpath: str
+
+
+@rule
+async def setup_assembly_pre_compilation(
+    request: AssemblyPreCompilationRequest,
+) -> AssemblyPreCompilation:
+    # 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 = await Get(Digest, CreateDigest([FileContent("go_asm.h", b"")]))
+    symabis_input_digest = await Get(
+        Digest, MergeDigests([request.compilation_input, go_asm_h_digest])
+    )
+    symabis_result = await Get(
+        ProcessResult,
+        GoSdkProcess(
+            input_digest=symabis_input_digest,
+            command=(
+                "tool",
+                "asm",
+                "-I",
+                "go/pkg/include",
+                "-gensymabis",
+                "-o",
+                "symabis",
+                "--",
+                *(f"./{request.source_files_subpath}/{name}" for name in request.s_files),
+            ),
+            description="Generate symabis metadata for assembly files.",
+            output_files=("symabis",),
+        ),
+    )
+    merged = await Get(
+        Digest,
+        MergeDigests([request.compilation_input, symabis_result.output_digest]),
+    )
+
+    assembly_results = await MultiGet(
+        Get(
+            ProcessResult,
+            GoSdkProcess(
+                input_digest=request.compilation_input,
+                command=(
+                    "tool",
+                    "asm",
+                    "-I",
+                    "go/pkg/include",
+                    "-o",
+                    f"./{request.source_files_subpath}/{PurePath(s_file).with_suffix('.o')}",
+                    f"./{request.source_files_subpath}/{s_file}",
+                ),
+                description=f"Assemble {s_file}",
+                output_files=(
+                    f"./{request.source_files_subpath}/{PurePath(s_file).with_suffix('.o')}",
+                ),
+            ),
+        )
+        for s_file in request.s_files
+    )
+    return AssemblyPreCompilation(
+        merged, tuple(result.output_digest for result in assembly_results)
+    )
+
+
+@rule
+async def link_assembly_post_compilation(
+    request: AssemblyPostCompilationRequest,
+) -> AssemblyPostCompilation:
+    merged_digest = await Get(
+        Digest, MergeDigests([request.compilation_result, *request.assembly_digests])
+    )
+    pack_result = await Get(
+        ProcessResult,
+        GoSdkProcess(
+            input_digest=merged_digest,
+            command=(
+                "tool",
+                "pack",
+                "r",
+                "__pkg__.a",
+                *(
+                    f"./{request.source_files_subpath}/{PurePath(name).with_suffix('.o')}"
+                    for name in request.s_files
+                ),
+            ),
+            description="Link assembly files to Go package archive.",
+            output_files=("__pkg__.a",),
+        ),
+    )
+    return AssemblyPostCompilation(pack_result.output_digest)
+
+
+def rules():
+    return collect_rules()

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

@@ -0,0 +1,108 @@
+# Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from textwrap import dedent
+
+import pytest
+
+from pants.backend.go import target_type_rules
+from pants.backend.go.target_types import GoExternalPackageTarget, GoModule, GoPackage
+from pants.backend.go.util_rules import (
+    assembly,
+    build_go_pkg,
+    external_module,
+    go_mod,
+    go_pkg,
+    import_analysis,
+    sdk,
+)
+from pants.backend.go.util_rules.build_go_pkg import BuildGoPackageRequest, BuiltGoPackage
+from pants.core.util_rules import external_tool, source_files
+from pants.engine.addresses import Address
+from pants.engine.rules import QueryRule
+from pants.testutil.rule_runner import RuleRunner
+
+
+@pytest.fixture
+def rule_runner() -> RuleRunner:
+    rule_runner = RuleRunner(
+        rules=[
+            *external_tool.rules(),
+            *source_files.rules(),
+            *sdk.rules(),
+            *assembly.rules(),
+            *build_go_pkg.rules(),
+            *import_analysis.rules(),
+            *go_mod.rules(),
+            *go_pkg.rules(),
+            *external_module.rules(),
+            *target_type_rules.rules(),
+            QueryRule(BuiltGoPackage, [BuildGoPackageRequest]),
+        ],
+        target_types=[GoPackage, GoModule, GoExternalPackageTarget],
+    )
+    rule_runner.set_options(["--backend-packages=pants.backend.experimental.go"])
+    return rule_runner
+
+
+def test_build_package_with_assembly(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            "go.mod": dedent(
+                """\
+                module example.com/assembly
+                go 1.16
+                """
+            ),
+            "main.go": dedent(
+                """\
+                package main
+
+                import "fmt"
+
+                func main() {
+                    fmt.Println(add(1, 2))
+                }
+                """
+            ),
+            "add_amd64.go": "package main\nfunc add(x, y int64) int64",
+            "add_arm64.go": "package main\nfunc add(x, y int64) int64",
+            # Based on https://davidwong.fr/goasm/add.
+            "add_amd64.s": dedent(
+                """\
+                TEXT ·add(SB),$0-24
+                    MOVQ  x+0(FP), BX
+                    MOVQ  y+8(FP), BP
+
+                    ADDQ  BP, BX
+                    MOVQ  BX, ret+16(FP)
+                    RET
+                """
+            ),
+            # Based on combining https://davidwong.fr/goasm/add and `go tool compile -S` to get
+            # ARM instructions.
+            "add_arm64.s": dedent(
+                """\
+                TEXT ·add(SB),$0-24
+                    MOVD  x+0(FP), R0
+                    MOVD  y+8(FP), R1
+
+                    ADD   R1, R0, R0
+                    MOVD  R0, ret+16(FP)
+                    RET
+                """
+            ),
+            "BUILD": dedent(
+                """\
+                go_module(name="mod")
+                go_package(name="main")
+                """
+            ),
+        }
+    )
+
+    built_package = rule_runner.request(
+        BuiltGoPackage,
+        [BuildGoPackageRequest(Address("", target_name="main"))],
+    )
+    assert built_package.import_path == "example.com/assembly/"