feat(tvix/eval): implement 'builtins.filter' r/4963

This is a little ugly because the plain Iterator::filter method can
not be used (it does not support fallible primitives), so we need to
resort to an `Iterator::filter_map` and deal with the wrapping in
Options everywhere.

This prevents use of `?` which introduces the need for some matching,
but it's not *too* bad.

Change-Id: Ie2c3c0c9756c4c627176f64fb4e0054e717c26d1
Reviewed-on: https://cl.tvl.fyi/c/depot/+/6765
Tested-by: BuildkiteCI
Reviewed-by: sterni <sternenseemann@systemli.org>

4 files changed, 44 insertions, 1 deletions

diff --git a/tvix/eval/src/builtins/mod.rs b/tvix/eval/src/builtins/mod.rs
index 872033648a..17094a6f94 100644
--- a/tvix/eval/src/builtins/mod.rs
+++ b/tvix/eval/src/builtins/mod.rs
@@ -132,6 +132,35 @@ fn pure_builtins() -> Vec<Builtin> {
                 }
             }
         }),
+        Builtin::new("filter", &[true, true], |args, vm| {
+            let list: NixList = args[1].to_list()?;
+
+            list.into_iter()
+                .filter_map(|elem| {
+                    vm.push(elem.clone());
+
+                    let result = match vm.call_value(&args[0]) {
+                        Err(err) => return Some(Err(err)),
+                        Ok(result) => result,
+                    };
+
+                    // Must be assigned to a local to avoid a borrowcheck
+                    // failure related to the ForceResult destructor.
+                    let result = match result.force(vm) {
+                        Err(err) => Some(Err(vm.error(err))),
+                        Ok(value) => match value.as_bool() {
+                            Ok(true) => Some(Ok(elem)),
+                            Ok(false) => None,
+                            Err(err) => Some(Err(vm.error(err))),
+                        },
+                    };
+
+                    result
+                })
+                .collect::<Result<Vec<Value>, _>>()
+                .map(|list| Value::List(NixList::from(list)))
+                .map_err(Into::into)
+        }),
         Builtin::new("getAttr", &[true, true], |args, _| {
             let k = args[0].to_str()?;
             let xs = args[1].to_attrs()?;
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-builtins-filter.exp b/tvix/eval/src/tests/tvix_tests/eval-okay-builtins-filter.exp
new file mode 100644
index 0000000000..fb94ebaa49
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-builtins-filter.exp
@@ -0,0 +1 @@
+[ [ 1 2 3 4 5 ] [ ] [ 2 2 2 ] [ [ 1 2 ] [ 3 4 ] ] ]
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-builtins-filter.nix b/tvix/eval/src/tests/tvix_tests/eval-okay-builtins-filter.nix
new file mode 100644
index 0000000000..b621fdb43e
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-builtins-filter.nix
@@ -0,0 +1,13 @@
+[
+  (builtins.filter (_: true) [ 1 2 3 4 5 ])
+  (builtins.filter (_: false) [ 1 2 3 4 5 ])
+  (builtins.filter (x: x == 2) [ 1 2 1 2 1 2 ])
+
+  (builtins.filter (x: (builtins.length x) > 0) [
+    [ ]
+    [ 1 2 ]
+    [ ]
+    [ ]
+    [ 3 4 ]
+  ])
+]
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index d6a24ebf86..547772b0fe 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -175,7 +175,7 @@ impl<'o> VM<'o> {
 
     /// Construct an error from the given ErrorKind and the source
     /// span of the current instruction.
-    fn error(&self, kind: ErrorKind) -> Error {
+    pub fn error(&self, kind: ErrorKind) -> Error {
         Error {
             kind,
             span: self.current_span(),