add module extension for nodejs toolchain

testing/nodejs/MODULE.bazel

@@ -28,22 +28,16 @@ bazel_dep(name = "bazel_skylib", version = "1.0.3")
 bazel_dep(name = "rules_cc", version = "0.0.4")
 bazel_dep(name = "rules_nodejs", version = "5.5.3")
 
+nix_repo = use_extension("@rules_nixpkgs_core//extensions:repository.bzl", "nix_repo")
+nix_repo.default(name = "nixpkgs")
+use_repo(nix_repo, "nixpkgs")
+
 # TODO[AH] Remove these transitive dependencies once nixpkgs_java_configure has
 #   become a module extension in rules_nixpkgs_java.
 bazel_dep(name = "platforms", version = "0.0.4")
 bazel_dep(name = "rules_java", version = "6.5.2")
 
 non_module_deps = use_extension("//:non_module_deps.bzl", "non_module_deps")
-use_repo(non_module_deps, "nixpkgs")
-[
-    (
-        use_repo(non_module_deps, "nixpkgs_nodejs_{}_{}_toolchain".format(os, arch)),
-        register_toolchains("@nixpkgs_nodejs_{}_{}_toolchain//:all".format(os, arch)),
-    )
-    for os in ["linux", "darwin"]
-    for arch in ["amd64", "arm64"]
-]
-
 use_repo(non_module_deps, "nixpkgs_config_cc")
 use_repo(non_module_deps, "nixpkgs_config_cc_info")
 use_repo(non_module_deps, "nixpkgs_config_cc_toolchains")

testing/nodejs/tests/nixpkgs_repositories.bzl

@@ -4,11 +4,18 @@ load("@rules_nixpkgs_java//:java.bzl", "nixpkgs_java_configure")
 load("@rules_nixpkgs_nodejs//:nodejs.bzl", "nixpkgs_nodejs_configure_platforms")
 
 def nixpkgs_repositories(*, bzlmod):
-    nixpkgs_local_repository(
-        name = "nixpkgs",
-        nix_file = "//:nixpkgs.nix",
-        nix_file_deps = ["//:flake.lock"],
-    )
+    if not bzlmod:
+        nixpkgs_local_repository(
+            name = "nixpkgs",
+            nix_file = "//:nixpkgs.nix",
+            nix_file_deps = ["//:flake.lock"],
+        )
+
+        nixpkgs_nodejs_configure_platforms(
+            name = "nixpkgs_nodejs",
+            repository = "@nixpkgs",
+            register = not bzlmod,
+        )
 
     nixpkgs_cc_configure(
         name = "nixpkgs_config_cc",
@@ -25,9 +32,3 @@ def nixpkgs_repositories(*, bzlmod):
         toolchain_name = "nixpkgs_java",
         toolchain_version = "11",
     )
-
-    nixpkgs_nodejs_configure_platforms(
-        name = "nixpkgs_nodejs",
-        repository = "@nixpkgs",
-        register = not bzlmod,
-    )

toolchains/nodejs/BUILD.bazel

@@ -4,7 +4,10 @@ package(default_visibility = ["//visibility:public"])
 
 filegroup(
     name = "srcs",
-    srcs = glob(["**"]),
+    srcs = glob(["**"]) + [
+        "//extensions:srcs",
+        "//private:srcs",
+    ],
     visibility = ["//visibility:public"],
 )
 
@@ -14,5 +17,6 @@ bzl_library(
     deps = [
         "@rules_nixpkgs_core//:nixpkgs",
         "@rules_nodejs//nodejs:bzl",
+        "//private:common",
     ],
 )

toolchains/nodejs/MODULE.bazel

@@ -4,5 +4,14 @@ module(
 )
 
 bazel_dep(name = "rules_nixpkgs_core", version = "0.11.1")
+bazel_dep(name = "platforms", version = "0.0.4")
 bazel_dep(name = "rules_nodejs", version = "5.5.3")
 bazel_dep(name = "bazel_skylib", version = "1.0.3")
+
+nix_repo = use_extension("@rules_nixpkgs_core//extensions:repository.bzl", "nix_repo")
+nix_repo.default(name = "nixpkgs")
+use_repo(nix_repo, "nixpkgs")
+
+nix_nodejs = use_extension("//extensions:toolchain.bzl", "nix_nodejs")
+use_repo(nix_nodejs, "nixpkgs_nodejs_toolchains")
+register_toolchains("@nixpkgs_nodejs_toolchains//:all")

toolchains/nodejs/extensions/BUILD.bazel

@@ -0,0 +1,13 @@
+load("@bazel_skylib//:bzl_library.bzl", "bzl_library")
+
+filegroup(
+    name = "srcs",
+    srcs = glob(["**"]),
+    visibility = ["//visibility:public"],
+)
+
+bzl_library(
+    name = "toolchain",
+    srcs = ["toolchain.bzl"],
+    visibility = ["//:__pkg__"],
+)

toolchains/nodejs/extensions/toolchain.bzl

@@ -0,0 +1,120 @@
+"""Defines the nix_nodejs module extension.
+"""
+
+load(
+    "//private:common.bzl",
+    "DEFAULT_PLATFORMS_MAPPING",
+    "nixpkgs_nodejs",
+    "nodejs_toolchain",
+)
+
+_DEFAULT_NIXPKGS = "@nixpkgs"
+_DEFAULT_ATTR = "nodejs"
+
+_TOOLCHAINS_REPO = "nixpkgs_nodejs_toolchains"
+_NODEJS_REPO = "nixpkgs_nodejs_{platform}"
+_NODEJS_LABEL = "@{repo}//:nodejs_nix_impl"
+
+def _nodejs_label(repo_name):
+    return _NODEJS_LABEL.format(repo = repo_name)
+
+def _toolchain_name(*, name, count, prefix_digits):
+    prefix = str(count)
+    prefix = "0" * (prefix_digits - len(prefix)) + prefix
+    return prefix + "-" + name
+
+def _toolchains_repo_impl(repository_ctx):
+    num_toolchains = len(repository_ctx.attr.labels)
+    prefix_digits = len(str(num_toolchains))
+
+    sequence = zip(
+        repository_ctx.attr.names,
+        repository_ctx.attr.labels,
+        repository_ctx.attr.exec_lengths,
+        repository_ctx.attr.target_lengths,
+    )
+
+    exec_offset = 0
+    target_offset = 0
+    builder = []
+
+    for count, (name, label, exec_length, target_length) in enumerate(sequence, start = 1):
+        name = _toolchain_name(
+            name = name,
+            count = count,
+            prefix_digits = prefix_digits,
+        )
+        exec_end = exec_offset + exec_length
+        exec_constraints = repository_ctx.attr.exec_constraints[exec_offset:exec_end]
+        exec_offset = exec_end
+        target_end = target_offset + target_length
+        target_constraints = repository_ctx.attr.target_constraints[target_offset:target_end]
+        target_offset = target_end
+        builder.append(nodejs_toolchain(
+            name = name,
+            label = label,
+            exec_constraints = exec_constraints,
+            target_constraints = target_constraints,
+        ))
+
+    repository_ctx.file(
+        "BUILD.bazel",
+        content = "\n".join(builder),
+        executable = False,
+    )
+
+_toolchains_repo = repository_rule(
+    _toolchains_repo_impl,
+    attrs = {
+        "names": attr.string_list(),
+        "labels": attr.string_list(),
+        "exec_constraints": attr.string_list(),
+        "exec_lengths": attr.int_list(),
+        "target_constraints": attr.string_list(),
+        "target_lengths": attr.int_list(),
+    },
+)
+
+def _nix_nodejs_impl(module_ctx):
+    toolchain_names = []
+    toolchain_labels = []
+    toolchain_exec_constraints = []
+    toolchain_exec_lengths = []
+    toolchain_target_constraints = []
+    toolchain_target_lengths = []
+
+    for nix_platform, bazel_platform in DEFAULT_PLATFORMS_MAPPING.items():
+        name = bazel_platform.rules_nodejs_platform
+        repo_name = _NODEJS_REPO.format(platform = bazel_platform.rules_nodejs_platform)
+        exec_constraints = [str(Label(c)) for c in bazel_platform.exec_constraints]
+        target_constraints = [str(Label(c)) for c in bazel_platform.target_constraints]
+
+        nixpkgs_nodejs(
+            name = repo_name,
+            nix_platform = nix_platform,
+            attribute_path = _DEFAULT_ATTR,
+            repository = _DEFAULT_NIXPKGS,
+        )
+
+        toolchain_names.append(name)
+        toolchain_labels.append(_nodejs_label(repo_name))
+        toolchain_exec_constraints.extend(exec_constraints)
+        toolchain_exec_lengths.append(len(exec_constraints))
+        toolchain_target_constraints.extend(target_constraints)
+        toolchain_target_lengths.append(len(target_constraints))
+
+    _toolchains_repo(
+        name = _TOOLCHAINS_REPO,
+        names = toolchain_names,
+        labels = toolchain_labels,
+        exec_constraints = toolchain_exec_constraints,
+        exec_lengths = toolchain_exec_lengths,
+        target_constraints = toolchain_target_constraints,
+        target_lengths = toolchain_target_lengths,
+    )
+
+nix_nodejs = module_extension(
+    _nix_nodejs_impl,
+    tag_classes = {
+    },
+)

toolchains/nodejs/nodejs.bzl

@@ -1,74 +1,20 @@
 load("@rules_nixpkgs_core//:nixpkgs.bzl", "nixpkgs_package")
 load("@rules_nixpkgs_core//:util.bzl", "ensure_constraints")
-load("@rules_nodejs//nodejs/private:toolchains_repo.bzl", "PLATFORMS")
-
-
-def _mk_mapping(rules_nodejs_platform_name):
-    constraints = PLATFORMS[rules_nodejs_platform_name].compatible_with
-    return struct(
-        rules_nodejs_platform = rules_nodejs_platform_name,
-        exec_constraints = constraints,
-        target_constraints = constraints,
-    )
-
-# obtained (and matched) from:
-# nixpkgs search: https://search.nixos.org/packages?channel=22.11&show=nodejs&from=0&size=50&sort=relevance&type=packages&query=nodejs
-# rules_nodejs: https://github.com/bazelbuild/rules_nodejs/blob/a5755eb458c2dd8e0e2cf9b92d8304d9e77ea117/nodejs/private/toolchains_repo.bzl#L20
-DEFAULT_PLATFORMS_MAPPING = {
-  "aarch64-darwin": _mk_mapping("darwin_arm64"),
-  "x86_64-linux": _mk_mapping("linux_amd64"),
-  "x86_64-darwin": _mk_mapping("darwin_amd64"),
-  "aarch64-linux": _mk_mapping("linux_arm64"),
-}
-
-_nodejs_nix_content = """\
-let
-    pkgs = import <nixpkgs> {{ config = {{}}; overlays = []; system = {nix_platform}; }};
-    nodejs = pkgs.{attribute_path};
-in
-pkgs.buildEnv {{
-  extraOutputsToInstall = ["out" "bin" "lib"];
-  name = "bazel-nodejs-toolchain";
-  paths  = [ nodejs ];
-  postBuild = ''
-    touch $out/ROOT
-    cat <<EOF > $out/BUILD
-
-    filegroup(
-        name = "nodejs",
-        srcs = ["bin/node"],
-        visibility = ["//visibility:public"],
-    )
-
-    load("@rules_nodejs//nodejs:toolchain.bzl", "node_toolchain")
-    node_toolchain(
-        name = "nodejs_nix_impl",
-        target_tool = ":nodejs",
-        visibility = ["//visibility:public"],
-    )
-
-    EOF
-  '';
-}}
-"""
-
-_nodejs_nix_toolchain = """
-toolchain(
-    name = "nodejs_nix",
-    toolchain = "@{toolchain_repo}//:nodejs_nix_impl",
-    toolchain_type = "@rules_nodejs//nodejs:toolchain_type",
-    exec_compatible_with = {exec_constraints},
-    target_compatible_with = {target_constraints},
+load(
+    "//private:common.bzl",
+    "DEFAULT_PLATFORMS_MAPPING",
+    "nixpkgs_nodejs",
+    "nodejs_toolchain",
 )
-"""
 
 def _nixpkgs_nodejs_toolchain_impl(repository_ctx):
     exec_constraints, target_constraints = ensure_constraints(repository_ctx)
     repository_ctx.file(
         "BUILD.bazel",
         executable = False,
-        content = _nodejs_nix_toolchain.format(
-            toolchain_repo = repository_ctx.attr.toolchain_repo,
+        content = nodejs_toolchain(
+            name = "nodejs_nix",
+            label = "@{}//:nodejs_nix_impl".format(repository_ctx.attr.toolchain_repo),
             exec_constraints = exec_constraints,
             target_constraints = target_constraints,
         ),
@@ -99,16 +45,10 @@ def nixpkgs_nodejs_configure(
   target_constraints = None,
   register = True,
 ):
-    if attribute_path == None:
-        fail("'attribute_path' is required.", "attribute_path")
-    if not nix_file and not nix_file_content:
-      nix_file_content = _nodejs_nix_content.format(
-        attribute_path = attribute_path,
-        nix_platform = "builtins.currentSystem" if nix_platform == None else repr(nix_platform),
-      )
-
-    nixpkgs_package(
+    nixpkgs_nodejs(
         name = name,
+        nix_platform = nix_platform,
+        attribute_path = attribute_path,
         repository = repository,
         repositories = repositories,
         nix_file = nix_file,

toolchains/nodejs/private/BUILD.bazel

@@ -0,0 +1,13 @@
+load("@bazel_skylib//:bzl_library.bzl", "bzl_library")
+
+filegroup(
+    name = "srcs",
+    srcs = glob(["**"]),
+    visibility = ["//visibility:public"],
+)
+
+bzl_library(
+    name = "common",
+    srcs = ["common.bzl"],
+    visibility = ["//:__pkg__"],
+)