feat(nix/eval): Implement builtins.functionArgs r/5155

Now that we're tracking formals on Lambda this ends up being quite easy;
we just pull them off of the Lambda for the argument closure and use
them to construct the result attribute set.

Change-Id: I811cb61ec34c6bef123a4043000b18c0e4ea0125
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7003
Reviewed-by: tazjin <tazjin@tvl.su>
Tested-by: BuildkiteCI

4 files changed, 100 insertions, 0 deletions

diff --git a/tvix/eval/src/builtins/mod.rs b/tvix/eval/src/builtins/mod.rs
index a5676a63f895..0b5911de85b6 100644
--- a/tvix/eval/src/builtins/mod.rs
+++ b/tvix/eval/src/builtins/mod.rs
@@ -273,6 +273,21 @@ fn pure_builtins() -> Vec<Builtin> {
                 Ok(res)
             },
         ),
+        Builtin::new("functionArgs", &[true], |args: Vec<Value>, _: &mut VM| {
+            let lambda = args[0].to_closure()?.lambda();
+            let formals = if let Some(formals) = &lambda.formals {
+                formals
+            } else {
+                return Ok(Value::attrs(NixAttrs::empty()));
+            };
+            Ok(Value::attrs(NixAttrs::from_map(
+                formals
+                    .arguments
+                    .iter()
+                    .map(|(k, v)| (k.clone(), (*v).into()))
+                    .collect(),
+            )))
+        }),
         Builtin::new("fromJSON", &[true], |args: Vec<Value>, _: &mut VM| {
             let json_str = args[0].to_str()?;
             let json: serde_json::Value = serde_json::from_str(&json_str)?;
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-functionargs.exp b/tvix/eval/src/tests/tvix_tests/eval-okay-functionargs.exp
new file mode 100644
index 000000000000..c1c9f8ffaf69
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-functionargs.exp
@@ -0,0 +1 @@
+[ "stdenv" "fetchurl" "aterm-stdenv" "aterm-stdenv2" "libX11" "libXv" "mplayer-stdenv2.libXv-libX11" "mplayer-stdenv2.libXv-libX11_2" "nix-stdenv-aterm-stdenv" "nix-stdenv2-aterm2-stdenv2" ]
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-functionargs.nix b/tvix/eval/src/tests/tvix_tests/eval-okay-functionargs.nix
new file mode 100644
index 000000000000..68dca62ee18d
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-functionargs.nix
@@ -0,0 +1,80 @@
+let
+
+  stdenvFun = { }: { name = "stdenv"; };
+  stdenv2Fun = { }: { name = "stdenv2"; };
+  fetchurlFun = { stdenv }: assert stdenv.name == "stdenv"; { name = "fetchurl"; };
+  atermFun = { stdenv, fetchurl }: { name = "aterm-${stdenv.name}"; };
+  aterm2Fun = { stdenv, fetchurl }: { name = "aterm2-${stdenv.name}"; };
+  nixFun = { stdenv, fetchurl, aterm }: { name = "nix-${stdenv.name}-${aterm.name}"; };
+  
+  mplayerFun =
+    { stdenv, fetchurl, enableX11 ? false, xorg ? null, enableFoo ? true, foo ? null  }:
+    assert stdenv.name == "stdenv2";
+    assert enableX11 -> xorg.libXv.name == "libXv";
+    assert enableFoo -> foo != null;
+    { name = "mplayer-${stdenv.name}.${xorg.libXv.name}-${xorg.libX11.name}"; };
+
+  makeOverridable = f: origArgs: f origArgs //
+    { override = newArgs:
+        makeOverridable f (origArgs // (if builtins.isFunction newArgs then newArgs origArgs else newArgs));
+    };
+    
+  callPackage_ = pkgs: f: args:
+    makeOverridable f ((builtins.intersectAttrs (builtins.functionArgs f) pkgs) // args);
+
+  allPackages =
+    { overrides ? (pkgs: pkgsPrev: { }) }:
+    let
+      callPackage = callPackage_ pkgs;
+      pkgs = pkgsStd // (overrides pkgs pkgsStd);
+      pkgsStd = {
+        inherit pkgs;
+        stdenv = callPackage stdenvFun { };
+        stdenv2 = callPackage stdenv2Fun { };
+        fetchurl = callPackage fetchurlFun { };
+        aterm = callPackage atermFun { };
+        xorg = callPackage xorgFun { };
+        mplayer = callPackage mplayerFun { stdenv = pkgs.stdenv2; enableFoo = false; };
+        nix = callPackage nixFun { };
+      };
+    in pkgs;
+
+  libX11Fun = { stdenv, fetchurl }: { name = "libX11"; };
+  libX11_2Fun = { stdenv, fetchurl }: { name = "libX11_2"; };
+  libXvFun = { stdenv, fetchurl, libX11 }: { name = "libXv"; };
+  
+  xorgFun =
+    { pkgs }:
+    let callPackage = callPackage_ (pkgs // pkgs.xorg); in
+    {
+      libX11 = callPackage libX11Fun { };
+      libXv = callPackage libXvFun { };
+    };
+
+in
+
+let
+
+  pkgs = allPackages { };
+  
+  pkgs2 = allPackages {
+    overrides = pkgs: pkgsPrev: {
+      stdenv = pkgs.stdenv2;
+      nix = pkgsPrev.nix.override { aterm = aterm2Fun { inherit (pkgs) stdenv fetchurl; }; };
+      xorg = pkgsPrev.xorg // { libX11 = libX11_2Fun { inherit (pkgs) stdenv fetchurl; }; };
+    };
+  };
+  
+in
+
+  [ pkgs.stdenv.name
+    pkgs.fetchurl.name
+    pkgs.aterm.name
+    pkgs2.aterm.name
+    pkgs.xorg.libX11.name
+    pkgs.xorg.libXv.name
+    pkgs.mplayer.name
+    pkgs2.mplayer.name
+    pkgs.nix.name
+    pkgs2.nix.name
+  ]
diff --git a/tvix/eval/src/value/attrs.rs b/tvix/eval/src/value/attrs.rs
index 6ee3efee679b..a67fd2f3e359 100644
--- a/tvix/eval/src/value/attrs.rs
+++ b/tvix/eval/src/value/attrs.rs
@@ -127,6 +127,10 @@ mod arbitrary {
 }
 
 impl NixAttrs {
+    pub fn empty() -> Self {
+        Self(AttrsRep::Empty)
+    }
+
     /// Return an attribute set containing the merge of the two
     /// provided sets. Keys from the `other` set have precedence.
     pub fn update(self, other: Self) -> Self {