feat: support treeartifacts (#630) (#631)

Co-authored-by: Sahin Yort <thesayyn@gmail.com>

lib/private/tar.bzl

@@ -1,7 +1,5 @@
 "Implementation of tar rule"
 
-load("@aspect_bazel_lib//lib:paths.bzl", "to_rlocation_path")
-
 _tar_attrs = {
     "args": attr.string_list(
         doc = "Additional flags permitted by BSD tar; see the man page.",
@@ -74,9 +72,6 @@ def _add_compress_options(compress, args):
     if compress == "zstd":
         args.add("--zstd")
 
-def _runfile_path(ctx, file, runfiles_dir):
-    return "/".join([runfiles_dir, to_rlocation_path(ctx, file)])
-
 def _calculate_runfiles_dir(default_info):
     manifest = default_info.files_to_run.runfiles_manifest
 
@@ -124,37 +119,75 @@ def _tar_impl(ctx):
 
     return DefaultInfo(files = depset([out]), runfiles = ctx.runfiles([out]))
 
-def _default_mtree_line(file):
-    # Functions passed to map_each cannot take optional arguments.
-    return _mtree_line(file.short_path, file.path, "dir" if file.is_directory else "file")
-
-def _mtree_line(file, content, type, uid = "0", gid = "0", time = "1672560000", mode = "0755"):
-    return " ".join([
+def _mtree_line(file, type, content = None, uid = "0", gid = "0", time = "1672560000", mode = "0755"):
+    spec = [
         file,
         "uid=" + uid,
         "gid=" + gid,
         "time=" + time,
         "mode=" + mode,
         "type=" + type,
-        "content=" + content,
-    ])
+    ]
+    if content:
+        spec.append("content=" + content)
+    return " ".join(spec)
+
+# This function exactly same as the one from "@aspect_bazel_lib//lib:paths.bzl"
+# except that it takes workspace_name directly instead of the ctx object.
+# Reason is the performance of Args.add_all closures where we use this function.
+# https://bazel.build/rules/lib/builtins/Args#add_all `allow_closure` explains this.
+def _to_rlocation_path(file, workspace):
+    if file.short_path.startswith("../"):
+        return file.short_path[3:]
+    else:
+        return workspace + "/" + file.short_path
+
+def _expand(file, expander, transform = lambda f: f.short_path):
+    expanded = expander.expand(file)
+    lines = []
+    for e in expanded:
+        path = transform(e)
+        segments = path.split("/")
+        for i in range(1, len(segments)):
+            parent = "/".join(segments[:i])
+            lines.append(_mtree_line(parent, "dir"))
+
+        lines.append(_mtree_line(path, "file", content = e.path))
+    return lines
 
 def _mtree_impl(ctx):
     out = ctx.outputs.out or ctx.actions.declare_file(ctx.attr.name + ".spec")
 
     content = ctx.actions.args()
     content.set_param_file_format("multiline")
-    content.add_all(ctx.files.srcs, map_each = _default_mtree_line)
+    content.add_all(
+        ctx.files.srcs,
+        map_each = _expand,
+        expand_directories = True,
+        uniquify = True,
+    )
 
     for s in ctx.attr.srcs:
         default_info = s[DefaultInfo]
         if not default_info.files_to_run.runfiles_manifest:
             continue
 
         runfiles_dir = _calculate_runfiles_dir(default_info)
-        for file in depset(transitive = [s.default_runfiles.files]).to_list():
-            destination = _runfile_path(ctx, file, runfiles_dir)
-            content.add(_mtree_line(destination, file.path, "file"))
+
+        # copy workspace name here just in case to prevent ctx
+        # to be transferred to execution phase.
+        workspace_name = str(ctx.workspace_name)
+
+        content.add(_mtree_line(runfiles_dir, type = "dir"))
+        content.add_all(
+            s.default_runfiles.files,
+            expand_directories = True,
+            uniquify = True,
+            format_each = "{}/%s".format(runfiles_dir),
+            # be careful about what you pass to _expand_for_runfiles as it will carry the data structures over to execution phase.
+            map_each = lambda f, e: _expand(f, e, lambda f: _to_rlocation_path(f, workspace_name)),
+            allow_closure = True,
+        )
 
     ctx.actions.write(out, content = content)
 

lib/tests/tar/BUILD.bazel

@@ -1,3 +1,4 @@
+load("@aspect_bazel_lib//lib:copy_directory.bzl", "copy_directory")
 load("@aspect_bazel_lib//lib:diff_test.bzl", "diff_test")
 load("@aspect_bazel_lib//lib:tar.bzl", "mtree_spec", "tar")
 load("@aspect_bazel_lib//lib:testing.bzl", "assert_archive_contains")
@@ -103,6 +104,9 @@ assert_tar_listing(
     name = "test_flags",
     actual = "tar_flags",
     expected = [
+        "drwxr-xr-x  0 0      0           0 Jan  1  2023 lib/",
+        "drwxr-xr-x  0 0      0           0 Jan  1  2023 lib/tests/",
+        "drwxr-xr-x  0 0      0           0 Jan  1  2023 lib/tests/tar/",
         "-rwxr-xr-x  0 0      0           7 Jan  1  2023 lib/tests/tar/a",
         "-rwxr-xr-x  0 0      0          21 Jan  1  2023 lib/tests/tar/src_file",
     ],
@@ -130,7 +134,15 @@ genrule(
     # lib/tests/tar/a uid=0 gid=0 time=1672560000 mode=0755 type=file content=bazel-out/darwin_arm64-opt/bin/lib/tests/tar/a
     # ->
     # a uid=0 gid=0 time=1672560000 mode=0755 type=file content=bazel-out/darwin_arm64-opt/bin/lib/tests/tar/a
-    cmd = "sed s#^{}/## <$< >$@".format(package_name()),
+    cmd = "sed '{}' <$< | sed '/^\\ /d' > $@".format(
+        "; ".join(reversed([
+            "s#^{s}/##; s#^{s}##".format(s = "/".join(package_name().split("/")[:i]))
+            for (i, _) in enumerate(
+                package_name().split("/"),
+                1,
+            )
+        ])),
+    ),
 )
 
 tar(
@@ -170,6 +182,7 @@ sh_binary(
 tar(
     name = "tar_runfiles",
     srcs = [":cat_src_file"],
+    out = "6.tar",
 )
 
 genrule(
@@ -199,3 +212,34 @@ diff_test(
     file1 = "src_file",
     file2 = "cat_src_file_output",
 )
+
+# Case 7: treeartifacts and source directories
+copy_directory(
+    name = "treeartifact",
+    src = "srcdir",
+    out = "treeartifact",
+)
+
+tar(
+    name = "dirs",
+    srcs = glob(["srcdir/**"]) + [
+        "treeartifact",
+    ],
+    out = "7.tar",
+)
+
+assert_tar_listing(
+    name = "test_dirs",
+    actual = "dirs",
+    expected = [
+        "drwxr-xr-x  0 0      0           0 Jan  1  2023 lib/",
+        "drwxr-xr-x  0 0      0           0 Jan  1  2023 lib/tests/",
+        "drwxr-xr-x  0 0      0           0 Jan  1  2023 lib/tests/tar/",
+        "drwxr-xr-x  0 0      0           0 Jan  1  2023 lib/tests/tar/srcdir/",
+        "-rwxr-xr-x  0 0      0           0 Jan  1  2023 lib/tests/tar/srcdir/info",
+        "-rwxr-xr-x  0 0      0           0 Jan  1  2023 lib/tests/tar/srcdir/pkg",
+        "drwxr-xr-x  0 0      0           0 Jan  1  2023 lib/tests/tar/treeartifact/",
+        "-rwxr-xr-x  0 0      0           0 Jan  1  2023 lib/tests/tar/treeartifact/info",
+        "-rwxr-xr-x  0 0      0           0 Jan  1  2023 lib/tests/tar/treeartifact/pkg",
+    ],
+)