2 files changed, 55 insertions, 37 deletions

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index d05989562f66..89b90f967268 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -781,7 +781,7 @@ impl Compiler<'_, '_> {
 
         // Deal with the body, then clean up the locals afterwards.
         self.compile(slot, node.body().unwrap());
-        self.end_scope(&node);
+        self.cleanup_scope(&node);
     }
 
     fn compile_ident(&mut self, slot: LocalIdx, node: ast::Ident) {
@@ -882,7 +882,7 @@ impl Compiler<'_, '_> {
 
         self.push_op(OpCode::OpPopWith, &node);
         self.scope_mut().pop_with();
-        self.end_scope(&node);
+        self.cleanup_scope(&node);
     }
 
     /// Compiles pattern function arguments, such as `{ a, b }: ...`.
@@ -1005,7 +1005,7 @@ impl Compiler<'_, '_> {
         }
 
         self.compile(slot, node.body().unwrap());
-        self.end_scope(&node);
+        self.cleanup_scope(&node);
 
         // TODO: determine and insert enclosing name, if available.
 
@@ -1060,7 +1060,7 @@ impl Compiler<'_, '_> {
         let slot = self.scope_mut().declare_phantom(span);
         self.begin_scope();
         content(self, node, slot);
-        self.end_scope(node);
+        self.cleanup_scope(node);
 
         let mut thunk = self.contexts.pop().unwrap();
         optimise_tail_call(&mut thunk.lambda.chunk);
@@ -1159,41 +1159,18 @@ impl Compiler<'_, '_> {
 
     /// 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");
-
-        // If this scope poisoned any builtins or special identifiers,
-        // they need to be reset.
-        let depth = self.scope().scope_depth;
-        self.scope_mut().unpoison(depth);
-
+    fn cleanup_scope<N: AstNode>(&mut self, node: &N) {
         // When ending a scope, all corresponding locals need to be
         // removed, but the value of the body needs to remain on the
         // stack. This is implemented by a separate instruction.
-        let mut pops = 0;
-
-        // TL;DR - iterate from the back while things belonging to the
-        // ended scope still exist.
-        while self.scope().locals.last().unwrap().depth == depth {
-            if let Some(local) = self.scope_mut().locals.pop() {
-                // pop the local from the stack if it was actually
-                // initialised
-                if local.initialised {
-                    pops += 1;
-                }
+        let (popcount, unused_spans) = self.scope_mut().end_scope();
 
-                // analyse whether the local was accessed during its
-                // lifetime, and emit a warning otherwise (unless the
-                // user explicitly chose to ignore it by prefixing the
-                // identifier with `_`)
-                if !local.used && !local.is_ignored() {
-                    self.emit_warning(local.span, WarningKind::UnusedBinding);
-                }
-            }
+        for span in unused_spans {
+            self.emit_warning(span, WarningKind::UnusedBinding);
         }
 
-        if pops > 0 {
-            self.push_op(OpCode::OpCloseScope(Count(pops)), node);
+        if popcount > 0 {
+            self.push_op(OpCode::OpCloseScope(Count(popcount)), node);
         }
 
         self.scope_mut().scope_depth -= 1;
diff --git a/tvix/eval/src/compiler/scope.rs b/tvix/eval/src/compiler/scope.rs
index b6fa61e7759b..7c1e65265e1d 100644
--- a/tvix/eval/src/compiler/scope.rs
+++ b/tvix/eval/src/compiler/scope.rs
@@ -188,11 +188,11 @@ impl Scope {
         self.poisoned_tokens.contains_key(name)
     }
 
-    /// "Unpoison" tokens that were poisoned at a given depth. Used
-    /// when scopes are closed.
-    pub fn unpoison(&mut self, depth: usize) {
+    /// "Unpoison" tokens that were poisoned at the current depth.
+    /// Used when scopes are closed.
+    fn unpoison(&mut self) {
         self.poisoned_tokens
-            .retain(|_, poisoned_at| *poisoned_at != depth);
+            .retain(|_, poisoned_at| *poisoned_at != self.scope_depth);
     }
 
     /// Increase the `with`-stack size of this scope.
@@ -284,4 +284,45 @@ impl Scope {
 
         StackIdx(idx.0 - uninitialised_count)
     }
+
+    /// Decrease the scope depth and remove all locals still tracked
+    /// for the current scope.
+    ///
+    /// Returns the count of locals that were dropped while marked as
+    /// initialised (used by the compiler to determine whether to emit
+    /// scope cleanup operations), as well as the spans of the
+    /// definitions of unused locals (used by the compiler to emit
+    /// unused binding warnings).
+    pub fn end_scope(&mut self) -> (usize, Vec<codemap::Span>) {
+        debug_assert!(self.scope_depth != 0, "can not end top scope");
+
+        // If this scope poisoned any builtins or special identifiers,
+        // they need to be reset.
+        self.unpoison();
+
+        let mut pops = 0;
+        let mut unused_spans = vec![];
+
+        // TL;DR - iterate from the back while things belonging to the
+        // ended scope still exist.
+        while self.locals.last().unwrap().depth == self.scope_depth {
+            if let Some(local) = self.locals.pop() {
+                // pop the local from the stack if it was actually
+                // initialised
+                if local.initialised {
+                    pops += 1;
+                }
+
+                // analyse whether the local was accessed during its
+                // lifetime, and emit a warning otherwise (unless the
+                // user explicitly chose to ignore it by prefixing the
+                // identifier with `_`)
+                if !local.used && !local.is_ignored() {
+                    unused_spans.push(local.span);
+                }
+            }
+        }
+
+        (pops, unused_spans)
+    }
 }