feat(tvix/eval): implement builtins.partition r/5148

Change-Id: I8b591f3057c68c1542046fc5a771973f2238c9df
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7020
Autosubmit: j4m3s <james.landrein@gmail.com>
Reviewed-by: tazjin <tazjin@tvl.su>
Tested-by: BuildkiteCI

3 files changed, 39 insertions, 0 deletions

diff --git a/tvix/eval/src/builtins/mod.rs b/tvix/eval/src/builtins/mod.rs
index 03466af65f..e9ad4a8ee8 100644
--- a/tvix/eval/src/builtins/mod.rs
+++ b/tvix/eval/src/builtins/mod.rs
@@ -447,6 +447,31 @@ fn pure_builtins() -> Vec<Builtin> {
             ]))))
         }),
         Builtin::new(
+            "partition",
+            &[true, true],
+            |args: Vec<Value>, vm: &mut VM| {
+                let mut right: Vec<Value> = vec![];
+                let mut wrong: Vec<Value> = vec![];
+
+                let list: NixList = args[1].to_list()?;
+                for elem in list.into_iter() {
+                    let result = vm.call_with(&args[0], [elem.clone()])?;
+
+                    if result.force(vm)?.as_bool()? {
+                        right.push(elem);
+                    } else {
+                        wrong.push(elem);
+                    };
+                }
+
+                let mut res: BTreeMap<NixString, Value> = BTreeMap::new();
+                res.insert("right".into(), Value::List(right.into()));
+                res.insert("wrong".into(), Value::List(wrong.into()));
+
+                Ok(Value::attrs(NixAttrs::from_map(res)))
+            },
+        ),
+        Builtin::new(
             "removeAttrs",
             &[true, true],
             |args: Vec<Value>, _: &mut VM| {
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-builtins-partition.exp b/tvix/eval/src/tests/tvix_tests/eval-okay-builtins-partition.exp
new file mode 100644
index 0000000000..d2390db4f5
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-builtins-partition.exp
@@ -0,0 +1 @@
+[ { right = [ 1 2 3 4 5 ]; wrong = [ ]; } { right = [ ]; wrong = [ 1 2 3 4 5 ]; } { right = [ 2 ]; wrong = [ 1 3 4 5 ]; } { right = [ [ 1 2 ] [ 3 4 ] ]; wrong = [ [ 1 ] [ 2 ] [ 3 ] ]; } ]
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-builtins-partition.nix b/tvix/eval/src/tests/tvix_tests/eval-okay-builtins-partition.nix
new file mode 100644
index 0000000000..0587330ff9
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-builtins-partition.nix
@@ -0,0 +1,13 @@
+[
+  (builtins.partition (_: true) [ 1 2 3 4 5])
+  (builtins.partition (_: false) [ 1 2 3 4 5])
+  (builtins.partition (x: x == 2) [ 1 2 3 4 5])
+
+  (builtins.partition (x: (builtins.length x) > 1) [
+    [ 1 ]
+    [ 1 2 ]
+    [ 2 ]
+    [ 3 ]
+    [ 3 4 ]
+  ])
+]