feat: introduce zstd toolchain (#831)

MODULE.bazel

@@ -21,9 +21,10 @@ bazel_lib_toolchains.jq()
 bazel_lib_toolchains.yq()
 bazel_lib_toolchains.coreutils()
 bazel_lib_toolchains.tar()
+bazel_lib_toolchains.zstd()
 bazel_lib_toolchains.expand_template()
 bazel_lib_toolchains.bats()
-use_repo(bazel_lib_toolchains, "bats_toolchains", "bsd_tar_toolchains", "copy_directory_toolchains", "copy_to_directory_toolchains", "coreutils_toolchains", "expand_template_toolchains", "jq_toolchains", "yq_toolchains")
+use_repo(bazel_lib_toolchains, "bats_toolchains", "bsd_tar_toolchains", "copy_directory_toolchains", "copy_to_directory_toolchains", "coreutils_toolchains", "expand_template_toolchains", "jq_toolchains", "yq_toolchains", "zstd_toolchains")
 
 register_toolchains(
     "@copy_directory_toolchains//:all",

lib/BUILD.bazel

@@ -64,6 +64,10 @@ toolchain_type(
     name = "tar_toolchain_type",
 )
 
+toolchain_type(
+    name = "zstd_toolchain_type",
+)
+
 toolchain_type(
     name = "bats_toolchain_type",
 )
@@ -254,6 +258,7 @@ bzl_library(
         "//lib/private:source_toolchains_repo",
         "//lib/private:tar_toolchain",
         "//lib/private:yq_toolchain",
+        "//lib/private:zstd_toolchain",
     ],
 )
 

lib/extensions.bzl

@@ -14,6 +14,7 @@ load(
     "DEFAULT_TAR_REPOSITORY",
     "DEFAULT_YQ_REPOSITORY",
     "DEFAULT_YQ_VERSION",
+    "DEFAULT_ZSTD_REPOSITORY",
     "register_bats_toolchains",
     "register_copy_directory_toolchains",
     "register_copy_to_directory_toolchains",
@@ -22,6 +23,7 @@ load(
     "register_jq_toolchains",
     "register_tar_toolchains",
     "register_yq_toolchains",
+    "register_zstd_toolchains",
 )
 load("//lib/private:extension_utils.bzl", "extension_utils")
 load("//lib/private:host_repo.bzl", "host_repo")
@@ -89,6 +91,15 @@ def _toolchains_extension_impl(mctx):
         get_version_fn = lambda attr: None,
     )
 
+    extension_utils.toolchain_repos_bfs(
+        mctx = mctx,
+        get_tag_fn = lambda tags: tags.zstd,
+        toolchain_name = "zstd",
+        default_repository = DEFAULT_ZSTD_REPOSITORY,
+        toolchain_repos_fn = lambda name, version: register_zstd_toolchains(name = name, register = False),
+        get_version_fn = lambda attr: None,
+    )
+
     extension_utils.toolchain_repos_bfs(
         mctx = mctx,
         get_tag_fn = lambda tags: tags.expand_template,
@@ -115,6 +126,7 @@ toolchains = module_extension(
         "yq": tag_class(attrs = {"name": attr.string(default = DEFAULT_YQ_REPOSITORY), "version": attr.string(default = DEFAULT_YQ_VERSION)}),
         "coreutils": tag_class(attrs = {"name": attr.string(default = DEFAULT_COREUTILS_REPOSITORY), "version": attr.string(default = DEFAULT_COREUTILS_VERSION)}),
         "tar": tag_class(attrs = {"name": attr.string(default = DEFAULT_TAR_REPOSITORY)}),
+        "zstd": tag_class(attrs = {"name": attr.string(default = DEFAULT_ZSTD_REPOSITORY)}),
         "expand_template": tag_class(attrs = {"name": attr.string(default = DEFAULT_EXPAND_TEMPLATE_REPOSITORY)}),
         "bats": tag_class(attrs = {
             "name": attr.string(default = DEFAULT_BATS_REPOSITORY),

lib/private/BUILD.bazel

@@ -356,3 +356,9 @@ bzl_library(
     srcs = ["strings.bzl"],
     visibility = ["//lib:__subpackages__"],
 )
+
+bzl_library(
+    name = "zstd_toolchain",
+    srcs = ["zstd_toolchain.bzl"],
+    visibility = ["//lib:__subpackages__"],
+)

lib/private/zstd_toolchain.bzl

@@ -0,0 +1,177 @@
+"Provide access to ZSTD"
+
+ZSTD_PLATFORMS = {
+    "darwin_amd64": struct(
+        compatible_with = [
+            "@platforms//os:osx",
+            "@platforms//cpu:x86_64",
+        ],
+    ),
+    "darwin_arm64": struct(
+        compatible_with = [
+            "@platforms//os:osx",
+            "@platforms//cpu:aarch64",
+        ],
+    ),
+    "linux_amd64": struct(
+        compatible_with = [
+            "@platforms//os:linux",
+            "@platforms//cpu:x86_64",
+        ],
+    ),
+    "linux_arm64": struct(
+        compatible_with = [
+            "@platforms//os:linux",
+            "@platforms//cpu:aarch64",
+        ],
+    ),
+}
+
+ZSTD_PREBUILT = {
+    "darwin_amd64": (
+        "https://github.com/aspect-build/zstd-prebuilt/releases/download/v1.5.6/zstd_darwin_amd64",
+        "e4d517212005cf26f8b8d657455d1380318b071cb52a3ffd9dfbdf4c2ba71a13",
+    ),
+    "darwin_arm64": (
+        "https://github.com/aspect-build/zstd-prebuilt/releases/download/v1.5.6/zstd_darwin_arm64",
+        "6e210eeae08fb6ba38c3ac2d1857075c28113aef68296f7e396f1180f7e894b9",
+    ),
+    "linux_amd64": (
+        "https://github.com/aspect-build/zstd-prebuilt/releases/download/v1.5.6/zstd_linux_amd64",
+        "0f0bd1193509a598629d7fa745c4b0b6d5fa6719e0c94c01ef0f20e466d801a7",
+    ),
+    "linux_arm64": (
+        "https://github.com/aspect-build/zstd-prebuilt/releases/download/v1.5.6/zstd_linux_arm64",
+        "82aacf8f1c67ff3c94e04afb0721a848bbba70fbf8249ee4bc4c9085afb84548",
+    ),
+}
+
+def _zstd_binary_repo(rctx):
+    (url, sha256) = ZSTD_PREBUILT[rctx.attr.platform]
+    rctx.download(
+        url = url,
+        output = "zstd",
+        executable = True,
+        sha256 = sha256,
+    )
+    binary = "zstd"
+
+    rctx.file("BUILD.bazel", """\
+# @generated by @aspect_bazel_lib//lib/private:zstd_toolchain.bzl
+
+load("@aspect_bazel_lib//lib/private:zstd_toolchain.bzl", "zstd_toolchain")
+
+package(default_visibility = ["//visibility:public"])
+
+zstd_toolchain(name = "zstd_toolchain", binary = "{}")
+""".format(binary))
+
+zstd_binary_repo = repository_rule(
+    implementation = _zstd_binary_repo,
+    attrs = {
+        "platform": attr.string(mandatory = True, values = ZSTD_PLATFORMS.keys()),
+    },
+)
+
+ZstdInfo = provider(
+    doc = "Provide info for executing zstd",
+    fields = {
+        "binary": "zstd executable",
+    },
+)
+
+def _zstd_toolchain_impl(ctx):
+    binary = ctx.executable.binary
+
+    # Make the $(ZSTD_BIN) variable available in places like genrules.
+    # See https://docs.bazel.build/versions/main/be/make-variables.html#custom_variables
+    template_variables = platform_common.TemplateVariableInfo({
+        "ZSTD_BIN": binary.path,
+    })
+
+    default_info = DefaultInfo(
+        files = depset(ctx.files.binary + ctx.files.files),
+    )
+    zstdinfo = ZstdInfo(
+        binary = binary,
+    )
+
+    # Export all the providers inside our ToolchainInfo
+    # so the resolved_toolchain rule can grab and re-export them.
+    toolchain_info = platform_common.ToolchainInfo(
+        zstdinfo = zstdinfo,
+        template_variables = template_variables,
+        default = default_info,
+    )
+
+    return [toolchain_info, template_variables, default_info]
+
+zstd_toolchain = rule(
+    implementation = _zstd_toolchain_impl,
+    attrs = {
+        "binary": attr.label(
+            doc = "a command to find on the system path",
+            allow_files = True,
+            executable = True,
+            cfg = "exec",
+        ),
+        "files": attr.label_list(allow_files = True),
+    },
+)
+
+def _zstd_toolchains_repo_impl(rctx):
+    # Expose a concrete toolchain which is the result of Bazel resolving the toolchain
+    # for the execution or zstdget platform.
+    # Workaround for https://github.com/bazelbuild/bazel/issues/14009
+    starlark_content = """\
+# @generated by @aspect_bazel_lib//lib/private:zstd_toolchain.bzl
+
+# Forward all the providers
+def _resolved_toolchain_impl(ctx):
+    toolchain_info = ctx.toolchains["@aspect_bazel_lib//lib:zstd_toolchain_type"]
+    return [
+        toolchain_info,
+        toolchain_info.default,
+        toolchain_info.zstdinfo,
+        toolchain_info.template_variables,
+    ]
+
+# Copied from java_toolchain_alias
+# https://cs.opensource.google/bazel/bazel/+/master:tools/jdk/java_toolchain_alias.bzl
+resolved_toolchain = rule(
+    implementation = _resolved_toolchain_impl,
+    toolchains = ["@aspect_bazel_lib//lib:zstd_toolchain_type"],
+    incompatible_use_toolchain_transition = True,
+)
+"""
+    rctx.file("defs.bzl", starlark_content)
+
+    build_content = """# @generated by @aspect_bazel_lib//lib/private:zstd_toolchain.bzl
+load(":defs.bzl", "resolved_toolchain")
+load("@local_config_platform//:constraints.bzl", "HOST_CONSTRAINTS")
+
+resolved_toolchain(name = "resolved_toolchain", visibility = ["//visibility:public"])"""
+
+    for [platform, meta] in ZSTD_PLATFORMS.items():
+        build_content += """
+toolchain(
+    name = "{platform}_toolchain",
+    exec_compatible_with = {compatible_with},
+    toolchain = "@{user_repository_name}_{platform}//:zstd_toolchain",
+    toolchain_type = "@aspect_bazel_lib//lib:zstd_toolchain_type",
+)
+""".format(
+            platform = platform,
+            user_repository_name = rctx.attr.user_repository_name,
+            compatible_with = meta.compatible_with,
+        )
+
+    rctx.file("BUILD.bazel", build_content)
+
+zstd_toolchains_repo = repository_rule(
+    _zstd_toolchains_repo_impl,
+    doc = """Creates a repository that exposes a zstd_toolchain_type zstdget.""",
+    attrs = {
+        "user_repository_name": attr.string(doc = "Base name for toolchains repository"),
+    },
+)

lib/repositories.bzl

@@ -10,6 +10,7 @@ load("//lib/private:jq_toolchain.bzl", "JQ_PLATFORMS", "jq_host_alias_repo", "jq
 load("//lib/private:source_toolchains_repo.bzl", "source_toolchains_repo")
 load("//lib/private:tar_toolchain.bzl", "BSDTAR_PLATFORMS", "bsdtar_binary_repo", "tar_toolchains_repo")
 load("//lib/private:yq_toolchain.bzl", "YQ_PLATFORMS", "yq_host_alias_repo", "yq_platform_repo", "yq_toolchains_repo", _DEFAULT_YQ_VERSION = "DEFAULT_YQ_VERSION")
+load("//lib/private:zstd_toolchain.bzl", "ZSTD_PLATFORMS", "zstd_binary_repo", "zstd_toolchains_repo")
 load("//tools:version.bzl", "IS_PRERELEASE")
 
 # buildifier: disable=unnamed-macro
@@ -104,6 +105,29 @@ def register_tar_toolchains(name = DEFAULT_TAR_REPOSITORY, register = True):
         user_repository_name = name,
     )
 
+DEFAULT_ZSTD_REPOSITORY = "zstd"
+
+def register_zstd_toolchains(name = DEFAULT_ZSTD_REPOSITORY, register = True):
+    """Registers zstd toolchain and repositories
+
+    Args:
+        name: override the prefix for the generated toolchain repositories
+        register: whether to call through to native.register_toolchains.
+            Should be True for WORKSPACE users, but false when used under bzlmod extension
+    """
+    for [platform, _] in ZSTD_PLATFORMS.items():
+        zstd_binary_repo(
+            name = "%s_%s" % (name, platform),
+            platform = platform,
+        )
+        if register:
+            native.register_toolchains("@%s_toolchains//:%s_toolchain" % (name, platform))
+
+    zstd_toolchains_repo(
+        name = "%s_toolchains" % name,
+        user_repository_name = name,
+    )
+
 DEFAULT_BATS_REPOSITORY = "bats"
 
 DEFAULT_BATS_CORE_VERSION = "v1.10.0"
@@ -324,4 +348,5 @@ def aspect_bazel_lib_register_toolchains():
     register_jq_toolchains()
     register_yq_toolchains()
     register_tar_toolchains()
+    register_zstd_toolchains()
     register_bats_toolchains()

lib/tests/zstd/BUILD.bazel

@@ -0,0 +1,30 @@
+load("@aspect_bazel_lib//lib:testing.bzl", "assert_archive_contains")
+
+# Case 1: Can decompress gzip archive
+genrule(
+    name = "tar",
+    srcs = [
+        "srcfile",
+    ],
+    outs = ["1.tar.gz"],
+    cmd = "$(BSDTAR_BIN) --create --gzip --dereference --file $@ -s '#$(BINDIR)##' $(execpath srcfile)",
+    toolchains = ["@bsd_tar_toolchains//:resolved_toolchain"],
+)
+
+genrule(
+    name = "decompress_tar",
+    srcs = [
+        ":tar",
+    ],
+    outs = ["1.tar"],
+    cmd = "$(ZSTD_BIN) -f --decompress $(execpath :tar) --stdout > $@",
+    toolchains = ["@zstd_toolchains//:resolved_toolchain"],
+)
+
+assert_archive_contains(
+    name = "test_decompressed",
+    archive = "1.tar",
+    expected = [
+        "lib/tests/zstd/srcfile",
+    ],
+)