Support cargo_bazel_bootstrap with bzlmod.

This is required for crate_universe.

MODULE.bazel

@@ -8,10 +8,26 @@ print("WARNING: The rules_rust Bazel module is still highly experimental and sub
 bazel_dep(name = "platforms", version = "0.0.5")
 bazel_dep(name = "rules_cc", version = "0.0.1")
 bazel_dep(name = "bazel_skylib", version = "1.2.0")
-bazel_dep(name = "apple_support", version = "1.3.1")
+bazel_dep(
+    name = "apple_support",
+    version = "1.3.1",
+    repo_name = "build_bazel_apple_support",
+)
 
-internal_deps = use_extension("//rust/private:extensions.bzl", "internal_deps")
+internal_deps = use_extension("//rust/private/module_extensions:internal_deps.bzl", "internal_deps")
 use_repo(
     internal_deps,
     "rules_rust_tinyjson",
 )
+
+# This is the default host tools configuration, but if a user defines it in
+# their repo, it's overridden.
+rust = use_extension("//rust:extensions.bzl", "rust")
+rust.host_tools(edition = "2021")
+use_repo(rust, "rust_host_tools")
+
+cargo_bazel_bootstrap = use_extension("//crate_universe/private/module_extensions:cargo_bazel_bootstrap.bzl", "cargo_bazel_bootstrap")
+use_repo(
+    cargo_bazel_bootstrap,
+    "cargo_bazel_bootstrap",
+)

crate_universe/deps_bootstrap.bzl

@@ -6,12 +6,13 @@ load("//crate_universe/private:srcs.bzl", "CARGO_BAZEL_SRCS")
 # buildifier: disable=bzl-visibility
 load("//rust/private:common.bzl", "rust_common")
 
-def cargo_bazel_bootstrap(name = "cargo_bazel_bootstrap", rust_version = rust_common.default_version):
+def cargo_bazel_bootstrap(name = "cargo_bazel_bootstrap", rust_version = rust_common.default_version, **kwargs):
     """An optional repository which bootstraps `cargo-bazel` for use with `crates_repository`
 
     Args:
         name (str, optional): The name of the `cargo_bootstrap_repository`.
         rust_version (str, optional): The rust version to use. Defaults to the default of `cargo_bootstrap_repository`.
+        **kwargs: kwargs to pass through to cargo_bootstrap_repository.
     """
     cargo_bootstrap_repository(
         name = name,
@@ -22,4 +23,5 @@ def cargo_bazel_bootstrap(name = "cargo_bazel_bootstrap", rust_version = rust_co
         version = rust_version,
         # The increased timeout helps avoid flakes in CI
         timeout = 900,
+        **kwargs
     )

crate_universe/private/module_extensions/cargo_bazel_bootstrap.bzl

@@ -0,0 +1,49 @@
+"""Module extension for bootstrapping cargo-bazel."""
+
+load("//crate_universe:deps_bootstrap.bzl", _cargo_bazel_bootstrap_repo_rule = "cargo_bazel_bootstrap")
+
+def _cargo_bazel_bootstrap_impl(_):
+    _cargo_bazel_bootstrap_repo_rule(
+        rust_toolchain_cargo_template = "@rust_host_tools//:bin/{tool}",
+        rust_toolchain_rustc_template = "@rust_host_tools//:bin/{tool}",
+    )
+
+cargo_bazel_bootstrap = module_extension(
+    implementation = _cargo_bazel_bootstrap_impl,
+)
+
+def get_cargo_bazel_runner(module_ctx):
+    """A helper function to allow executing cargo_bazel in module extensions.
+
+    Args:
+        module_ctx: The module extension's context.
+
+    Returns:
+        A function that can be called to execute cargo_bazel.
+    """
+    cargo_path = str(module_ctx.path(Label("@rust_host_tools//:bin/cargo")))
+    rustc_path = str(module_ctx.path(Label("@rust_host_tools//:bin/rustc")))
+    cargo_bazel = module_ctx.path(Label("@cargo_bazel_bootstrap//:cargo-bazel"))
+
+    def run(args, env = {}, timeout = 600):
+        final_args = [cargo_bazel]
+        final_args.extend(args)
+        final_args.extend([
+            "--cargo",
+            cargo_path,
+            "--rustc",
+            rustc_path,
+        ])
+        result = module_ctx.execute(
+            final_args,
+            environment = dict(CARGO = cargo_path, RUSTC = rustc_path, **env),
+            timeout = timeout,
+        )
+        if result.return_code != 0:
+            if result.stdout:
+                print("Stdout:", result.stdout)  # buildifier: disable=print
+            pretty_args = " ".join([str(arg) for arg in final_args])
+            fail("%s returned with exit code %d:\n%s" % (pretty_args, result.return_code, result.stderr))
+        return result
+
+    return run

rust/extensions.bzl

@@ -1,42 +1,5 @@
 "Module extensions for using rules_rust with bzlmod"
 
-load(
-    "//rust/private:repository_utils.bzl",
-    "DEFAULT_EXTRA_TARGET_TRIPLES",
-    "DEFAULT_NIGHTLY_VERSION",
-    "DEFAULT_STATIC_RUST_URL_TEMPLATES",
-)
-load(":repositories.bzl", "rust_register_toolchains")
+load("//rust/private/module_extensions:toolchain.bzl", _rust = "rust")
 
-def _rust_impl(ctx):
-    mod = ctx.modules[0]
-    for toolchain in mod.tags.toolchain:
-        rust_register_toolchains(
-            dev_components = toolchain.dev_components,
-            edition = toolchain.edition,
-            allocator_library = toolchain.allocator_library,
-            rustfmt_version = toolchain.rustfmt_version,
-            rust_analyzer_version = toolchain.rust_analyzer_version,
-            sha256s = toolchain.sha256s,
-            extra_target_triples = toolchain.extra_target_triples,
-            urls = toolchain.urls,
-            versions = toolchain.versions,
-            register_toolchains = False,
-        )
-
-rust_toolchain = tag_class(attrs = {
-    "allocator_library": attr.string(),
-    "dev_components": attr.bool(default = False),
-    "edition": attr.string(),
-    "extra_target_triples": attr.string_list(default = DEFAULT_EXTRA_TARGET_TRIPLES),
-    "rust_analyzer_version": attr.string(),
-    "rustfmt_version": attr.string(default = DEFAULT_NIGHTLY_VERSION),
-    "sha256s": attr.string_dict(),
-    "urls": attr.string_list(default = DEFAULT_STATIC_RUST_URL_TEMPLATES),
-    "versions": attr.string_list(default = []),
-})
-
-rust = module_extension(
-    implementation = _rust_impl,
-    tag_classes = {"toolchain": rust_toolchain},
-)
+rust = _rust

rust/private/module_extensions/toolchain.bzl

@@ -0,0 +1,77 @@
+"""Module extension for rust toolchain."""
+
+load("//rust:defs.bzl", "rust_common")
+load("//rust:repositories.bzl", "rust_register_toolchains", "rust_toolchain_tools_repository")
+load("//rust/platform:triple.bzl", "get_host_triple")
+load(
+    "//rust/private:repository_utils.bzl",
+    "DEFAULT_EXTRA_TARGET_TRIPLES",
+    "DEFAULT_NIGHTLY_VERSION",
+    "DEFAULT_STATIC_RUST_URL_TEMPLATES",
+)
+
+def _rust_impl(module_ctx):
+    # Allow the root module to define host tools. Otherwise, we'll fall back to
+    # the one defined in rules_rust.
+    for mod in module_ctx.modules:
+        if mod.tags.host_tools:
+            host_tools = mod.tags.host_tools[0]
+            host_triple = get_host_triple(module_ctx)
+
+            rust_toolchain_tools_repository(
+                name = "rust_host_tools",
+                exec_triple = host_triple.str,
+                target_triple = host_triple.str,
+                allocator_library = host_tools.allocator_library,
+                dev_components = host_tools.dev_components,
+                edition = host_tools.edition,
+                rustfmt_version = host_tools.rustfmt_version,
+                sha256s = host_tools.sha256s,
+                urls = host_tools.urls,
+                version = host_tools.version,
+            )
+            break
+
+    mod = module_ctx.modules[0]
+    for toolchain in mod.tags.toolchain:
+        rust_register_toolchains(
+            dev_components = toolchain.dev_components,
+            edition = toolchain.edition,
+            allocator_library = toolchain.allocator_library,
+            rustfmt_version = toolchain.rustfmt_version,
+            rust_analyzer_version = toolchain.rust_analyzer_version,
+            sha256s = toolchain.sha256s,
+            extra_target_triples = toolchain.extra_target_triples,
+            urls = toolchain.urls,
+            versions = toolchain.versions,
+            register_toolchains = False,
+        )
+
+common_kwargs = dict(
+    allocator_library = attr.string(),
+    dev_components = attr.bool(default = False),
+    edition = attr.string(),
+    rustfmt_version = attr.string(default = DEFAULT_NIGHTLY_VERSION),
+    sha256s = attr.string_dict(),
+    urls = attr.string_list(default = DEFAULT_STATIC_RUST_URL_TEMPLATES),
+)
+
+rust_toolchain = tag_class(attrs = dict(
+    extra_target_triples = attr.string_list(default = DEFAULT_EXTRA_TARGET_TRIPLES),
+    rust_analyzer_version = attr.string(),
+    versions = attr.string_list(default = []),
+    **common_kwargs
+))
+
+rust_host_tools = tag_class(attrs = dict(
+    version = attr.string(default = rust_common.default_version),
+    **common_kwargs
+))
+
+rust = module_extension(
+    implementation = _rust_impl,
+    tag_classes = {
+        "host_tools": rust_host_tools,
+        "toolchain": rust_toolchain,
+    },
+)