Put all the toolchain proxies into one repository

This does change the public interface a bit, and isn't implemented for rust_analyzer yet, and it breaks non-bzlmod usage currently, but the idea is there.

examples/bzlmod/hello_world/MODULE.bazel

@@ -19,11 +19,4 @@ use_repo(
     "rust_toolchains",
 )
 
-register_toolchains(
-    "@rust_toolchains//:rust_linux_x86_64__x86_64-unknown-linux-gnu",
-    "@rust_toolchains//:rust_darwin_aarch64__aarch64-apple-darwin",
-    "@rust_toolchains//:rust_linux_aarch64__aarch64-unknown-linux-gnu",
-    "@rust_toolchains//:rust_darwin_x86_64__x86_64-apple-darwin",
-    "@rust_toolchains//:rust_windows_x86_64__x86_64-pc-windows-msvc",
-    "@rust_toolchains//:rust_freebsd_x86_64__x86_64-unknown-freebsd",
-)
+register_toolchains("@rust_toolchains//:all")

extensions.bzl

@@ -1,30 +1,19 @@
 "Module extensions for using rules_rust with bzlmod"
 
-load("//rust:repositories.bzl", "rust_register_toolchains", "get_toolchain_repositories", "DEFAULT_TOOLCHAIN_TRIPLES")
+load("//rust:repositories.bzl", "rust_register_toolchains")
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 
 def _toolchains_impl(ctx):
     mod = ctx.modules[0]
     for toolchain in mod.tags.toolchain:
         rust_register_toolchains(edition = toolchain.edition, register_toolchains = False)
-    toolchain_repos = []
-    for exec_triple, name in DEFAULT_TOOLCHAIN_TRIPLES.items():
-        toolchain_repos += get_toolchain_repositories(name = name, exec_triple = exec_triple, extra_target_triples = [])
-    rust_toolchains_repo(name="rust_toolchains", repos=[tr.name for tr in toolchain_repos])
 
 toolchains_toolchain = tag_class(attrs = {"edition": attr.string()})
 toolchains = module_extension(
     implementation = _toolchains_impl,
     tag_classes = {"toolchain": toolchains_toolchain},
 )
 
-
-def _rust_toolchains_repo_impl(ctx):
-  ctx.file("WORKSPACE")
-  ctx.file("BUILD", "\n".join(["alias(name='%s', actual='@%s//:toolchain', visibility = [\"//visibility:public\"])" % (repo,repo) for repo in ctx.attr.repos]))
-rust_toolchains_repo = repository_rule(_rust_toolchains_repo_impl, attrs={"repos": attr.string_list()})
-
-
 def _create_build_file_content(name):
     return """
 load("@rules_rust//rust/private:rust.bzl", "rust_library_without_process_wrapper")

rust/private/repository_utils.bzl

@@ -305,18 +305,18 @@ toolchain(
 """
 
 def BUILD_for_toolchain(
-        name,
-        toolchain,
+        toolchain_names,
+        toolchain_labels,
         toolchain_type,
         target_compatible_with,
         exec_compatible_with):
-    return _build_file_for_toolchain_template.format(
-        name = name,
-        exec_constraint_sets_serialized = exec_compatible_with,
-        target_constraint_sets_serialized = target_compatible_with,
-        toolchain = toolchain,
+    return "\n".join([_build_file_for_toolchain_template.format(
+        name = toolchain_name,
+        exec_constraint_sets_serialized = json.encode(exec_compatible_with[toolchain_name]),
+        target_constraint_sets_serialized = json.encode(target_compatible_with[toolchain_name]),
+        toolchain = toolchain_labels[toolchain_name],
         toolchain_type = toolchain_type,
-    )
+    ) for toolchain_name in toolchain_names])
 
 def load_rustfmt(ctx):
     """Loads a rustfmt binary and yields corresponding BUILD for it

rust/repositories.bzl

@@ -159,6 +159,11 @@ def rust_register_toolchains(
             rust_analyzer_repo_name,
         ))
 
+    toolchain_names = []
+    toolchain_labels = {}
+    exec_compatible_with_by_toolchain = {}
+    target_compatible_with_by_toolchain = {}
+
     for exec_triple, name in DEFAULT_TOOLCHAIN_TRIPLES.items():
         maybe(
             rust_repository_set,
@@ -177,6 +182,22 @@ def rust_register_toolchains(
             version = version,
         )
 
+        for toolchain in get_toolchain_repositories(name, exec_triple, extra_target_triples):
+            toolchain_names.append(toolchain.name)
+            toolchain_labels[toolchain.name] = "@{}//:{}".format(toolchain.name + "_tools", "rust_toolchain")
+            exec_compatible_with_by_toolchain[toolchain.name] = triple_to_constraint_set(exec_triple)
+            target_compatible_with_by_toolchain[toolchain.name] = triple_to_constraint_set(toolchain.target_triple)
+
+
+    toolchain_repository_proxy(
+        name = "rust_toolchains",
+        toolchain_names = toolchain_names,
+        toolchain_labels = toolchain_labels,
+        toolchain_type = "@rules_rust//rust:toolchain",
+        exec_compatible_with = exec_compatible_with_by_toolchain,
+        target_compatible_with = target_compatible_with_by_toolchain,
+    )
+
 # buildifier: disable=unnamed-macro
 def rust_repositories(**kwargs):
     """**Deprecated**: Use [rules_rust_dependencies](#rules_rust_dependencies) \
@@ -312,11 +333,11 @@ def _toolchain_repository_proxy_impl(repository_ctx):
     ))
 
     repository_ctx.file("BUILD.bazel", BUILD_for_toolchain(
-        name = "toolchain",
-        toolchain = repository_ctx.attr.toolchain,
+        toolchain_names = repository_ctx.attr.toolchain_names,
+        toolchain_labels = repository_ctx.attr.toolchain_labels,
         toolchain_type = repository_ctx.attr.toolchain_type,
-        target_compatible_with = json.encode(repository_ctx.attr.target_compatible_with),
-        exec_compatible_with = json.encode(repository_ctx.attr.exec_compatible_with),
+        target_compatible_with = repository_ctx.attr.target_compatible_with,
+        exec_compatible_with =repository_ctx.attr.exec_compatible_with,
     ))
 
 toolchain_repository_proxy = repository_rule(
@@ -325,13 +346,16 @@ toolchain_repository_proxy = repository_rule(
         "rust_toolchain_repository."
     ),
     attrs = {
-        "exec_compatible_with": attr.string_list(
-            doc = "A list of constraints for the execution platform for this toolchain.",
+        "toolchain_names": attr.string_list(mandatory = True),
+        "exec_compatible_with": attr.string_list_dict(
+            doc = "A list of constraints for the execution platform for this toolchain, keyed by toolchain name.",
+            mandatory = True,
         ),
-        "target_compatible_with": attr.string_list(
-            doc = "A list of constraints for the target platform for this toolchain.",
+        "target_compatible_with": attr.string_list_dict(
+            doc = "A list of constraints for the target platform for this toolchain, keyed by toolchain name.",
+            mandatory = True,
         ),
-        "toolchain": attr.string(
+        "toolchain_labels": attr.string_dict(
             doc = "The name of the toolchain implementation target.",
             mandatory = True,
         ),
@@ -412,14 +436,6 @@ def rust_toolchain_repository(
         auth = auth,
     )
 
-    toolchain_repository_proxy(
-        name = name,
-        toolchain = "@{}//:{}".format(name + "_tools", "rust_toolchain"),
-        toolchain_type = "@rules_rust//rust:toolchain",
-        exec_compatible_with = exec_compatible_with,
-        target_compatible_with = target_compatible_with,
-    )
-
 def _rust_analyzer_toolchain_srcs_repository_impl(repository_ctx):
     load_rust_src(repository_ctx)
 
@@ -493,13 +509,14 @@ def rust_analyzer_toolchain_repository(
         auth = auth,
     )
 
-    toolchain_repository_proxy(
-        name = name,
-        toolchain = "@{}//:{}".format(name + "_srcs", "rust_analyzer_toolchain"),
-        toolchain_type = "@rules_rust//rust/rust_analyzer:toolchain_type",
-        exec_compatible_with = exec_compatible_with,
-        target_compatible_with = target_compatible_with,
-    )
+    # TODO: Implement this
+    # toolchain_repository_proxy(
+    #     name = name,
+    #     toolchain = "@{}//:{}".format(name + "_srcs", "rust_analyzer_toolchain"),
+    #     toolchain_type = "@rules_rust//rust/rust_analyzer:toolchain_type",
+    #     exec_compatible_with = exec_compatible_with,
+    #     target_compatible_with = target_compatible_with,
+    # )
 
     return "@{}//:toolchain".format(
         name,