fix(tvix/eval): correctly thread through dynamic upvalues r/4658

This puts together the puzzle pieces for threading dynamic
upvalues (that is, upvalues resolved from the `with`-stack) all the
way through.

Reading the test case enclosed in this commit and walking through it
is recommended to understand what problem is being tackled here.

In short, because the compiler can not statically know *which*
with-scope a dynamic argument is resolved from it needs to lay the
groundwork for resolving from *all* possible scopes.

There are multiple different approaches to doing this. The approach
chosen in this commit is that if a dynamic upvalue is detected, the
compiler will emit instructions to close over this dynamic value
in *all* enclosing lambda contexts.

It uses a new instruction for this that will leave around a sentinel
value in case an identifier could not be resolved, and wire the
location of this found value (or sentinel) up through the upvalues to
the next level of nesting.

In this tradeoff, tvix potentially closes over more upvalues than are
needed (but in practice, how often do people create *really* deep
`with`-stacks? and in *this* kind of code situation? maybe we should
even warn for this!) but avoids keeping the entire attribute sets
themselves around.

Looking at the test case, each surrounding closure will close
over *all* dynamic identifiers that are referenced later on visible to
it, but only the last one for each identifier will actually end up
being used.

This also covers our bases for an additional edge-case this creates,
in which an identifier potentially resolves to a dynamic upvalue *and*
to a dynamic value within the function's own scope (again, would
anyone really do this?) by introducing a resolution instruction for
that particular case.

There is likely some potential for cleaning up this code which is
quite ugly in some parts, but as this implementation is now carefully
calibrated to work I decided it is time to commit it and clean it up
in subsequent commits.

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

6 files changed, 112 insertions, 12 deletions

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index 11d6d8231f09..78fe76ca01a9 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -702,6 +702,15 @@ impl Compiler {
                 if let Some(idx) =
                     self.resolve_dynamic_upvalue(self.contexts.len() - 1, ident.text())
                 {
+                    // Edge case: Current scope *also* has a non-empty
+                    // `with`-stack. This means we need to resolve
+                    // both in this scope, and in the upvalues.
+                    if self.scope().has_with() {
+                        self.emit_constant(Value::String(ident.text().into()));
+                        self.chunk().push_op(OpCode::OpResolveWithOrUpvalue(idx));
+                        return;
+                    }
+
                     self.chunk().push_op(OpCode::OpGetUpvalue(idx));
                     return;
                 }
@@ -810,11 +819,18 @@ impl Compiler {
         self.chunk().push_op(OpCode::OpClosure(closure_idx));
         for upvalue in compiled.scope.upvalues {
             match upvalue {
-                Upvalue::Stack(idx) => self.chunk().push_op(OpCode::DataLocalIdx(idx)),
-                Upvalue::Upvalue(idx) => self.chunk().push_op(OpCode::DataUpvalueIdx(idx)),
-                Upvalue::Dynamic { name, .. } => {
+                Upvalue::Stack(idx) => {
+                    self.chunk().push_op(OpCode::DataLocalIdx(idx));
+                }
+                Upvalue::Upvalue(idx) => {
+                    self.chunk().push_op(OpCode::DataUpvalueIdx(idx));
+                }
+                Upvalue::Dynamic { name, up } => {
                     let idx = self.chunk().push_constant(Value::String(name.into()));
-                    self.chunk().push_op(OpCode::DataDynamicIdx(idx))
+                    self.chunk().push_op(OpCode::DataDynamicIdx(idx));
+                    if let Some(up) = up {
+                        self.chunk().push_op(OpCode::DataDynamicAncestor(up));
+                    }
                 }
             };
         }
diff --git a/tvix/eval/src/opcode.rs b/tvix/eval/src/opcode.rs
index edf7a8a5dc4a..5f1f33f2b3e5 100644
--- a/tvix/eval/src/opcode.rs
+++ b/tvix/eval/src/opcode.rs
@@ -82,6 +82,7 @@ pub enum OpCode {
     OpPushWith(StackIdx),
     OpPopWith,
     OpResolveWith,
+    OpResolveWithOrUpvalue(UpvalueIdx),
 
     // Lists
     OpList(Count),
@@ -116,4 +117,5 @@ pub enum OpCode {
     DataLocalIdx(StackIdx),
     DataUpvalueIdx(UpvalueIdx),
     DataDynamicIdx(ConstantIdx),
+    DataDynamicAncestor(UpvalueIdx),
 }
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-deeply-nested-with-closure.exp b/tvix/eval/src/tests/tvix_tests/eval-okay-deeply-nested-with-closure.exp
new file mode 100644
index 000000000000..3bed31f76e3f
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-deeply-nested-with-closure.exp
@@ -0,0 +1 @@
+[ 1 2 3 4 ]
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-deeply-nested-with-closure.nix b/tvix/eval/src/tests/tvix_tests/eval-okay-deeply-nested-with-closure.nix
new file mode 100644
index 000000000000..7f13f1f27030
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-deeply-nested-with-closure.nix
@@ -0,0 +1,18 @@
+# This convoluted test constructs a situation in which dynamically
+# resolved upvalues refer `with` blocks introduced at different lambda
+# context boundaries, i.e. the access to a, b in the innermost closure
+# must be threaded through upvalues in several levels.
+
+(_:
+with { a = 1; b = 1; };
+
+_:
+with { b = 2; c = 2; };
+
+_:
+with { c = 3; d = 3; };
+
+_:
+with { d = 4; };
+
+[ a b c d ]) null null null null
diff --git a/tvix/eval/src/value/mod.rs b/tvix/eval/src/value/mod.rs
index eb23605bd0b2..54211e8ba313 100644
--- a/tvix/eval/src/value/mod.rs
+++ b/tvix/eval/src/value/mod.rs
@@ -34,6 +34,7 @@ pub enum Value {
     // are never returned to or created directly by users.
     AttrPath(Vec<NixString>),
     AttrNotFound,
+    DynamicUpvalueMissing(NixString),
 }
 
 impl Value {
@@ -54,7 +55,9 @@ impl Value {
             Value::Closure(_) | Value::Builtin(_) => "lambda",
 
             // Internal types
-            Value::AttrPath(_) | Value::AttrNotFound => "internal",
+            Value::AttrPath(_) | Value::AttrNotFound | Value::DynamicUpvalueMissing(_) => {
+                "internal"
+            }
         }
     }
 
@@ -163,6 +166,9 @@ impl Display for Value {
             // internal types
             Value::AttrPath(path) => write!(f, "internal[attrpath({})]", path.len()),
             Value::AttrNotFound => f.write_str("internal[not found]"),
+            Value::DynamicUpvalueMissing(name) => {
+                write!(f, "internal[no_dyn_upvalue({name})]")
+            }
         }
     }
 }
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index 23d9c3555153..5d5cb57dbabf 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -5,8 +5,8 @@ use std::{cell::Ref, rc::Rc};
 
 use crate::{
     chunk::Chunk,
-    errors::{ErrorKind, EvalResult},
-    opcode::{Count, JumpOffset, OpCode, StackIdx},
+    errors::{Error, ErrorKind, EvalResult},
+    opcode::{ConstantIdx, Count, JumpOffset, OpCode, StackIdx},
     value::{Closure, Lambda, NixAttrs, NixList, Value},
 };
 
@@ -100,6 +100,10 @@ impl VM {
         op
     }
 
+    fn peek_op(&self) -> OpCode {
+        self.chunk().code[self.frame().ip]
+    }
+
     fn pop(&mut self) -> Value {
         self.stack.pop().expect("runtime stack empty")
     }
@@ -337,6 +341,25 @@ impl VM {
                     self.push(value)
                 }
 
+                OpCode::OpResolveWithOrUpvalue(idx) => {
+                    let ident = self.pop().to_string()?;
+                    match self.resolve_with(ident.as_str()) {
+                        // Variable found in local `with`-stack.
+                        Ok(value) => self.push(value),
+
+                        // Variable not found => check upvalues.
+                        Err(Error {
+                            kind: ErrorKind::UnknownDynamicVariable(_),
+                            ..
+                        }) => {
+                            let value = self.frame().closure.upvalue(idx).clone();
+                            self.push(value);
+                        }
+
+                        Err(err) => return Err(err),
+                    }
+                }
+
                 OpCode::OpAssert => {
                     if !self.pop().as_bool()? {
                         return Err(ErrorKind::AssertionFailed.into());
@@ -389,10 +412,8 @@ impl VM {
                             }
 
                             OpCode::DataDynamicIdx(ident_idx) => {
-                                let chunk = self.chunk();
-                                let ident = chunk.constant(ident_idx).as_str()?.to_string();
-                                drop(chunk); // some lifetime trickery due to cell::Ref
-                                closure.push_upvalue(self.resolve_with(&ident)?);
+                                let value = self.resolve_dynamic_upvalue(ident_idx)?;
+                                closure.push_upvalue(value);
                             }
 
                             _ => panic!("compiler error: missing closure operand"),
@@ -402,7 +423,10 @@ impl VM {
 
                 // Data-carrying operands should never be executed,
                 // that is a critical error in the VM.
-                OpCode::DataLocalIdx(_) | OpCode::DataUpvalueIdx(_) | OpCode::DataDynamicIdx(_) => {
+                OpCode::DataLocalIdx(_)
+                | OpCode::DataUpvalueIdx(_)
+                | OpCode::DataDynamicIdx(_)
+                | OpCode::DataDynamicAncestor(_) => {
                     panic!("VM bug: attempted to execute data-carrying operand")
                 }
             }
@@ -452,6 +476,39 @@ impl VM {
         Ok(())
     }
 
+    fn resolve_dynamic_upvalue(&mut self, ident_idx: ConstantIdx) -> EvalResult<Value> {
+        let chunk = self.chunk();
+        let ident = chunk.constant(ident_idx).as_str()?.to_string();
+        drop(chunk); // some lifetime trickery due to cell::Ref
+
+        // Peek at the current instruction (note: IP has already
+        // advanced!) to see if it is actually data indicating a
+        // "fallback upvalue" in case the dynamic could not be
+        // resolved at this level.
+        let up = match self.peek_op() {
+            OpCode::DataDynamicAncestor(idx) => {
+                // advance ip past this data
+                self.inc_ip();
+                Some(idx)
+            }
+            _ => None,
+        };
+
+        match self.resolve_with(&ident) {
+            Ok(v) => Ok(v),
+
+            Err(Error {
+                kind: ErrorKind::UnknownDynamicVariable(_),
+                ..
+            }) => match up {
+                Some(idx) => Ok(self.frame().closure.upvalue(idx).clone()),
+                None => Ok(Value::DynamicUpvalueMissing(ident.into())),
+            },
+
+            Err(err) => Err(err),
+        }
+    }
+
     /// Resolve a dynamic identifier through the with-stack at runtime.
     fn resolve_with(&self, ident: &str) -> EvalResult<Value> {
         for idx in self.with_stack.iter().rev() {