From e2f0967d3fd44cac78ac50425bc2dbe65fd4a8c4 Mon Sep 17 00:00:00 2001
From: Griffin Smith <grfn@gws.fyi>
Date: Tue, 18 Oct 2022 06:32:35 -0400
Subject: [PATCH] feat(nix/eval): Implement builtins.groupBy

Change-Id: I3e0aa017a7100cbeb86d2e5747471b36affcc102
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7038
Autosubmit: grfn <grfn@gws.fyi>
Reviewed-by: tazjin <tazjin@tvl.su>
Tested-by: BuildkiteCI
---
 tvix/eval/src/builtins/mod.rs | 11 +++++++++++
 tvix/eval/src/value/list.rs   |  8 ++++++++
 tvix/eval/src/value/mod.rs    | 15 +++++++++++++++
 3 files changed, 34 insertions(+)

diff --git a/tvix/eval/src/builtins/mod.rs b/tvix/eval/src/builtins/mod.rs
index ecd19213c..2209edd96 100644
--- a/tvix/eval/src/builtins/mod.rs
+++ b/tvix/eval/src/builtins/mod.rs
@@ -323,6 +323,17 @@ fn pure_builtins() -> Vec<Builtin> {
                 }),
             }
         }),
+        Builtin::new("groupBy", &[true, true], |args: Vec<Value>, vm: &mut VM| {
+            let mut res: BTreeMap<NixString, Value> = BTreeMap::new();
+            for val in args[1].to_list()? {
+                let key = vm.call_with(&args[0], [val.clone()])?.force(vm)?.to_str()?;
+                res.entry(key)
+                    .or_insert_with(|| Value::List(NixList::new()))
+                    .as_list_mut()?
+                    .push(val);
+            }
+            Ok(Value::attrs(NixAttrs::from_map(res)))
+        }),
         Builtin::new("hasAttr", &[true, true], |args: Vec<Value>, _: &mut VM| {
             let k = args[0].to_str()?;
             let xs = args[1].to_attrs()?;
diff --git a/tvix/eval/src/value/list.rs b/tvix/eval/src/value/list.rs
index d13e220bb..66f7eb810 100644
--- a/tvix/eval/src/value/list.rs
+++ b/tvix/eval/src/value/list.rs
@@ -49,6 +49,14 @@ mod arbitrary {
 }
 
 impl NixList {
+    pub fn new() -> Self {
+        Self(vec![])
+    }
+
+    pub fn push(&mut self, val: Value) {
+        self.0.push(val)
+    }
+
     pub fn concat(&self, other: &Self) -> Self {
         let mut ret = self.clone();
         ret.0.extend_from_slice(&other.0);
diff --git a/tvix/eval/src/value/mod.rs b/tvix/eval/src/value/mod.rs
index c065d5c1b..14b1a5c61 100644
--- a/tvix/eval/src/value/mod.rs
+++ b/tvix/eval/src/value/mod.rs
@@ -79,6 +79,19 @@ macro_rules! gen_cast {
     };
 }
 
+/// Generate an `as_*_mut/to_*_mut` accessor method that returns either the
+/// expected type, or a type error.
+macro_rules! gen_cast_mut {
+    ( $name:ident, $type:ty, $expected:expr, $variant:ident) => {
+        pub fn $name(&mut self) -> Result<&mut $type, ErrorKind> {
+            match self {
+                Value::$variant(x) => Ok(x),
+                other => Err(type_error($expected, &other)),
+            }
+        }
+    };
+}
+
 /// Generate an `is_*` type-checking method.
 macro_rules! gen_is {
     ( $name:ident, $variant:pat ) => {
@@ -284,6 +297,8 @@ impl Value {
     gen_cast!(to_list, NixList, "list", Value::List(l), l.clone());
     gen_cast!(to_closure, Closure, "lambda", Value::Closure(c), c.clone());
 
+    gen_cast_mut!(as_list_mut, NixList, "list", List);
+
     gen_is!(is_path, Value::Path(_));
     gen_is!(is_number, Value::Integer(_) | Value::Float(_));
     gen_is!(is_bool, Value::Bool(_));