Implement env_flags sub_target in buildscript_run

* Collect build script `cargo:rustc-env=NAME=val` output as `--env=NAME=val` flags
* Export flags as an Artifact, which can be passed to cargo.rust_library/binary with
  `env_flags = ["@$(location :my_crate-0.0.0-build-script-run[env_flags])"],`

Signed-off-by: Yuxuan Dai <[email protected]>

Author: yxdai-nju

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

@@ -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,23 @@ 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:
+                key = cargo_env_flag_match.group(1)
+                val = cargo_env_flag_match.group(2)
+                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__":