Add support for cargo:rustc-env outputs from build scripts (draft)

prelude/decls/rust_common.bzl

@@ -112,6 +112,15 @@ def _env_arg():
 """),
     }
 
+def _env_flags_arg():
+    return {
+        "env_flags": attrs.list(attrs.arg(), default = [], doc = """
+    A sequence of "--env=NAME=VAL" flags for additional environment variables for this rule's
+     invocations of rustc. For example `env_flags = ["--env=NAME1=val1", "--env=NAME2=val2"]`.
+     The environment variable values may include macros which are expanded.
+"""),
+    }
+
 def _run_env_arg():
     return {
         "run_env": attrs.dict(key = attrs.string(), value = attrs.arg(), sorted = False, default = {}, doc = """
@@ -186,6 +195,7 @@ rust_common = struct(
     crate_root = _crate_root,
     default_roots_arg = _default_roots_arg,
     env_arg = _env_arg,
+    env_flags_arg = _env_flags_arg,
     run_env_arg = _run_env_arg,
     build_and_run_env_arg = _build_and_run_env_arg,
     mapped_srcs_arg = _mapped_srcs_arg,

prelude/decls/rust_rules.bzl

@@ -121,6 +121,7 @@ rust_binary = prelude_rule(
         rust_common.crate(crate_type = attrs.option(attrs.string(), default = None)) |
         rust_common.crate_root() |
         rust_common.env_arg() |
+        rust_common.env_flags_arg() |
         _rust_binary_attrs_group(prefix = "") |
         _rust_common_attributes(is_binary = True) |
         _RUST_EXECUTABLE_ATTRIBUTES |
@@ -186,6 +187,7 @@ rust_library = prelude_rule(
         rust_common.linker_flags_arg() |
         rust_common.exported_linker_flags_arg() |
         rust_common.env_arg() |
+        rust_common.env_flags_arg() |
         rust_common.crate(crate_type = attrs.option(attrs.string(), default = None)) |
         rust_common.crate_root() |
         native_common.preferred_linkage(preferred_linkage_type = attrs.enum(Linkage.values(), default = "any")) |

prelude/rust/build.bzl

@@ -1514,8 +1514,13 @@ def _rustc_invoke(
 
     toolchain_info = compile_ctx.toolchain_info
 
-    plain_env, path_env = process_env(compile_ctx, ctx.attrs.env, exec_is_windows)
+    # Read and process ctx.attrs.env_flags, produce an *-env-flags.txt file
+    # contatining processed environment variable flags.
+    env_flags_file = resolve_env_flags(ctx, compile_ctx, prefix, ctx.attrs.env_flags, exec_is_windows)
 
+    # Process environment variable flags from `ctx.attrs.env`
+    plain_env, path_env = process_env(compile_ctx, ctx.attrs.env, exec_is_windows)
+    # Process environment variable flags from `env`
     more_plain_env, more_path_env = process_env(compile_ctx, env, exec_is_windows)
     plain_env.update(more_plain_env)
     path_env.update(more_path_env)
@@ -1535,6 +1540,9 @@ def _rustc_invoke(
 
     for k, v in crate_map:
         compile_cmd.add(crate_map_arg(k, v))
+    # Feed the environment variable flags in *-env-flags.txt to the command.
+    compile_cmd.add(cmd_args(env_flags_file, format = "@{}"))
+    # Then feed environment variables from `ctx.attrs.env` and `env` to the command.
     for k, v in plain_env.items():
         compile_cmd.add(cmd_args("--env=", k, "=", v, delimiter = ""))
     for k, v in path_env.items():
@@ -1631,6 +1639,69 @@ def _long_command(
         ),
     )
 
+# Takes a list of environment flag as input. Returns an Artifact containing the
+# processed flags, suitable for `_long_command` macro (for rustc and rustdoc).
+def resolve_env_flags(
+        ctx: AnalysisContext,
+        compile_ctx: CompileContext,
+        prefix: str,
+        env_flags: list[str | ResolvedStringWithMacros | Artifact],
+        exec_is_windows: bool,
+        escape_for_rustc_action: bool = True) -> Artifact:
+    # Step 1: In most cases `env_flags = ["@$(location :my_target)"]`, and :my_target
+    #         is a file containing several "--env=NAME=VALUE" flags. This reads
+    #         all such flags and produces an *-env-lines.txt file containing each
+    #         environment variable with the variable name followed by its value.
+    #         For instance, `--env=USER=john --env=HOME=/home/john` (content of :my_target)
+    #         gets converted to:
+    #         ```
+    #         USER
+    #         john
+    #         HOME
+    #         /home/john
+    #         ```
+    env_lines_file = ctx.actions.declare_output("{}-env-lines.txt".format(prefix))
+    cmd = cmd_args(
+        compile_ctx.internal_tools_info.write_env_lines_action,
+        cmd_args(env_lines_file.as_output(), format = "--outfile={}"),
+        cmd_args(env_flags),
+    )
+    ctx.actions.run(cmd, category = "write_env_lines_action", identifier = prefix)
+    env_flags_outfile = ctx.actions.declare_output("{}-env-flags.txt".format(prefix))
+
+    def parse_env_lines_file(ctx: AnalysisContext, artifacts, outputs):
+        # Step 2: Read *-env-lines.txt, then create a mapping from it. For instance
+        #         `{"USER": "john", "HOME": "/home/john"}`.
+        lines = artifacts[env_lines_file].read_string().strip().split("\n")
+        env_map = {k: v for k, v in zip(lines[0::2], lines[1::2])}
+        # Step 3: Utilize 'process_env' to determine whether each environment variable
+        #         is "plain" or "with path", then produce an *-env-flags.txt file
+        #         containing environment variable flags suitable to pass to the
+        #         `_long_command` macro. For instance
+        #         ```
+        #         --env=USER=john
+        #         --path-env=HOME=/home/john
+        #         ```
+        plain_env, path_env = process_env(compile_ctx, env_map, exec_is_windows, escape_for_rustc_action)
+        cmd = cmd_args(
+            compile_ctx.internal_tools_info.write_env_flags_action,
+            cmd_args(outputs[env_flags_outfile].as_output(), format = "--outfile={}"),
+        )
+        for k, v in plain_env.items():
+            cmd.add(cmd_args("--env=", k, "=", v, delimiter = ""))
+        for k, v in path_env.items():
+            cmd.add(cmd_args("--path-env=", k, "=", v, delimiter = ""))
+        ctx.actions.run(cmd, category = "write_env_flags_action", identifier = prefix)
+
+    ctx.actions.dynamic_output(
+        dynamic = [env_lines_file],
+        inputs = [],
+        outputs = [env_flags_outfile.as_output()],
+        f = parse_env_lines_file
+    )
+    # Step 4: Return the Artifact referencing the *-env-flags.txt file.
+    return env_flags_outfile
+
 _DOUBLE_ESCAPED_NEWLINE_RE = regex("\\\\n")
 _ESCAPED_NEWLINE_RE = regex("\\n")
 _DIRECTORY_ENV = [

prelude/rust/cargo_buildscript.bzl

@@ -85,6 +85,7 @@ def _cargo_buildscript_impl(ctx: AnalysisContext) -> list[Provider]:
     cwd = ctx.actions.declare_output("cwd", dir = True)
     out_dir = ctx.actions.declare_output("OUT_DIR", dir = True)
     rustc_flags = ctx.actions.declare_output("rustc_flags")
+    env_flags = ctx.actions.declare_output("env_flags")
 
     if ctx.attrs.manifest_dir != None:
         manifest_dir = ctx.attrs.manifest_dir[DefaultInfo].default_outputs[0]
@@ -97,7 +98,8 @@ def _cargo_buildscript_impl(ctx: AnalysisContext) -> list[Provider]:
         cmd_args("--rustc-cfg=", ctx.attrs.rustc_cfg[DefaultInfo].default_outputs[0], delimiter = ""),
         cmd_args("--manifest-dir=", manifest_dir, delimiter = ""),
         cmd_args("--create-cwd=", cwd.as_output(), delimiter = ""),
-        cmd_args("--outfile=", rustc_flags.as_output(), delimiter = ""),
+        cmd_args("--rustc-flags-outfile=", rustc_flags.as_output(), delimiter = ""),
+        cmd_args("--env-flags-outfile=", env_flags.as_output(), delimiter = ""),
     ]
 
     # See https://doc.rust-lang.org/cargo/reference/environment-variables.html#environment-variables-cargo-sets-for-build-scripts
@@ -144,6 +146,7 @@ def _cargo_buildscript_impl(ctx: AnalysisContext) -> list[Provider]:
         sub_targets = {
             "out_dir": [DefaultInfo(default_output = out_dir)],
             "rustc_flags": [DefaultInfo(default_output = rustc_flags)],
+            "env_flags": [DefaultInfo(default_output = env_flags)],
         },
     )]
 

prelude/rust/tools/BUCK

@@ -59,6 +59,18 @@ prelude.python_bootstrap_binary(
     visibility = ["PUBLIC"],
 )
 
+prelude.python_bootstrap_binary(
+    name = "write_env_flags_action",
+    main = "write_env_flags_action.py",
+    visibility = ["PUBLIC"],
+)
+
+prelude.python_bootstrap_binary(
+    name = "write_env_lines_action",
+    main = "write_env_lines_action.py",
+    visibility = ["PUBLIC"],
+)
+
 prelude.python_bootstrap_binary(
     name = "buildscript_run",
     main = "buildscript_run.py",

prelude/rust/tools/attrs.bzl

@@ -22,6 +22,8 @@ _internal_tool_attrs = {
     "rustdoc_test_with_resources": _internal_tool("prelude//rust/tools:rustdoc_test_with_resources"),
     "symlink_only_dir_entry": _internal_tool("prelude//rust/tools:symlink_only_dir_entry"),
     "transitive_dependency_symlinks_tool": _internal_tool("prelude//rust/tools:transitive_dependency_symlinks"),
+    "write_env_flags_action": _internal_tool("prelude//rust/tools:write_env_flags_action"),
+    "write_env_lines_action": _internal_tool("prelude//rust/tools:write_env_lines_action"),
 }
 
 RustInternalToolsInfo = provider(fields = {

prelude/rust/tools/buildscript_run.py

@@ -181,7 +181,8 @@ class Args(NamedTuple):
     rustc_host_tuple: Optional[Path]
     manifest_dir: Path
     create_cwd: Path
-    outfile: IO[str]
+    rustc_flags_outfile: IO[str]
+    env_flags_outfile: IO[str]
 
 
 def arg_parse() -> Args:
@@ -191,7 +192,8 @@ def arg_parse() -> Args:
     parser.add_argument("--rustc-host-tuple", type=Path)
     parser.add_argument("--manifest-dir", type=Path, required=True)
     parser.add_argument("--create-cwd", type=Path, required=True)
-    parser.add_argument("--outfile", type=argparse.FileType("w"), required=True)
+    parser.add_argument("--rustc-flags-outfile", type=argparse.FileType("w"), required=True)
+    parser.add_argument("--env-flags-outfile", type=argparse.FileType("w"), required=True)
 
     return Args(**vars(parser.parse_args()))
 
@@ -227,14 +229,26 @@ def main() -> None:  # noqa: C901
     script_output = run_buildscript(args.buildscript, env=env, cwd=cwd)
 
     cargo_rustc_cfg_pattern = re.compile("^cargo:rustc-cfg=(.*)")
-    flags = ""
+    cargo_env_flag_pattern = re.compile("^cargo:rustc-env=(.+?)=(.*)")
+    rustc_flags = ""
+    env_flags = ""
     for line in script_output.split("\n"):
         cargo_rustc_cfg_match = cargo_rustc_cfg_pattern.match(line)
         if cargo_rustc_cfg_match:
-            flags += "--cfg={}\n".format(cargo_rustc_cfg_match.group(1))
+            rustc_flags += "--cfg={}\n".format(
+                cargo_rustc_cfg_match.group(1),
+            )
         else:
-            print(line, end="\n")
-    args.outfile.write(flags)
+            cargo_env_flag_match = cargo_env_flag_pattern.match(line)
+            if cargo_env_flag_match:
+                env_flags += "--env={}={}\n".format(
+                    cargo_env_flag_match.group(1),
+                    cargo_env_flag_match.group(2),
+                )
+            else:
+                print(line, end="\n")
+    args.rustc_flags_outfile.write(rustc_flags)
+    args.env_flags_outfile.write(env_flags)
 
 
 if __name__ == "__main__":

prelude/rust/tools/write_env_flags_action.py

@@ -0,0 +1,82 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under both the MIT license found in the
+# LICENSE-MIT file in the root directory of this source tree and the Apache
+# License, Version 2.0 found in the LICENSE-APACHE file in the root directory
+# of this source tree.
+
+"""
+A simple script that accepts environment variable key-value pairs via --env and --path-env flags
+and writes them to an output file in the original command format.
+
+Examples:
+    write_env_flags_action.py --env=USER=john --path-env=HOME=/home/john --outfile output.txt
+
+    output.txt:
+    ```
+    --env=USER=john
+    --path-env=HOME=/home/john
+
+    ```
+"""
+
+import argparse
+import os
+import sys
+
+
+def key_value_arg(s: str):
+    """Parse a key=value argument."""
+    key_value = s.split("=", 1)
+    assert len(key_value) == 2, f"expected the form 'key=value' for '{s}'"
+    return (key_value[0], key_value[1])
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Write command flags to a file in their original format."
+    )
+    parser.add_argument(
+        "--env",
+        action="append",
+        type=key_value_arg,
+        metavar="NAME=VALUE",
+        help="Environment variable in NAME=VALUE format",
+        default=[]
+    )
+    parser.add_argument(
+        "--path-env",
+        action="append",
+        type=key_value_arg,
+        metavar="NAME=VALUE",
+        help="Path environment variable in NAME=VALUE format",
+        default=[]
+    )
+    parser.add_argument(
+        "--outfile",
+        required=True,
+        help="Output file to write commands to"
+    )
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    outfile_dir = os.path.dirname(args.outfile)
+    if outfile_dir:
+        os.makedirs(outfile_dir, exist_ok=True)
+
+    with open(args.outfile, "w") as f:
+        for name, value in args.path_env:
+            f.write(f"--path-env={name}={value}\n")
+        for name, value in args.env:
+            f.write(f"--env={name}={value}\n")
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

prelude/rust/tools/write_env_lines_action.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under both the MIT license found in the
+# LICENSE-MIT file in the root directory of this source tree and the Apache
+# License, Version 2.0 found in the LICENSE-APACHE file in the root directory
+# of this source tree.
+
+"""
+A simple script that accepts environment variable key-value pairs via --env flags
+and writes them to an output file.
+
+Each environment variable and its value are written on separate lines, with the
+variable name followed by its value.
+
+Examples:
+    write_env_lines_action.py --env=USER=john --env=HOME=/home/john --outfile output.txt
+
+    output.txt:
+    ```
+    USER
+    john
+    HOME
+    /home/john
+    
+    ```
+"""
+
+import argparse
+import os
+import sys
+
+
+def key_value_arg(s: str):
+    """Parse a key=value argument."""
+    key_value = s.split("=", 1)
+    assert len(key_value) == 2, f"expected the form 'key=value' for '{s}'"
+    return (key_value[0], key_value[1])
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        fromfile_prefix_chars="@",
+        description="Write environment variables to a file with each name and value on separate lines."
+    )
+    parser.add_argument(
+        "--env",
+        action="append",
+        type=key_value_arg,
+        metavar="NAME=VALUE",
+        help="Environment variable in NAME=VALUE format",
+        default=[]
+    )
+    parser.add_argument(
+        "--outfile",
+        required=True,
+        help="Output file to write environment variables to"
+    )
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    outfile_dir = os.path.dirname(args.outfile)
+    os.makedirs(outfile_dir, exist_ok=True)
+
+    with open(args.outfile, "w") as f:
+        for name, value in args.env:
+            f.write(f"{name}\n{value}\n")
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

prelude/utils/lazy.bzl

@@ -30,7 +30,7 @@ def _is_any(predicate, iterable):
 def _is_all(predicate, iterable):
     """
     This expression lazily iterates the container with 0 new allocations.
-    In the event that the iterable is empty, it will return False.
+    In the event that the iterable is empty, it will return True.
 
     For scenarios like this: