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

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

@@ -1540,6 +1540,7 @@ def _rustc_invoke(
         compile_cmd.add(cmd_args("--env=", k, "=", v, delimiter = ""))
     for k, v in path_env.items():
         compile_cmd.add(cmd_args("--path-env=", k, "=", v, delimiter = ""))
+    compile_cmd.add(cmd_args(ctx.attrs.env_flags))
 
     build_status = None
     if infallible_diagnostics:

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/buildscript_run.py

@@ -19,6 +19,7 @@
 
 
 IS_WINDOWS: bool = os.name == "nt"
+TOOL_CWD: str = os.getcwd()
 
 
 def eprint(*args: Any, **kwargs: Any) -> None:
@@ -181,7 +182,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 +193,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 +230,36 @@ 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))
+            # Case 1: Output is "cargo:rustc-cfg=VAL"
+            val = cargo_rustc_cfg_match.group(1)
+            rustc_flags += "--cfg={}\n".format(val)
         else:
-            print(line, end="\n")
-    args.outfile.write(flags)
+            cargo_env_flag_match = cargo_env_flag_pattern.match(line)
+            if cargo_env_flag_match:
+                # Case 2: Output is "cargo:rustc-env=KEY=VAL"
+                key = cargo_env_flag_match.group(1)
+                val = cargo_env_flag_match.group(2)
+                prefix = TOOL_CWD + os.path.sep
+                if val.startswith(prefix):
+                    # Case 2a: Output is "cargo:rustc-env=KEY=VAL", where VAL is an
+                    # absolute path to an artifact (under `buck-out`). This for example
+                    # turns `cargo:rustc-env=BIND_PATH=/absolute/path/to/buck-out/.../__target__/OUT_DIR/bind.rs`
+                    # into `--path-env=BIND_PATH=buck-out/.../__target__/OUT_DIR/bind.rs`
+                    env_flags += "--path-env={}={}\n".format(key, val.removeprefix(prefix))
+                else:
+                    # Case 2b: Output is "cargo:rustc-env=KEY=VAL", where VAL is not
+                    # an absolute path to an artifact.
+                    env_flags += "--env={}={}\n".format(key, val)
+            else:
+                print(line, end="\n")
+    args.rustc_flags_outfile.write(rustc_flags)
+    args.env_flags_outfile.write(env_flags)
 
 
 if __name__ == "__main__":