From b8ad79081a4e4e736460ab21a83dca916fac6bcb Mon Sep 17 00:00:00 2001
From: Brandon Jones <brandon.jones@arm.com>
Date: Thu, 24 Oct 2024 09:37:22 +0100
Subject: [PATCH] feat: local zig_sdk repo that symlinks to exec based on host

---
 MODULE.bazel         |  4 +++-
 rules/zig_binary.bzl |  4 ++--
 toolchain/defs.bzl   | 41 ++++++++++++++++++++++++++++++++---------
 toolchain/ext.bzl    | 15 ++++++++-------
 4 files changed, 45 insertions(+), 19 deletions(-)

diff --git a/MODULE.bazel b/MODULE.bazel
index 4527b13a..71177594 100644
--- a/MODULE.bazel
+++ b/MODULE.bazel
@@ -30,12 +30,14 @@ use_repo(
 )
 
 toolchains = use_extension("//toolchain:ext.bzl", "toolchains")
-use_repo(toolchains, "zig_sdk-linux-amd64", "zig_sdk-linux-arm64", "zig_sdk-windows-amd64")
+use_repo(toolchains, "zig_sdk" , "zig_sdk-linux-amd64", "zig_sdk-linux-arm64", "zig_sdk-windows-amd64", "zig_sdk-macos-amd64", "zig_sdk-macos-arm64")
 
 _COMMON_EXEC_PLATFORMS = [
     ("linux", "amd64"),
     ("linux", "arm64"),
     ("windows", "amd64"),
+    ("macos", "arm64"),
+    ("macos", "amd64"),
 ]
 
 [register_toolchains(
diff --git a/rules/zig_binary.bzl b/rules/zig_binary.bzl
index 6be40ab6..1a96dd50 100644
--- a/rules/zig_binary.bzl
+++ b/rules/zig_binary.bzl
@@ -32,11 +32,11 @@ zig_binary = rule(
             allow_single_file = [".zig"],
         ),
         "_zig": attr.label(
-            default = "@zig_sdk-linux-amd64//:tools/zig-wrapper",
+            default = "@zig_sdk//:tools/zig-wrapper",
             allow_single_file = True,
         ),
         "_zig_sdk": attr.label(
-            default = "@zig_sdk-linux-amd64//:all",
+            default = "@zig_sdk//:all",
             allow_files = True,
         ),
         "_macos_constraint": attr.label(
diff --git a/toolchain/defs.bzl b/toolchain/defs.bzl
index 5c39407c..4a236da0 100644
--- a/toolchain/defs.bzl
+++ b/toolchain/defs.bzl
@@ -57,12 +57,12 @@ After commenting on the issue, `rm -fr {cache_prefix}` and re-run your command.
 """
 
 def toolchains(
-        exec,
+        exec_os,
+        exec_arch,
         version = VERSION,
         url_formats = [],
         host_platform_sha256 = HOST_PLATFORM_SHA256,
-        host_platform_ext = _HOST_PLATFORM_EXT,
-        ):
+        host_platform_ext = _HOST_PLATFORM_EXT):
     """
         Download zig toolchain and declare bazel toolchains.
         The platforms are not registered automatically, that should be done by
@@ -80,21 +80,21 @@ def toolchains(
         url_formats = [mirror_format, original_format]
 
     zig_repository(
-        name = "zig_sdk-{}".format(exec),
+        name = "zig_sdk-{}-{}".format(exec_os, exec_arch),
         version = version,
         url_formats = url_formats,
         host_platform_sha256 = host_platform_sha256,
         host_platform_ext = host_platform_ext,
-        exec = exec,
+        exec_os = exec_os,
+        exec_arch = exec_arch,
     )
 
 def _quote(s):
     return "'" + s.replace("'", "'\\''") + "'"
 
 def _zig_repository_impl(repository_ctx):
-    arch_os = repository_ctx.attr.exec.split("-")
-    exec_os = arch_os[0]
-    exec_arch = arch_os[1]
+    exec_os = repository_ctx.attr.exec_os
+    exec_arch = repository_ctx.attr.exec_arch
     host_os = repository_ctx.os.name
     host_arch = repository_ctx.os.arch
     if exec_arch == "amd64":
@@ -264,12 +264,35 @@ zig_repository = repository_rule(
         "host_platform_sha256": attr.string_dict(),
         "url_formats": attr.string_list(allow_empty = False),
         "host_platform_ext": attr.string_dict(),
-        "exec": attr.string(),
+        "exec_os": attr.string(),
+        "exec_arch": attr.string(),
     },
     environ = ["HERMETIC_CC_TOOLCHAIN_CACHE_PREFIX"],
     implementation = _zig_repository_impl,
 )
 
+def _host_zig_repository_impl(repository_ctx):
+    host_os = repository_ctx.os.name
+    host_arch = repository_ctx.os.arch
+    if host_arch == "amd64":
+        host_arch = "x86_64"
+
+    if host_os.startswith("mac os"):
+        host_os = "macos"
+
+    if host_os.startswith("windows"):
+        host_os = "windows"
+
+    host_platform_compatible_zig = Label("@zig_sdk-{}-{}//:WORKSPACE".format(host_os, host_arch))
+    compatible_zig_path = repository_ctx.path(host_platform_compatible_zig)
+    repository_ctx.delete(".")
+    repository_ctx.symlink(str(compatible_zig_path) + "/..", ".")
+
+host_zig_repository = repository_rule(
+    implementation = _host_zig_repository_impl,
+    local = True,
+)
+
 def filegroup(name, **kwargs):
     native.filegroup(name = name, **kwargs)
     return ":" + name
diff --git a/toolchain/ext.bzl b/toolchain/ext.bzl
index c6d30eaa..bb994a20 100644
--- a/toolchain/ext.bzl
+++ b/toolchain/ext.bzl
@@ -1,20 +1,21 @@
-load("@hermetic_cc_toolchain//toolchain:defs.bzl", zig_toolchains = "toolchains")
+load("@hermetic_cc_toolchain//toolchain:defs.bzl", "host_zig_repository", zig_toolchains = "toolchains")
 
 _COMMON_EXEC_PLATFORMS = [
     ("linux", "amd64"),
     ("linux", "arm64"),
     ("windows", "amd64"),
+    ("macos", "arm64"),
+    ("macos", "amd64"),
 ]
 
 def _toolchains_impl(ctx):
     for os, arch in _COMMON_EXEC_PLATFORMS:
-        zig_toolchains(exec = "{}-{}".format(os, arch))
-
+        zig_toolchains(exec_os = os, exec_arch = arch)
 
+    host_zig_repository(name = "zig_sdk")
     return ctx.extension_metadata(
-        root_module_direct_deps= ["zig_sdk-{}-{}".format(os, arch) for os, arch in _COMMON_EXEC_PLATFORMS],
-        root_module_direct_dev_deps=[],
-        )
-
+        root_module_direct_deps = ["zig_sdk"] + ["zig_sdk-{}-{}".format(os, arch) for os, arch in _COMMON_EXEC_PLATFORMS],
+        root_module_direct_dev_deps = [],
+    )
 
 toolchains = module_extension(implementation = _toolchains_impl)