fix(tvix/eval): account for attrset temporaries during construction r/4796

The temporaries left on the stack as operands to `OpAttrs` must be
accounted for in the locals array in order for operations within them
to receive correct slots.

Some test cases that were previously broken have been added.

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

5 files changed, 68 insertions, 2 deletions

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index e119c866bc49..b6a3f329a0d6 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -473,6 +473,8 @@ impl Compiler<'_, '_> {
                         // First emit the identifier itself (this
                         // becomes the new key).
                         self.emit_literal_ident(&ident);
+                        let ident_span = self.span_for(&ident);
+                        self.scope_mut().declare_phantom(ident_span, true);
 
                         // Then emit the node that we're inheriting
                         // from.
@@ -482,22 +484,27 @@ impl Compiler<'_, '_> {
                         // instruction followed by a merge, rather
                         // than pushing/popping the same attrs
                         // potentially a lot of times.
-                        self.compile(slot, from.expr().unwrap());
+                        let val_idx = self.scope_mut().declare_phantom(ident_span, false);
+                        self.compile(val_idx, from.expr().unwrap());
                         self.emit_force(&from.expr().unwrap());
                         self.emit_literal_ident(&ident);
                         self.push_op(OpCode::OpAttrsSelect, &ident);
+                        self.scope_mut().mark_initialised(val_idx);
                     }
                 }
 
                 None => {
                     for ident in inherit.idents() {
+                        let ident_span = self.span_for(&ident);
                         count += 1;
 
                         // Emit the key to use for OpAttrs
                         self.emit_literal_ident(&ident);
+                        self.scope_mut().declare_phantom(ident_span, true);
 
                         // Emit the value.
                         self.compile_ident(slot, ident);
+                        self.scope_mut().declare_phantom(ident_span, true);
                     }
                 }
             }
@@ -519,6 +526,10 @@ impl Compiler<'_, '_> {
             todo!("recursive attribute sets are not yet implemented")
         }
 
+        // Open a scope to track the positions of the temporaries used
+        // by the `OpAttrs` instruction.
+        self.begin_scope();
+
         let mut count = self.compile_inherit_attrs(slot, node.inherits());
 
         for kv in node.attrpath_values() {
@@ -530,9 +541,32 @@ impl Compiler<'_, '_> {
             // runtime value representing the attribute path is
             // emitted.
             let mut key_count = 0;
+            let key_span = self.span_for(&kv.attrpath().unwrap());
+            let key_idx = self.scope_mut().declare_phantom(key_span, false);
+
             for fragment in kv.attrpath().unwrap().attrs() {
+                // Key fragments can contain dynamic expressions,
+                // which makes accounting for their stack slots very
+                // tricky.
+                //
+                // In order to ensure the locals are correctly cleaned
+                // up, we essentially treat any key fragment after the
+                // first one (which has a locals index that will
+                // become that of the final key itself) as being part
+                // of a separate scope, which results in the somewhat
+                // annoying setup logic below.
+                let fragment_slot = match key_count {
+                    0 => key_idx,
+                    1 => {
+                        self.begin_scope();
+                        self.scope_mut().declare_phantom(key_span, false)
+                    }
+                    _ => self.scope_mut().declare_phantom(key_span, false),
+                };
+
                 key_count += 1;
-                self.compile_attr(slot, fragment);
+                self.compile_attr(fragment_slot, fragment);
+                self.scope_mut().mark_initialised(fragment_slot);
             }
 
             // We're done with the key if there was only one fragment,
@@ -543,15 +577,26 @@ impl Compiler<'_, '_> {
                     OpCode::OpAttrPath(Count(key_count)),
                     &kv.attrpath().unwrap(),
                 );
+
+                // Close the temporary scope that was set up for the
+                // key fragments.
+                self.scope_mut().end_scope();
             }
 
             // The value is just compiled as normal so that its
             // resulting value is on the stack when the attribute set
             // is constructed at runtime.
+            let value_span = self.span_for(&kv.value().unwrap());
+            let value_slot = self.scope_mut().declare_phantom(value_span, false);
             self.compile(slot, kv.value().unwrap());
+            self.scope_mut().mark_initialised(value_slot);
         }
 
         self.push_op(OpCode::OpAttrs(Count(count)), &node);
+
+        // Remove the temporary scope, but do not emit any additional
+        // cleanup (OpAttrs consumes all of these locals).
+        self.scope_mut().end_scope();
     }
 
     fn compile_select(&mut self, slot: LocalIdx, node: ast::Select) {
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-thunked-with.exp b/tvix/eval/src/tests/tvix_tests/eval-okay-thunked-with.exp
new file mode 100644
index 000000000000..d81cc0710eb6
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-thunked-with.exp
@@ -0,0 +1 @@
+42
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-thunked-with.nix b/tvix/eval/src/tests/tvix_tests/eval-okay-thunked-with.nix
new file mode 100644
index 000000000000..6f32660c4c33
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-thunked-with.nix
@@ -0,0 +1,7 @@
+# Creates a `with` across multiple thunk boundaries.
+
+let
+  set = {
+    a = with { b = 42; }; b;
+  };
+in set.a
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-with-in-dynamic-key.exp b/tvix/eval/src/tests/tvix_tests/eval-okay-with-in-dynamic-key.exp
new file mode 100644
index 000000000000..d81cc0710eb6
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-with-in-dynamic-key.exp
@@ -0,0 +1 @@
+42
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-with-in-dynamic-key.nix b/tvix/eval/src/tests/tvix_tests/eval-okay-with-in-dynamic-key.nix
new file mode 100644
index 000000000000..c44455a5bf83
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-with-in-dynamic-key.nix
@@ -0,0 +1,12 @@
+# Tests correct tracking of stack indices within construction of an
+# attribute set. Dynamic keys can be any expression, so something that
+# is extremely sensitive to stack offsets (like `with`) can be tricky.
+
+let
+  set1 = { key = "b"; };
+  set2 = {
+    a = 20;
+    ${with set1; key} = 20;
+    ${with { key = "c"; }; key} = 2;
+  };
+in set2.a + set2.b + set2.c