From 15a4b35cb6e148f7f65475bbef1e2dfcee12d928 Mon Sep 17 00:00:00 2001
From: Cameron Martin <cameronmartin123@gmail.com>
Date: Sat, 18 Feb 2023 15:33:04 +0000
Subject: [PATCH] Support all toolchain parameters in module extension

---
 extensions.bzl                    | 31 +++++++++++++++++++++++++++++--
 rust/private/repository_utils.bzl |  3 +++
 rust/repositories.bzl             | 10 +++++-----
 3 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/extensions.bzl b/extensions.bzl
index 9116a4f43a..fda329f1df 100644
--- a/extensions.bzl
+++ b/extensions.bzl
@@ -1,13 +1,40 @@
 "Module extensions for using rules_rust with bzlmod"
 
 load("//rust:repositories.bzl", "rust_register_toolchains")
+load(
+    "//rust/private:repository_utils.bzl",
+    "DEFAULT_EXTRA_TARGET_TRIPLES",
+    "DEFAULT_NIGHTLY_VERSION",
+    "DEFAULT_STATIC_RUST_URL_TEMPLATES",
+)
 
 def _toolchains_impl(ctx):
     mod = ctx.modules[0]
     for toolchain in mod.tags.toolchain:
-        rust_register_toolchains(edition = toolchain.edition, register_toolchains = False)
+        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,
+        )
 
-toolchains_toolchain = tag_class(attrs = {"edition": attr.string()})
+toolchains_toolchain = tag_class(attrs = {
+    "dev_components": attr.bool(default = False),
+    "edition": attr.string(),
+    "allocator_library": attr.string(),
+    "rustfmt_version": attr.string(default = DEFAULT_NIGHTLY_VERSION),
+    "rust_analyzer_version": attr.string(),
+    "sha256s": attr.string_dict(),
+    "extra_target_triples": attr.string_list(default = DEFAULT_EXTRA_TARGET_TRIPLES),
+    "urls": attr.string_list(default = DEFAULT_STATIC_RUST_URL_TEMPLATES),
+    "versions": attr.string_list(default = []),
+})
 toolchains = module_extension(
     implementation = _toolchains_impl,
     tag_classes = {"toolchain": toolchains_toolchain},
diff --git a/rust/private/repository_utils.bzl b/rust/private/repository_utils.bzl
index 6a66e2ba5c..d160b72c20 100644
--- a/rust/private/repository_utils.bzl
+++ b/rust/private/repository_utils.bzl
@@ -1,5 +1,6 @@
 """Utility macros for use in rules_rust repository rules"""
 
+load("//rust/private:common.bzl", "DEFAULT_NIGHTLY_ISO_DATE")
 load("//rust:known_shas.bzl", "FILE_KEY_TO_SHA")
 load(
     "//rust/platform:triple_mappings.bzl",
@@ -11,6 +12,8 @@ load(
 
 DEFAULT_TOOLCHAIN_NAME_PREFIX = "toolchain_for"
 DEFAULT_STATIC_RUST_URL_TEMPLATES = ["https://static.rust-lang.org/dist/{}.tar.gz"]
+DEFAULT_NIGHTLY_VERSION = "nightly/{}".format(DEFAULT_NIGHTLY_ISO_DATE)
+DEFAULT_EXTRA_TARGET_TRIPLES = ["wasm32-unknown-unknown", "wasm32-wasi"]
 
 _build_file_for_compiler_template = """\
 filegroup(
diff --git a/rust/repositories.bzl b/rust/repositories.bzl
index 2c47539c72..6e13c02ed3 100644
--- a/rust/repositories.bzl
+++ b/rust/repositories.bzl
@@ -12,6 +12,8 @@ load(
     "BUILD_for_rust_toolchain",
     "BUILD_for_rustfmt_toolchain",
     "BUILD_for_toolchain",
+    "DEFAULT_EXTRA_TARGET_TRIPLES",
+    "DEFAULT_NIGHTLY_VERSION",
     "DEFAULT_STATIC_RUST_URL_TEMPLATES",
     "check_version_valid",
     "includes_rust_analyzer_proc_macro_srv",
@@ -91,11 +93,9 @@ def rules_rust_dependencies():
         build_file = "@rules_rust//util/process_wrapper:BUILD.tinyjson.bazel",
     )
 
-_DEFAULT_NIGHTLY_VERSION = "nightly/{}".format(DEFAULT_NIGHTLY_ISO_DATE)
-
 _RUST_TOOLCHAIN_VERSIONS = [
     rust_common.default_version,
-    _DEFAULT_NIGHTLY_VERSION,
+    DEFAULT_NIGHTLY_VERSION,
 ]
 
 # buildifier: disable=unnamed-macro
@@ -105,10 +105,10 @@ def rust_register_toolchains(
         allocator_library = None,
         iso_date = None,
         register_toolchains = True,
-        rustfmt_version = _DEFAULT_NIGHTLY_VERSION,
+        rustfmt_version = DEFAULT_NIGHTLY_VERSION,
         rust_analyzer_version = None,
         sha256s = None,
-        extra_target_triples = ["wasm32-unknown-unknown", "wasm32-wasi"],
+        extra_target_triples = DEFAULT_EXTRA_TARGET_TRIPLES,
         urls = DEFAULT_STATIC_RUST_URL_TEMPLATES,
         version = None,
         versions = []):