MIR: fix free-vars lexical scoping

src/mir/builder/vars/free_vars.rs

@@ -1,71 +1,141 @@
 use crate::ast::ASTNode;
 use std::collections::HashSet;
 
-/// Collect free variables used in `node` into `used`, excluding names present in `locals`.
-/// `locals` is updated as new local declarations are encountered.
 #[allow(dead_code)]
 pub(in crate::mir::builder) fn collect_free_vars(
     node: &ASTNode,
     used: &mut HashSet<String>,
     locals: &mut HashSet<String>,
 ) {
+    let mut collector = FreeVarCollector::new(used, locals);
+    collector.walk(node, true);
+    *locals = collector.take_root_locals();
+}
+
+struct FreeVarCollector<'a> {
+    used: &'a mut HashSet<String>,
+    root_locals: HashSet<String>,
+    scope_stack: Vec<HashSet<String>>,
+}
+
+impl<'a> FreeVarCollector<'a> {
+    fn new(used: &'a mut HashSet<String>, root_locals: &HashSet<String>) -> Self {
+        Self {
+            used,
+            root_locals: root_locals.clone(),
+            scope_stack: vec![root_locals.clone()],
+        }
+    }
+
+    fn take_root_locals(self) -> HashSet<String> {
+        self.root_locals
+    }
+
+    fn is_declared(&self, name: &str) -> bool {
+        self.scope_stack.iter().rev().any(|s| s.contains(name))
+    }
+
+    fn declare_in_current_scope(&mut self, name: &str) {
+        if let Some(scope) = self.scope_stack.last_mut() {
+            scope.insert(name.to_string());
+        }
+        if self.scope_stack.len() == 1 {
+            self.root_locals.insert(name.to_string());
+        }
+    }
+
+    fn with_child_scope(&mut self, f: impl FnOnce(&mut Self)) {
+        self.scope_stack.push(HashSet::new());
+        f(self);
+        self.scope_stack.pop();
+    }
+
+    fn walk_block(&mut self, statements: &[ASTNode]) {
+        for st in statements {
+            self.walk(st, false);
+        }
+    }
+
+    fn walk(&mut self, node: &ASTNode, is_root: bool) {
         match node {
             ASTNode::Variable { name, .. } => {
-            if name != "me" && name != "this" && !locals.contains(name) {
+                if name != "me" && name != "this" && !self.is_declared(name) {
-                used.insert(name.clone());
+                    self.used.insert(name.clone());
+                }
+            }
+            ASTNode::Local {
+                variables,
+                initial_values,
+                ..
+            } => {
+                for init in initial_values {
+                    if let Some(expr) = init {
+                        self.walk(expr, false);
                     }
                 }
-        ASTNode::Local { variables, .. } => {
                 for v in variables {
-                locals.insert(v.clone());
+                    self.declare_in_current_scope(v);
+                }
+            }
+            ASTNode::Outbox {
+                variables,
+                initial_values,
+                ..
+            } => {
+                for init in initial_values {
+                    if let Some(expr) = init {
+                        self.walk(expr, false);
+                    }
+                }
+                for v in variables {
+                    self.declare_in_current_scope(v);
                 }
             }
             ASTNode::Assignment { target, value, .. } => {
-            collect_free_vars(target, used, locals);
+                self.walk(target, false);
-            collect_free_vars(value, used, locals);
+                self.walk(value, false);
             }
-        // Phase 152-A: Grouped assignment expression
             ASTNode::GroupedAssignmentExpr { rhs, .. } => {
-            collect_free_vars(rhs, used, locals);
+                self.walk(rhs, false);
             }
             ASTNode::BinaryOp { left, right, .. } => {
-            collect_free_vars(left, used, locals);
+                self.walk(left, false);
-            collect_free_vars(right, used, locals);
+                self.walk(right, false);
             }
             ASTNode::UnaryOp { operand, .. } => {
-            collect_free_vars(operand, used, locals);
+                self.walk(operand, false);
             }
             ASTNode::MethodCall {
                 object, arguments, ..
             } => {
-            collect_free_vars(object, used, locals);
+                self.walk(object, false);
                 for a in arguments {
-                collect_free_vars(a, used, locals);
+                    self.walk(a, false);
                 }
             }
             ASTNode::FunctionCall { arguments, .. } => {
                 for a in arguments {
-                collect_free_vars(a, used, locals);
+                    self.walk(a, false);
                 }
             }
             ASTNode::Call {
                 callee, arguments, ..
             } => {
-            collect_free_vars(callee, used, locals);
+                self.walk(callee, false);
                 for a in arguments {
-                collect_free_vars(a, used, locals);
+                    self.walk(a, false);
                 }
             }
             ASTNode::FieldAccess { object, .. } => {
-            collect_free_vars(object, used, locals);
+                self.walk(object, false);
             }
             ASTNode::Index { target, index, .. } => {
-            collect_free_vars(target, used, locals);
+                self.walk(target, false);
-            collect_free_vars(index, used, locals);
+                self.walk(index, false);
             }
             ASTNode::New { arguments, .. } => {
                 for a in arguments {
-                collect_free_vars(a, used, locals);
+                    self.walk(a, false);
                 }
             }
             ASTNode::If {
@@ -74,23 +144,17 @@ pub(in crate::mir::builder) fn collect_free_vars(
                 else_body,
                 ..
             } => {
-            collect_free_vars(condition, used, locals);
+                self.walk(condition, false);
-            for st in then_body {
+                self.with_child_scope(|this| this.walk_block(then_body));
-                collect_free_vars(st, used, locals);
-            }
                 if let Some(eb) = else_body {
-                for st in eb {
+                    self.with_child_scope(|this| this.walk_block(eb));
-                    collect_free_vars(st, used, locals);
-                }
                 }
             }
             ASTNode::Loop {
                 condition, body, ..
             } => {
-            collect_free_vars(condition, used, locals);
+                self.walk(condition, false);
-            for st in body {
+                self.with_child_scope(|this| this.walk_block(body));
-                collect_free_vars(st, used, locals);
-            }
             }
             ASTNode::TryCatch {
                 try_body,
@@ -98,33 +162,27 @@ pub(in crate::mir::builder) fn collect_free_vars(
                 finally_body,
                 ..
             } => {
-            for st in try_body {
+                self.with_child_scope(|this| this.walk_block(try_body));
-                collect_free_vars(st, used, locals);
-            }
                 for c in catch_clauses {
-                for st in &c.body {
+                    self.with_child_scope(|this| this.walk_block(&c.body));
-                    collect_free_vars(st, used, locals);
-                }
                 }
                 if let Some(fb) = finally_body {
-                for st in fb {
+                    self.with_child_scope(|this| this.walk_block(fb));
-                    collect_free_vars(st, used, locals);
-                }
                 }
             }
             ASTNode::Throw { expression, .. } => {
-            collect_free_vars(expression, used, locals);
+                self.walk(expression, false);
             }
             ASTNode::Print { expression, .. } => {
-            collect_free_vars(expression, used, locals);
+                self.walk(expression, false);
             }
             ASTNode::Return { value, .. } => {
                 if let Some(v) = value {
-                collect_free_vars(v, used, locals);
+                    self.walk(v, false);
                 }
             }
             ASTNode::AwaitExpression { expression, .. } => {
-            collect_free_vars(expression, used, locals);
+                self.walk(expression, false);
             }
             ASTNode::MatchExpr {
                 scrutinee,
@@ -132,26 +190,104 @@ pub(in crate::mir::builder) fn collect_free_vars(
                 else_expr,
                 ..
             } => {
-            collect_free_vars(scrutinee, used, locals);
+                self.walk(scrutinee, false);
                 for (_, e) in arms {
-                collect_free_vars(e, used, locals);
+                    self.walk(e, false);
                 }
-            collect_free_vars(else_expr, used, locals);
+                self.walk(else_expr, false);
             }
             ASTNode::Program { statements, .. } => {
-            for st in statements {
+                if is_root {
-                collect_free_vars(st, used, locals);
+                    self.walk_block(statements);
+                } else {
+                    self.with_child_scope(|this| this.walk_block(statements));
+                }
+            }
+            ASTNode::ScopeBox { body, .. } => {
+                if is_root {
+                    self.walk_block(body);
+                } else {
+                    self.with_child_scope(|this| this.walk_block(body));
                 }
             }
             ASTNode::FunctionDeclaration { params, body, .. } => {
-            let mut inner = locals.clone();
+                self.with_child_scope(|this| {
                     for p in params {
-                inner.insert(p.clone());
+                        this.declare_in_current_scope(p);
-            }
-            for st in body {
-                collect_free_vars(st, used, &mut inner);
                     }
+                    this.walk_block(body);
+                });
             }
             _ => {}
         }
     }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::collect_free_vars;
+    use crate::ast::ASTNode;
+    use crate::ast::Span;
+    use std::collections::HashSet;
+
+    fn var(name: &str) -> ASTNode {
+        ASTNode::Variable {
+            name: name.to_string(),
+            span: Span::unknown(),
+        }
+    }
+
+    fn local(name: &str) -> ASTNode {
+        ASTNode::Local {
+            variables: vec![name.to_string()],
+            initial_values: vec![None],
+            span: Span::unknown(),
+        }
+    }
+
+    fn block(statements: Vec<ASTNode>) -> ASTNode {
+        ASTNode::Program {
+            statements,
+            span: Span::unknown(),
+        }
+    }
+
+    fn scopebox(statements: Vec<ASTNode>) -> ASTNode {
+        ASTNode::ScopeBox {
+            body: statements,
+            span: Span::unknown(),
+        }
+    }
+
+    #[test]
+    fn block_local_does_not_leak_to_outer_scope() {
+        let node = block(vec![block(vec![local("y")]), var("y")]);
+        let mut used = HashSet::new();
+        let mut locals = HashSet::new();
+        collect_free_vars(&node, &mut used, &mut locals);
+        assert!(used.contains("y"));
+    }
+
+    #[test]
+    fn scopebox_local_does_not_leak_to_outer_scope() {
+        let node = block(vec![scopebox(vec![local("y")]), var("y")]);
+        let mut used = HashSet::new();
+        let mut locals = HashSet::new();
+        collect_free_vars(&node, &mut used, &mut locals);
+        assert!(used.contains("y"));
+    }
+
+    #[test]
+    fn local_initializer_is_counted_as_read_before_declare() {
+        let node = block(vec![ASTNode::Local {
+            variables: vec!["x".to_string()],
+            initial_values: vec![Some(Box::new(var("y")))],
+            span: Span::unknown(),
+        }]);
+        let mut used = HashSet::new();
+        let mut locals = HashSet::new();
+        collect_free_vars(&node, &mut used, &mut locals);
+        assert!(used.contains("y"));
+        assert!(locals.contains("x"));
+    }
+}