From b6089fb1e516f967bc3f24e9a65c196bed0a9659 Mon Sep 17 00:00:00 2001
From: Griffin Smith <root@gws.fyi>
Date: Sat, 8 Oct 2022 15:12:32 -0400
Subject: [PATCH] feat(tvix/eval): Implement builtins.concatMap

Change-Id: I08bfd040a242aa43b64760c19f48a28303f206ac
Reviewed-on: https://cl.tvl.fyi/c/depot/+/6900
Autosubmit: grfn <grfn@gws.fyi>
Tested-by: BuildkiteCI
Reviewed-by: tazjin <tazjin@tvl.su>
---
 tvix/eval/src/builtins/mod.rs                       | 13 +++++++++++++
 .../src/tests/tvix_tests/eval-okay-concatmap.exp    |  1 +
 .../src/tests/tvix_tests/eval-okay-concatmap.nix    |  1 +
 3 files changed, 15 insertions(+)
 create mode 100644 tvix/eval/src/tests/tvix_tests/eval-okay-concatmap.exp
 create mode 100644 tvix/eval/src/tests/tvix_tests/eval-okay-concatmap.nix

diff --git a/tvix/eval/src/builtins/mod.rs b/tvix/eval/src/builtins/mod.rs
index b859e2844..e047ef599 100644
--- a/tvix/eval/src/builtins/mod.rs
+++ b/tvix/eval/src/builtins/mod.rs
@@ -163,6 +163,19 @@ fn pure_builtins() -> Vec<Builtin> {
                 lists.into_iter().flatten().collect::<Vec<Value>>(),
             )))
         }),
+        Builtin::new(
+            "concatMap",
+            &[true, true],
+            |args: Vec<Value>, vm: &mut VM| {
+                let list = args[1].to_list()?;
+                let mut res = Vec::new();
+                for val in list {
+                    vm.push(val);
+                    res.extend(vm.call_value(&args[0])?.force(vm)?.to_list()?);
+                }
+                Ok(Value::List(res.into()))
+            },
+        ),
         Builtin::new(
             "div",
             &[false, false],
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-concatmap.exp b/tvix/eval/src/tests/tvix_tests/eval-okay-concatmap.exp
new file mode 100644
index 000000000..14d804aa2
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-concatmap.exp
@@ -0,0 +1 @@
+[ "a" "z" "b" "z" ]
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-concatmap.nix b/tvix/eval/src/tests/tvix_tests/eval-okay-concatmap.nix
new file mode 100644
index 000000000..149a0722b
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-concatmap.nix
@@ -0,0 +1 @@
+(builtins.concatMap (x: [x] ++ ["z"]) ["a" "b"])