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

13 files changed, 87 insertions, 1 deletions

diff --git a/tvix/eval/src/compiler.rs b/tvix/eval/src/compiler.rs
index 668ec842e81b..4aadfaba80e7 100644
--- a/tvix/eval/src/compiler.rs
+++ b/tvix/eval/src/compiler.rs
@@ -136,6 +136,7 @@ impl Compiler {
             BinOpKind::Mul => OpCode::OpMul,
             BinOpKind::Div => OpCode::OpDiv,
             BinOpKind::Equal => OpCode::OpEqual,
+            BinOpKind::Update => OpCode::OpAttrsUpdate,
             _ => todo!(),
         };
 
@@ -187,6 +188,10 @@ impl Compiler {
     // 2. Keys can refer to nested attribute sets.
     // 3. Attribute sets can (optionally) be recursive.
     fn compile_attr_set(&mut self, node: rnix::types::AttrSet) -> EvalResult<()> {
+        if node.recursive() {
+            todo!("recursive attribute sets are not yet implemented")
+        }
+
         let mut count = 0;
 
         for kv in node.entries() {
diff --git a/tvix/eval/src/opcode.rs b/tvix/eval/src/opcode.rs
index 622a02ac85f8..f682cfc8b83f 100644
--- a/tvix/eval/src/opcode.rs
+++ b/tvix/eval/src/opcode.rs
@@ -33,6 +33,7 @@ pub enum OpCode {
     // Attribute sets
     OpAttrs(usize),
     OpAttrPath(usize),
+    OpAttrsUpdate,
 
     // Lists
     OpList(usize),
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update-empty-lhs.exp b/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update-empty-lhs.exp
new file mode 100644
index 000000000000..fedf8f25a693
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update-empty-lhs.exp
@@ -0,0 +1 @@
+{ a = "ok"; }
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update-empty-lhs.nix b/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update-empty-lhs.nix
new file mode 100644
index 000000000000..9596be22b831
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update-empty-lhs.nix
@@ -0,0 +1 @@
+{} // { a = "ok"; }
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update-empty-rhs.exp b/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update-empty-rhs.exp
new file mode 100644
index 000000000000..fedf8f25a693
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update-empty-rhs.exp
@@ -0,0 +1 @@
+{ a = "ok"; }
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update-empty-rhs.nix b/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update-empty-rhs.nix
new file mode 100644
index 000000000000..117c01141357
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update-empty-rhs.nix
@@ -0,0 +1 @@
+{ a = "ok"; } // {}
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update-kv-lhs.exp b/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update-kv-lhs.exp
new file mode 100644
index 000000000000..c2234a47e2b8
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update-kv-lhs.exp
@@ -0,0 +1 @@
+{ name = "foo"; other = 42; value = "bar"; }
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update-kv-lhs.nix b/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update-kv-lhs.nix
new file mode 100644
index 000000000000..6f71684902e5
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update-kv-lhs.nix
@@ -0,0 +1 @@
+{ name = "foo"; value = "bar"; } // { other = 42; }
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update.exp b/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update.exp
new file mode 100644
index 000000000000..57f4d541bd85
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update.exp
@@ -0,0 +1 @@
+{ a = 15; b = "works"; }
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update.nix b/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update.nix
new file mode 100644
index 000000000000..735602fe02d5
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-attrs-update.nix
@@ -0,0 +1 @@
+{ a = 15; } // { b = "works"; }
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 {
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index 9a65668caa23..7a1082344e34 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -117,8 +117,17 @@ impl VM {
                 OpCode::OpNull => self.push(Value::Null),
                 OpCode::OpTrue => self.push(Value::Bool(true)),
                 OpCode::OpFalse => self.push(Value::Bool(false)),
+
                 OpCode::OpAttrs(count) => self.run_attrset(count)?,
                 OpCode::OpAttrPath(count) => self.run_attr_path(count)?,
+
+                OpCode::OpAttrsUpdate => {
+                    let rhs = self.pop().as_attrs()?;
+                    let lhs = self.pop().as_attrs()?;
+
+                    self.push(Value::Attrs(Rc::new(lhs.update(&rhs))))
+                }
+
                 OpCode::OpList(count) => self.run_list(count)?,
                 OpCode::OpInterpolate(count) => self.run_interpolate(count)?,
             }