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>

1 files changed, 17 insertions, 1 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) {