feat(tvix/eval): make NixList::clone() cheap r/5301

When we start unrecursivifying (sp?) things, Rust's borrow checker
is going to be a headache; its magic only works when you use the CPU
stack as your call stack.

Fixing the borrow checker issues usually involves adding lots of
`clone()`s.  Right now `NixList` is the only variant of `Value` that
isn't cheap to clone() -- all the others are either a wrapper around
Rc or else are of bounded size.

Note that this requires dropping the `DerefMut for NixList` instance
and using `Vec<Value>` instead in those situations.

Change-Id: I5a47df66855342aa2064f8f3cb7934ff422d26bd
Signed-off-by: Adam Joseph <adam@westernsemico.com>
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7359
Reviewed-by: tazjin <tazjin@tvl.su>
Tested-by: BuildkiteCI

3 files changed, 36 insertions, 38 deletions

diff --git a/tvix/eval/src/builtins/mod.rs b/tvix/eval/src/builtins/mod.rs
index 89e7c37d3f20..7c4f40479bc2 100644
--- a/tvix/eval/src/builtins/mod.rs
+++ b/tvix/eval/src/builtins/mod.rs
@@ -360,7 +360,7 @@ mod pure_builtins {
 
         let operator = attrs.select_required("operator")?;
 
-        let mut res = NixList::new();
+        let mut res: Vec<Value> = vec![];
         let mut done_keys: Vec<Value> = vec![];
 
         let mut insert_key = |k: Value, vm: &mut VM| -> Result<bool, ErrorKind> {
@@ -387,7 +387,7 @@ mod pure_builtins {
             work_set.extend(op_result.into_iter());
         }
 
-        Ok(Value::List(res))
+        Ok(Value::List(NixList::from(res)))
     }
 
     #[builtin("genList")]
@@ -415,15 +415,16 @@ mod pure_builtins {
 
     #[builtin("groupBy")]
     fn builtin_group_by(vm: &mut VM, f: Value, list: Value) -> Result<Value, ErrorKind> {
-        let mut res: BTreeMap<NixString, Value> = BTreeMap::new();
+        let mut res: BTreeMap<NixString, Vec<Value>> = BTreeMap::new();
         for val in list.to_list()? {
             let key = vm.call_with(&f, [val.clone()])?.force(vm)?.to_str()?;
-            res.entry(key)
-                .or_insert_with(|| Value::List(NixList::new()))
-                .as_list_mut()?
-                .push(val);
+            res.entry(key).or_insert_with(|| vec![]).push(val);
         }
-        Ok(Value::attrs(NixAttrs::from_map(res)))
+        Ok(Value::attrs(NixAttrs::from_map(
+            res.into_iter()
+                .map(|(k, v)| (k, Value::List(NixList::from(v))))
+                .collect(),
+        )))
     }
 
     #[builtin("hasAttr")]
@@ -745,7 +746,7 @@ mod pure_builtins {
         let re: Regex = Regex::new(re.as_str()).unwrap();
         let mut capture_locations = re.capture_locations();
         let num_captures = capture_locations.len();
-        let mut ret = NixList::new();
+        let mut ret: Vec<Value> = vec![];
         let mut pos = 0;
 
         while let Some(thematch) = re.captures_read_at(&mut capture_locations, text, pos) {
@@ -770,12 +771,12 @@ mod pure_builtins {
         // push the unmatched characters following the last match
         ret.push(Value::from(&text[pos..]));
 
-        Ok(Value::List(ret))
+        Ok(Value::List(NixList::from(ret)))
     }
 
     #[builtin("sort")]
     fn builtin_sort(vm: &mut VM, comparator: Value, list: Value) -> Result<Value, ErrorKind> {
-        let mut list = list.to_list()?;
+        let mut list = list.to_list()?.into_vec();
 
         // Used to let errors "escape" from the sorting closure. If anything
         // ends up setting an error, it is returned from this function.
@@ -806,7 +807,7 @@ mod pure_builtins {
         });
 
         match error {
-            None => Ok(Value::List(list)),
+            None => Ok(Value::List(NixList::from(list))),
             Some(e) => Err(e),
         }
     }
diff --git a/tvix/eval/src/value/list.rs b/tvix/eval/src/value/list.rs
index 085c85346b65..f90d625ce353 100644
--- a/tvix/eval/src/value/list.rs
+++ b/tvix/eval/src/value/list.rs
@@ -1,5 +1,6 @@
 //! This module implements Nix lists.
-use std::ops::{Deref, DerefMut};
+use std::ops::Deref;
+use std::rc::Rc;
 
 use crate::errors::ErrorKind;
 use crate::vm::VM;
@@ -10,13 +11,13 @@ use super::Value;
 
 #[repr(transparent)]
 #[derive(Clone, Debug)]
-pub struct NixList(Vec<Value>);
+pub struct NixList(Rc<Vec<Value>>);
 
 impl TotalDisplay for NixList {
     fn total_fmt(&self, f: &mut std::fmt::Formatter<'_>, set: &mut ThunkSet) -> std::fmt::Result {
         f.write_str("[ ")?;
 
-        for v in &self.0 {
+        for v in self.0.as_ref() {
             v.total_fmt(f, set)?;
             f.write_str(" ")?;
         }
@@ -27,7 +28,7 @@ impl TotalDisplay for NixList {
 
 impl From<Vec<Value>> for NixList {
     fn from(vs: Vec<Value>) -> Self {
-        Self(vs)
+        Self(Rc::new(vs))
     }
 }
 
@@ -45,24 +46,14 @@ mod arbitrary {
         type Strategy = BoxedStrategy<Self>;
 
         fn arbitrary_with(args: Self::Parameters) -> Self::Strategy {
-            any_with::<Vec<Value>>(args).prop_map(Self).boxed()
+            any_with::<Rc<Vec<Value>>>(args).prop_map(Self).boxed()
         }
     }
 }
 
 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);
-        ret
+        Self(Rc::new(vec![]))
     }
 
     pub fn len(&self) -> usize {
@@ -81,15 +72,22 @@ impl NixList {
             stack_slice.len(),
         );
 
-        NixList(stack_slice)
+        NixList(Rc::new(stack_slice))
     }
 
     pub fn iter(&self) -> std::slice::Iter<Value> {
         self.0.iter()
     }
 
+    pub fn ptr_eq(&self, other: &Self) -> bool {
+        Rc::ptr_eq(&self.0, &other.0)
+    }
+
     /// Compare `self` against `other` for equality using Nix equality semantics
     pub fn nix_eq(&self, other: &Self, vm: &mut VM) -> Result<bool, ErrorKind> {
+        if self.ptr_eq(other) {
+            return Ok(true);
+        }
         if self.len() != other.len() {
             return Ok(false);
         }
@@ -107,6 +105,10 @@ impl NixList {
     pub fn force_elements(&self, vm: &mut VM) -> Result<(), ErrorKind> {
         self.iter().try_for_each(|v| v.force(vm).map(|_| ()))
     }
+
+    pub fn into_vec(self) -> Vec<Value> {
+        crate::unwrap_or_clone_rc(self.0)
+    }
 }
 
 impl IntoIterator for NixList {
@@ -114,7 +116,7 @@ impl IntoIterator for NixList {
     type IntoIter = std::vec::IntoIter<Value>;
 
     fn into_iter(self) -> std::vec::IntoIter<Value> {
-        self.0.into_iter()
+        self.into_vec().into_iter()
     }
 }
 
@@ -135,9 +137,3 @@ impl Deref for NixList {
         &self.0
     }
 }
-
-impl DerefMut for NixList {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.0
-    }
-}
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index 17a7e0ce2758..400d85adac30 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -538,8 +538,9 @@ impl<'o> VM<'o> {
 
             OpCode::OpConcat => {
                 let rhs = fallible!(self, self.pop().to_list());
-                let lhs = fallible!(self, self.pop().to_list());
-                self.push(Value::List(lhs.concat(&rhs)))
+                let mut lhs = fallible!(self, self.pop().to_list()).into_vec();
+                lhs.extend_from_slice(&rhs);
+                self.push(Value::List(NixList::from(lhs)))
             }
 
             OpCode::OpInterpolate(Count(count)) => self.run_interpolate(count)?,