Add wrap extension tag to implement go_wrap_sdk

docs/go/core/bzlmod.md

@@ -104,7 +104,6 @@ go_sdk.nogo(
 ### Not yet supported
 
 -   `go_local_sdk`
--   `go_wrap_sdk`
 
 ## Generating BUILD files
 

go/private/extensions.bzl

@@ -14,7 +14,7 @@
 
 load("@io_bazel_rules_go_bazel_features//:features.bzl", "bazel_features")
 load("//go/private:nogo.bzl", "DEFAULT_NOGO", "NOGO_DEFAULT_EXCLUDES", "NOGO_DEFAULT_INCLUDES", "go_register_nogo")
-load("//go/private:sdk.bzl", "detect_host_platform", "go_download_sdk_rule", "go_host_sdk_rule", "go_multiple_toolchains")
+load("//go/private:sdk.bzl", "detect_host_platform", "go_download_sdk_rule", "go_host_sdk_rule", "go_multiple_toolchains", "go_wrap_sdk_rule")
 
 def host_compatible_toolchain_impl(ctx):
     ctx.file("BUILD.bazel")
@@ -93,6 +93,25 @@ Uses the same format as 'visibility', i.e., every entry must be a label that end
     },
 )
 
+_wrap_tag = tag_class(
+    attrs = {
+        "root_file": attr.label(
+            mandatory = False,
+            doc = "A file in the SDK root directory. Use to determine GOROOT.",
+        ),
+        "root_files": attr.string_dict(
+            mandatory = False,
+            doc = "A set of mappings from the host platform to a file in the SDK's root directory.",
+        ),
+        "version": attr.string(),
+        "experiments": attr.string_list(
+            doc = "Go experiments to enable via GOEXPERIMENT.",
+        ),
+        "goos": attr.string(),
+        "goarch": attr.string(),
+    },
+)
+
 # A list of (goos, goarch) pairs that are commonly used for remote executors in cross-platform
 # builds (where host != exec platform). By default, we register toolchains for all of these
 # platforms in addition to the host platform.
@@ -158,6 +177,30 @@ def _go_sdk_impl(ctx):
     host_detected_goos, host_detected_goarch = detect_host_platform(ctx)
     toolchains = []
     for module in ctx.modules:
+        # Apply wrapped toolchains first to override specific platforms from the
+        # default toolchain or any downloads.
+        for index, wrap_tag in enumerate(module.tags.wrap):
+            name = _default_go_sdk_name(
+                module = module,
+                multi_version = multi_version_module[module.name],
+                tag_type = "wrap",
+                index = index,
+            )
+            go_wrap_sdk_rule(
+                name = name,
+                root_file = wrap_tag.root_file,
+                root_files = wrap_tag.root_files,
+                version = wrap_tag.version,
+                experiments = wrap_tag.experiments,
+            )
+            toolchains.append(struct(
+                goos = wrap_tag.goos,
+                goarch = wrap_tag.goarch,
+                sdk_repo = name,
+                sdk_type = "remote",
+                sdk_version = wrap_tag.version,
+            ))
+
         for index, download_tag in enumerate(module.tags.download):
             # SDKs without an explicit version are fetched even when not selected by toolchain
             # resolution. This is acceptable if brought in by the root module, but transitive
@@ -339,6 +382,7 @@ go_sdk = module_extension(
         "download": _download_tag,
         "host": _host_tag,
         "nogo": _nogo_tag,
+        "wrap": _wrap_tag,
     },
     **go_sdk_extra_kwargs
 )

go/private/sdk.bzl

@@ -381,7 +381,7 @@ def _go_wrap_sdk_impl(ctx):
     _sdk_build_file(ctx, platform, version, ctx.attr.experiments)
     _local_sdk(ctx, goroot)
 
-_go_wrap_sdk = repository_rule(
+go_wrap_sdk_rule = repository_rule(
     implementation = _go_wrap_sdk_impl,
     attrs = {
         "root_file": attr.label(
@@ -405,7 +405,7 @@ _go_wrap_sdk = repository_rule(
 def go_wrap_sdk(name, register_toolchains = True, **kwargs):
     goos = kwargs.pop("goos", None)
     goarch = kwargs.pop("goarch", None)
-    _go_wrap_sdk(name = name, **kwargs)
+    go_wrap_sdk_rule(name = name, **kwargs)
     _go_toolchains(
         name = name + "_toolchains",
         sdk_repo = name,
@@ -646,7 +646,7 @@ def go_register_toolchains(version = None, nogo = None, go_version = None, exper
     if not version:
         version = go_version  # old name
 
-    sdk_kinds = ("go_download_sdk_rule", "go_host_sdk_rule", "_go_local_sdk", "_go_wrap_sdk")
+    sdk_kinds = ("go_download_sdk_rule", "go_host_sdk_rule", "_go_local_sdk", "go_wrap_sdk_rule")
     existing_rules = native.existing_rules()
     sdk_rules = [r for r in existing_rules.values() if r["kind"] in sdk_kinds]
     if len(sdk_rules) == 0 and "go_sdk" in existing_rules:

tests/bcr/.bazelrc

@@ -1 +1,2 @@
 common --enable_bzlmod
+build --@my_rules_go//go/toolchain:sdk_version=1.23.1

tests/bcr/BUILD.bazel

@@ -7,6 +7,7 @@ load(
     "go_test",
     "nogo",
 )
+load("//:transition.bzl", "sdk_transition_test")
 
 nogo(
     name = "my_nogo",
@@ -69,3 +70,15 @@ sh_test(
     env = {"GO_TOOL_RLOCATION": "$(rlocationpath @my_rules_go//go)"},
     deps = ["@bazel_tools//tools/bash/runfiles"],
 )
+
+go_test(
+    name = "wrap_test",
+    srcs = ["wrap_test.go"],
+    tags = ["manual"],
+)
+
+sdk_transition_test(
+    name = "wrap_sdk_test",
+    binary = ":wrap_test",
+    sdk_version = "1.23.6",
+)

tests/bcr/MODULE.bazel

@@ -33,6 +33,9 @@ go_sdk.download(
     ],
     version = "1.23.1",
 )
+go_sdk.wrap(
+    root_file = "@go_default_sdk//:README.md",
+)
 
 # Request an invalid SDK to verify that it isn't fetched since the first tag takes precedence.
 go_sdk.host(version = "3.0.0")

tests/bcr/transition.bzl

@@ -0,0 +1,34 @@
+def _transition_sdk_impl(_, attr):
+    return {"@my_rules_go//go/toolchain:sdk_version": str(attr.sdk_version)}
+
+_transition_sdk = transition(
+    implementation = _transition_sdk_impl,
+    inputs = [],
+    outputs = ["@my_rules_go//go/toolchain:sdk_version"],
+)
+
+def _sdk_transition_impl(ctx):
+    executable = ctx.actions.declare_file(ctx.file.binary.basename)
+    ctx.actions.symlink(
+        output = executable,
+        target_file = ctx.file.binary,
+        is_executable = True,
+    )
+    return DefaultInfo(executable = executable)
+
+sdk_transition_test = rule(
+    _sdk_transition_impl,
+    attrs = {
+        "sdk_version": attr.string(mandatory = True),
+        "binary": attr.label(
+            allow_single_file = True,
+            mandatory = True,
+            cfg = "target",
+        ),
+        "_allowlist_function_transition": attr.label(
+            default = "@bazel_tools//tools/allowlists/function_transition_allowlist",
+        ),
+    },
+    cfg = _transition_sdk,
+    test = True,
+)

tests/bcr/wrap_test.go

@@ -0,0 +1,13 @@
+package main
+
+import (
+	"runtime"
+	"strings"
+	"testing"
+)
+
+func TestSdkVersion(t *testing.T) {
+	if !strings.Contains(runtime.Version(), "1.23.6") {
+		t.Fatal("Incorrect toolchain version", runtime.Version())
+	}
+}