feat(tvix/eval): implement upvalue resolution in `with` scopes r/4635

These need to be handled specially by the runtime if the compiler
determines that a given local must be resolved via `with`.

Note that this implementation has a bug: It currently allows `with`
inside of nested lambdas to shadow statically known identifiers. This
will be cleaned up in the next commit.

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

7 files changed, 47 insertions, 2 deletions

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index efcb83b8a4d4..f8b9ccc460ef 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -16,6 +16,7 @@
 use path_clean::PathClean;
 use rnix::ast::{self, AstToken, HasEntry};
 use rowan::ast::AstNode;
+use smol_str::SmolStr;
 use std::collections::{hash_map, HashMap};
 use std::path::{Path, PathBuf};
 use std::rc::Rc;
@@ -68,6 +69,10 @@ enum Upvalue {
 
     /// This upvalue captures an enclosing upvalue.
     Upvalue(UpvalueIdx),
+
+    /// This upvalue captures a dynamically resolved value (i.e.
+    /// `with`).
+    Dynamic(SmolStr),
 }
 
 /// Represents a scope known during compilation, which can be resolved
@@ -875,6 +880,10 @@ 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()));
+                    self.chunk().push_op(OpCode::DataDynamicIdx(idx))
+                }
             };
         }
     }
@@ -1008,11 +1017,18 @@ impl Compiler {
             return None;
         }
 
+        // Determine whether the upvalue is a local in the enclosing context.
         if let Some(idx) = self.contexts[ctx_idx - 1].scope.resolve_local(name) {
             return Some(self.add_upvalue(ctx_idx, Upvalue::Stack(idx)));
         }
 
-        // If the upvalue comes from an enclosing context, we need to
+        // Determine whether the upvalue is a dynamic variable in the
+        // enclosing context.
+        if !self.contexts[ctx_idx - 1].scope.with_stack.is_empty() {
+            return Some(self.add_upvalue(ctx_idx, Upvalue::Dynamic(SmolStr::new(name))));
+        }
+
+        // If the upvalue comes from even further up, we need to
         // recurse to make sure that the upvalues are created at each
         // level.
         if let Some(idx) = self.resolve_upvalue(ctx_idx - 1, name) {
diff --git a/tvix/eval/src/opcode.rs b/tvix/eval/src/opcode.rs
index ca99602862d5..edf7a8a5dc4a 100644
--- a/tvix/eval/src/opcode.rs
+++ b/tvix/eval/src/opcode.rs
@@ -115,4 +115,5 @@ pub enum OpCode {
     // according to the count.
     DataLocalIdx(StackIdx),
     DataUpvalueIdx(UpvalueIdx),
+    DataDynamicIdx(ConstantIdx),
 }
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-with-closure.exp b/tvix/eval/src/tests/tvix_tests/eval-okay-with-closure.exp
new file mode 100644
index 000000000000..fa8f08cb6ff8
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-with-closure.exp
@@ -0,0 +1 @@
+150
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-with-closure.nix b/tvix/eval/src/tests/tvix_tests/eval-okay-with-closure.nix
new file mode 100644
index 000000000000..7e2f7c073bfc
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-with-closure.nix
@@ -0,0 +1,5 @@
+# Upvalues from `with` require special runtime handling. Do they work?
+let
+  f = with { a = 15; }; n: n * a;
+in
+f 10
diff --git a/tvix/eval/src/value/mod.rs b/tvix/eval/src/value/mod.rs
index f6583db5ab97..eb23605bd0b2 100644
--- a/tvix/eval/src/value/mod.rs
+++ b/tvix/eval/src/value/mod.rs
@@ -80,6 +80,17 @@ impl Value {
         }
     }
 
+    pub fn as_str(&self) -> EvalResult<&str> {
+        match self {
+            Value::String(s) => Ok(s.as_str()),
+            other => Err(ErrorKind::TypeError {
+                expected: "string",
+                actual: other.type_of(),
+            }
+            .into()),
+        }
+    }
+
     pub fn to_string(self) -> EvalResult<NixString> {
         match self {
             Value::String(s) => Ok(s),
diff --git a/tvix/eval/src/value/string.rs b/tvix/eval/src/value/string.rs
index 51b0f0345457..65022b6cc49c 100644
--- a/tvix/eval/src/value/string.rs
+++ b/tvix/eval/src/value/string.rs
@@ -46,6 +46,12 @@ impl From<String> for NixString {
     }
 }
 
+impl From<SmolStr> for NixString {
+    fn from(s: SmolStr) -> Self {
+        NixString(StringRepr::Smol(s))
+    }
+}
+
 impl Hash for NixString {
     fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
         self.as_str().hash(state)
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index 4ce3505cdb9a..707d2f4b26b3 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -379,6 +379,11 @@ impl VM {
                                     .push(self.frame().closure.upvalues[upv_idx].clone());
                             }
 
+                            OpCode::DataDynamicIdx(ident_idx) => {
+                                let ident = self.chunk().constant(ident_idx).as_str()?;
+                                closure.upvalues.push(self.resolve_with(ident)?);
+                            }
+
                             _ => panic!("compiler error: missing closure operand"),
                         }
                     }
@@ -388,7 +393,7 @@ impl VM {
 
                 // Data-carrying operands should never be executed,
                 // that is a critical error in the VM.
-                OpCode::DataLocalIdx(_) | OpCode::DataUpvalueIdx(_) => {
+                OpCode::DataLocalIdx(_) | OpCode::DataUpvalueIdx(_) | OpCode::DataDynamicIdx(_) => {
                     panic!("VM bug: attempted to execute data-carrying operand")
                 }
             }