refactor(tvix/eval): thread dynamic upvalues through all contexts r/4657

With this change, encountering a dynamic upvalue will thread through
all contexts starting from the lowest context that has a non-empty
`with`-stack.

The additional upvalues are not actually used yet, so the effective
behaviour remains mostly the same. This is done in preparation for an
upcoming change, which will implement proper dynamic resolution for
complex cases of nested dynamic upvalues.

Yes, this whole upvalue + dynamic values thing is a little bit
mind-bending, but we would like to not give up being able to resolve a
large chunk of the scoping behaviour statically.

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

2 files changed, 65 insertions, 9 deletions

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index ffc7daf8e9..11d6d8231f 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -698,6 +698,14 @@ impl Compiler {
                     return;
                 }
 
+                // Even worse - are we dealing with a dynamic upvalue?
+                if let Some(idx) =
+                    self.resolve_dynamic_upvalue(self.contexts.len() - 1, ident.text())
+                {
+                    self.chunk().push_op(OpCode::OpGetUpvalue(idx));
+                    return;
+                }
+
                 if !self.scope().has_with() {
                     self.emit_error(node.syntax().clone(), ErrorKind::UnknownStaticVariable);
                     return;
@@ -804,8 +812,8 @@ impl Compiler {
             match upvalue {
                 Upvalue::Stack(idx) => self.chunk().push_op(OpCode::DataLocalIdx(idx)),
                 Upvalue::Upvalue(idx) => self.chunk().push_op(OpCode::DataUpvalueIdx(idx)),
-                Upvalue::Dynamic(s) => {
-                    let idx = self.chunk().push_constant(Value::String(s.into()));
+                Upvalue::Dynamic { name, .. } => {
+                    let idx = self.chunk().push_constant(Value::String(name.into()));
                     self.chunk().push_op(OpCode::DataDynamicIdx(idx))
                 }
             };
@@ -1000,11 +1008,52 @@ impl Compiler {
             return Some(self.add_upvalue(ctx_idx, Upvalue::Upvalue(idx)));
         }
 
-        // If the resolution of a statically known upvalue failed,
-        // attempt to resolve a dynamic one (i.e. search for enclosing
-        // `with` blocks and make that resolution dynamic).
-        if self.contexts[ctx_idx - 1].scope.has_with() {
-            return Some(self.add_upvalue(ctx_idx, Upvalue::Dynamic(SmolStr::new(name))));
+        None
+    }
+
+    /// If no static resolution for a potential upvalue was found,
+    /// finds the lowest lambda context that has a `with`-stack and
+    /// thread dynamic upvalues all the way through.
+    ///
+    /// At runtime, as closures are being constructed they either
+    /// capture a dynamically available upvalue, take an upvalue from
+    /// their "ancestor" or leave a sentinel value on the stack.
+    ///
+    /// As such an upvalue is actually accessed, an error is produced
+    /// when the sentinel is found. See the runtime's handling of
+    /// dynamic upvalues for details.
+    fn resolve_dynamic_upvalue(&mut self, at: usize, name: &str) -> Option<UpvalueIdx> {
+        if at == 0 {
+            // There can not be any upvalue at the outermost context.
+            return None;
+        }
+
+        if let Some((lowest_idx, _)) = self
+            .contexts
+            .iter()
+            .enumerate()
+            .find(|(_, c)| c.scope.has_with())
+        {
+            // An enclosing lambda context has dynamic values. Each
+            // context in the chain from that point on now needs to
+            // capture dynamic upvalues because we can not statically
+            // know at which level the correct one is located.
+            let name = SmolStr::new(name);
+            let mut upvalue_idx = None;
+
+            for idx in lowest_idx..=at {
+                upvalue_idx = Some(self.add_upvalue(
+                    idx,
+                    Upvalue::Dynamic {
+                        name: name.clone(),
+                        up: upvalue_idx,
+                    },
+                ));
+            }
+
+            // Return the outermost upvalue index (i.e. the one of the
+            // current context).
+            return upvalue_idx;
         }
 
         None
diff --git a/tvix/eval/src/compiler/scope.rs b/tvix/eval/src/compiler/scope.rs
index e295731e32..3061e08ead 100644
--- a/tvix/eval/src/compiler/scope.rs
+++ b/tvix/eval/src/compiler/scope.rs
@@ -81,7 +81,7 @@ pub enum LocalPosition {
 
 /// Represents the different ways in which upvalues can be captured in
 /// closures or thunks.
-#[derive(Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq)]
 pub enum Upvalue {
     /// This upvalue captures a local from the stack.
     Stack(StackIdx),
@@ -91,7 +91,14 @@ pub enum Upvalue {
 
     /// This upvalue captures a dynamically resolved value (i.e.
     /// `with`).
-    Dynamic(SmolStr),
+    ///
+    /// It stores the identifier with which to perform a dynamic
+    /// lookup, as well as the optional upvalue index in the enclosing
+    /// function (if any).
+    Dynamic {
+        name: SmolStr,
+        up: Option<UpvalueIdx>,
+    },
 }
 
 /// Represents a scope known during compilation, which can be resolved