fix(tvix/eval): inherit scope poisoning data in nested contexts r/4743

Scope poisoning must be inherited across lambda context boundaries,
e.g. if an outer scope has a poisoned `null`, any lambdas defined on
the same level must reference that poisoned identifier correctly.

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

4 files changed, 36 insertions, 4 deletions

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index cc98c2e852f2..2d499564077d 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -53,6 +53,13 @@ impl LambdaCtx {
             scope: Default::default(),
         }
     }
+
+    fn inherit(&self) -> Self {
+        LambdaCtx {
+            lambda: Lambda::new_anonymous(),
+            scope: self.scope.inherit(),
+        }
+    }
 }
 
 /// Alias for the map of globally available functions that should
@@ -836,9 +843,7 @@ impl Compiler<'_> {
     }
 
     fn compile_lambda(&mut self, slot: Option<LocalIdx>, node: ast::Lambda) {
-        // Open new lambda context in compiler, which has its own
-        // scope etc.
-        self.contexts.push(LambdaCtx::new());
+        self.new_context();
         self.begin_scope();
 
         // Compile the function itself
@@ -912,7 +917,7 @@ impl Compiler<'_> {
         N: AstNode + Clone,
         F: FnOnce(&mut Compiler, &N, Option<LocalIdx>),
     {
-        self.contexts.push(LambdaCtx::new());
+        self.new_context();
         self.begin_scope();
         content(self, node, slot);
         self.end_scope(node);
@@ -1011,10 +1016,14 @@ impl Compiler<'_> {
         }
     }
 
+    /// Increase the scope depth of the current function (e.g. within
+    /// a new bindings block, or `with`-scope).
     fn begin_scope(&mut self) {
         self.scope_mut().scope_depth += 1;
     }
 
+    /// Decrease scope depth of the current function and emit
+    /// instructions to clean up the stack at runtime.
     fn end_scope<N: AstNode>(&mut self, node: &N) {
         debug_assert!(self.scope().scope_depth != 0, "can not end top scope");
 
@@ -1055,6 +1064,15 @@ impl Compiler<'_> {
         }
     }
 
+    /// Open a new lambda context within which to compile a function,
+    /// closure or thunk.
+    fn new_context(&mut self) {
+        // This must inherit the scope-poisoning status of the parent
+        // in order for upvalue resolution to work correctly with
+        // poisoned identifiers.
+        self.contexts.push(self.context().inherit());
+    }
+
     /// Declare a local variable known in the scope that is being
     /// compiled by pushing it to the locals. This is used to
     /// determine the stack offset of variables.
diff --git a/tvix/eval/src/compiler/scope.rs b/tvix/eval/src/compiler/scope.rs
index b1ad4bfcd2b7..a50b3b79343e 100644
--- a/tvix/eval/src/compiler/scope.rs
+++ b/tvix/eval/src/compiler/scope.rs
@@ -152,6 +152,15 @@ impl Scope {
         }
     }
 
+    /// Inherit scope details from a parent scope (required for
+    /// correctly nesting scopes in lambdas and thunks when special
+    /// scope features like poisoning are present).
+    pub fn inherit(&self) -> Self {
+        let mut scope = Self::default();
+        scope.poisoned_tokens = self.poisoned_tokens.clone();
+        scope
+    }
+
     /// Check whether a given token is poisoned.
     pub fn is_poisoned(&self, name: &str) -> bool {
         self.poisoned_tokens.contains_key(name)
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-nested-poisoning.exp b/tvix/eval/src/tests/tvix_tests/eval-okay-nested-poisoning.exp
new file mode 100644
index 000000000000..d81cc0710eb6
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-nested-poisoning.exp
@@ -0,0 +1 @@
+42
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-nested-poisoning.nix b/tvix/eval/src/tests/tvix_tests/eval-okay-nested-poisoning.nix
new file mode 100644
index 000000000000..8d0280bb8973
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-nested-poisoning.nix
@@ -0,0 +1,4 @@
+let
+  null = 1;
+  f = n: n + null;
+in f 41