feat(tvix/eval): implement attrset update (`//`) operator r/4475

The underlying implementation does a few tricks based on which pair of
attrset representations is encountered.

Particularly the effect of short-circuiting the empty cases might be
relevant in nixpkgs/NixOS, due to the use of lib.optionalAttrs.

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

2 files changed, 64 insertions, 1 deletions

diff --git a/tvix/eval/src/value/attrs.rs b/tvix/eval/src/value/attrs.rs
index 51f4795c59eb..e7da6ee621d0 100644
--- a/tvix/eval/src/value/attrs.rs
+++ b/tvix/eval/src/value/attrs.rs
@@ -35,7 +35,7 @@ impl Display for NixAttrs {
             }
 
             NixAttrs::Map(map) => {
-                for (name, value) in map {
+                for (name, value) in map.iter() {
                     f.write_fmt(format_args!("{} = {}; ", name.ident_str(), value))?;
                 }
             }
@@ -54,6 +54,59 @@ impl PartialEq for NixAttrs {
 }
 
 impl NixAttrs {
+    // Update one attribute set with the values of the other.
+    pub fn update(&self, other: &Self) -> Self {
+        match (self, other) {
+            // Short-circuit on some optimal cases:
+            (NixAttrs::Empty, NixAttrs::Empty) => NixAttrs::Empty,
+            (NixAttrs::Empty, _) => other.clone(),
+            (_, NixAttrs::Empty) => self.clone(),
+            (NixAttrs::KV { .. }, NixAttrs::KV { .. }) => other.clone(),
+
+            // Slightly more advanced, but still optimised updates
+            (NixAttrs::Map(m), NixAttrs::KV { name, value }) => {
+                let mut m = m.clone();
+                m.insert(NixString::NAME, name.clone());
+                m.insert(NixString::VALUE, value.clone());
+                NixAttrs::Map(m)
+            }
+
+            (NixAttrs::KV { name, value }, NixAttrs::Map(m)) => {
+                let mut m = m.clone();
+
+                match m.entry(NixString::NAME) {
+                    std::collections::btree_map::Entry::Vacant(e) => {
+                        e.insert(name.clone());
+                    }
+
+                    std::collections::btree_map::Entry::Occupied(_) => {
+                        /* name from `m` has precedence */
+                    }
+                };
+
+                match m.entry(NixString::VALUE) {
+                    std::collections::btree_map::Entry::Vacant(e) => {
+                        e.insert(value.clone());
+                    }
+
+                    std::collections::btree_map::Entry::Occupied(_) => {
+                        /* value from `m` has precedence */
+                    }
+                };
+
+                NixAttrs::Map(m)
+            }
+
+            // Plain merge of maps.
+            (NixAttrs::Map(m1), NixAttrs::Map(m2)) => {
+                let mut m1 = m1.clone();
+                let mut m2 = m2.clone();
+                m1.append(&mut m2);
+                NixAttrs::Map(m1)
+            }
+        }
+    }
+
     /// Retrieve reference to a mutable map inside of an attrs,
     /// optionally changing the representation if required.
     fn map_mut(&mut self) -> &mut BTreeMap<NixString, Value> {
diff --git a/tvix/eval/src/value/mod.rs b/tvix/eval/src/value/mod.rs
index 0a430ae08cb3..2a89f1c35890 100644
--- a/tvix/eval/src/value/mod.rs
+++ b/tvix/eval/src/value/mod.rs
@@ -71,6 +71,16 @@ impl Value {
             }),
         }
     }
+
+    pub fn as_attrs(self) -> EvalResult<Rc<NixAttrs>> {
+        match self {
+            Value::Attrs(s) => Ok(s),
+            other => Err(Error::TypeError {
+                expected: "set",
+                actual: other.type_of(),
+            }),
+        }
+    }
 }
 
 impl Display for Value {