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>

1 files changed, 64 insertions, 7 deletions

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() {