[rust] First pass of function context analysis

This is a quick "good enough" first pass at computing function expression 
context variables. It definitely needs to be overhauled, but it's enough to make 
a lot of common cases work correctly. 

First, this PR adds a hand-rolled Visitor trait for `estree`. Long-term that 
should probably be code-generated, but there are some subtleties to it such as 
the `visit_lvalue(callback)` helper which has to be wrapped around various calls 
(or we need some other way to distinguish identifiers within lvalues from 
identifiers within rvaluess). So for now it makes sense to hand-roll it until we 
are more confident in exactly how it should work. 

Given that visitor, i was able to port part of the existing 
`gatherCapturedDeps()` to Rust with some modifications. Note that we assume 
_something_ has run name resolution on the estree to match up identifiers to the 
declaration they refer to. But we don't store information about parent scopes so 
we can't walk up to check where things were defined. Instead we do the 
following: 

* Build up a list of all referenced identifiers 

* Also build up a set of bindings defined in the function itself 

* After visiting, filter our the first list to only include identifiers not 
defined by the function itself, and to eliminate duplicates. 

This covers the majority of cases: the most obvious gap is nested function 
expressions though actually that should just work.

compiler/forget/crates/build-hir/src/build.rs

@@ -1,19 +1,20 @@
-use bumpalo::{
-    boxed::Box,
-    collections::{String, Vec},
-};
+use std::collections::HashSet;
+
+use bumpalo::{boxed::Box, collections::String};
 use estree::{
     AssignmentTarget, BinaryExpression, BlockStatement, Expression, ForInit, ForStatement,
-    FunctionExpression, IfStatement, JsValue, Literal, Pattern, Statement, VariableDeclarationKind,
+    Function, FunctionExpression, IfStatement, JsValue, Literal, Pattern, Statement,
+    VariableDeclarationKind,
 };
 use hir::{
-    ArrayElement, BlockKind, BranchTerminal, Environment, ForTerminal, Function, GotoKind,
-    IdentifierOperand, InstrIx, InstructionKind, InstructionValue, LValue, LoadGlobal, LoadLocal,
-    Operand, PrimitiveValue, TerminalValue,
+    ArrayElement, BlockKind, BranchTerminal, Environment, ForTerminal, GotoKind, IdentifierOperand,
+    InstrIx, InstructionKind, InstructionValue, LValue, LoadGlobal, LoadLocal, Operand,
+    PrimitiveValue, TerminalValue,
 };
 
 use crate::{
     builder::{Binding, Builder, LoopScope},
+    context::get_context_identifiers,
     error::DiagnosticError,
     BuildDiagnostic, ErrorSeverity,
 };
@@ -26,8 +27,8 @@ use crate::{
 /// that is not yet supported.
 pub fn build<'a>(
     env: &'a Environment<'a>,
-    fun: estree::Function,
-) -> Result<Box<'a, Function<'a>>, BuildDiagnostic> {
+    fun: Function,
+) -> Result<Box<'a, hir::Function<'a>>, BuildDiagnostic> {
     let mut builder = Builder::new(env);
 
     match fun.body {
@@ -78,7 +79,7 @@ pub fn build<'a>(
     );
 
     let body = builder.build()?;
-    Ok(env.box_new(Function {
+    Ok(env.box_new(hir::Function {
         id: fun
             .id
             .map(|id| String::from_str_in(&id.name, &env.allocator)),
@@ -179,7 +180,7 @@ fn lower_statement<'a>(
                 } else {
                     if let Pattern::Identifier(id) = declaration.id {
                         // TODO: handle unbound variables
-                        let binding = builder.resolve_binding(&id)?;
+                        let binding = builder.resolve_identifier(&id)?;
                         let identifier = match binding {
                             Binding::Local(identifier) => identifier,
                             _ => {
@@ -348,7 +349,7 @@ fn lower_expression<'a>(
     let value = match expr {
         Expression::Identifier(expr) => {
             // TODO: handle unbound variables
-            let binding = builder.resolve_binding(&expr)?;
+            let binding = builder.resolve_identifier(&expr)?;
             match binding {
                 Binding::Local(identifier) => {
                     let place = IdentifierOperand {
@@ -425,13 +426,7 @@ fn lower_expression<'a>(
         }
 
         Expression::FunctionExpression(expr) => {
-            let FunctionExpression { function, .. } = *expr;
-            let fun = build(env, function)?;
-            InstructionValue::Function(hir::FunctionExpression {
-                // TODO: collect dependencies!
-                dependencies: env.vec_new(),
-                lowered_function: fun,
-            })
+            InstructionValue::Function(lower_function(env, builder, *expr)?)
         }
 
         _ => todo!("Lower expr {expr:#?}"),
@@ -439,6 +434,38 @@ fn lower_expression<'a>(
     Ok(builder.push(value))
 }
 
+fn lower_function<'a>(
+    env: &'a Environment<'a>,
+    builder: &mut Builder<'a>,
+    expr: FunctionExpression,
+) -> Result<hir::FunctionExpression<'a>, BuildDiagnostic> {
+    let FunctionExpression { function, .. } = expr;
+    let context_identifiers = get_context_identifiers(env, &function);
+    let mut context = env.vec_new();
+    let mut seen = HashSet::new();
+    for identifier in context_identifiers {
+        match builder.resolve_identifier(identifier)? {
+            Binding::Local(identifier) => {
+                if !seen.insert(identifier.id) {
+                    continue;
+                }
+                context.push(IdentifierOperand {
+                    effect: None,
+                    identifier,
+                });
+            }
+            _ => {}
+        }
+    }
+    let mut fun = build(env, function)?;
+    fun.context = context;
+    Ok(hir::FunctionExpression {
+        // TODO: collect dependencies!
+        dependencies: env.vec_new(),
+        lowered_function: fun,
+    })
+}
+
 fn lower_assignment<'a>(
     env: &'a Environment<'a>,
     builder: &mut Builder<'a>,
@@ -470,7 +497,7 @@ fn lower_identifier_for_assignment<'a>(
     _kind: InstructionKind,
     identifier: estree::Identifier,
 ) -> Result<IdentifierOperand<'a>, BuildDiagnostic> {
-    let binding = builder.resolve_binding(&identifier)?;
+    let binding = builder.resolve_identifier(&identifier)?;
     match binding {
         Binding::Module(..) | Binding::Global => Err(BuildDiagnostic::new(
             DiagnosticError::ReassignedGlobal,

compiler/forget/crates/build-hir/src/builder.rs

@@ -325,7 +325,7 @@ impl<'a> Builder<'a> {
         ))
     }
 
-    pub(crate) fn resolve_binding(
+    pub(crate) fn resolve_identifier(
         &mut self,
         identifier: &estree::Identifier,
     ) -> Result<Binding<'a>, BuildDiagnostic> {

compiler/forget/crates/build-hir/src/context.rs

@@ -0,0 +1,102 @@
+use std::collections::HashSet;
+
+use estree::{Binding, BindingId, Expression, ExpressionOrSuper, Function, Identifier, Visitor};
+use hir::Environment;
+
+pub(crate) fn get_context_identifiers<'a, 'ast>(
+    _env: &'a Environment<'a>,
+    function: &'ast Function,
+) -> Vec<&'ast Identifier> {
+    let mut visitor = ContextVisitor::new();
+    visitor.visit_function(function);
+    let ContextVisitor {
+        free_variables,
+        defined,
+        ..
+    } = visitor;
+    free_variables
+        .into_iter()
+        .filter(|identifier| match &identifier.binding {
+            Some(Binding::Local(id)) if !defined.contains(id) => true,
+            _ => false,
+        })
+        .collect()
+}
+
+struct ContextVisitor<'ast> {
+    free_variables: Vec<&'ast Identifier>,
+    defined: HashSet<BindingId>,
+    lvalue: bool,
+}
+
+impl<'ast> ContextVisitor<'ast> {
+    fn new() -> Self {
+        Self {
+            free_variables: Default::default(),
+            defined: Default::default(),
+            lvalue: false,
+        }
+    }
+}
+
+impl<'ast> Visitor<'ast> for ContextVisitor<'ast> {
+    fn visit_lvalue<F>(&mut self, f: F)
+    where
+        F: FnOnce(&mut Self) -> (),
+    {
+        let prev_lvalue = self.lvalue;
+        self.lvalue = true;
+        f(self);
+        self.lvalue = prev_lvalue;
+    }
+
+    fn visit_function(&mut self, function: &'ast Function) {
+        self.visit_lvalue(|visitor| {
+            for param in &function.params {
+                visitor.visit_pattern(param);
+            }
+        });
+        self.default_visit_function(function);
+    }
+
+    fn visit_expression(&mut self, expr: &'ast Expression) {
+        if let Expression::AssignmentExpression(expr) = expr {
+            self.visit_lvalue(|visitor| visitor.visit_assignment_target(&expr.left));
+            self.visit_expression(&expr.right);
+            return;
+        }
+        let mut object = expr;
+        while let Expression::MemberExpression(expr) = object {
+            match &expr.object {
+                ExpressionOrSuper::Super(_) => return,
+                ExpressionOrSuper::Expression(expr) => {
+                    object = expr;
+                }
+            }
+            if expr.computed {
+                self.visit_expression(&expr.property);
+            }
+        }
+        if let Expression::Identifier(identifier) = object {
+            self.visit_identifier(identifier);
+        } else {
+            self.default_visit_expression(expr);
+        }
+    }
+
+    fn visit_identifier(&mut self, identifier: &'ast Identifier) {
+        let binding = identifier.binding.unwrap();
+        match binding {
+            Binding::Local(binding_id) => {
+                if self.lvalue {
+                    // println!("lvalue {identifier:?}");
+                    self.defined.insert(binding_id);
+                } else {
+                    // println!("rvalue {identifier:?}");
+                    self.free_variables.push(identifier);
+                }
+            }
+            _ => {}
+        }
+    }
+}

compiler/forget/crates/build-hir/src/lib.rs

@@ -1,5 +1,6 @@
 mod build;
 mod builder;
+mod context;
 mod error;
 
 pub use build::build;

compiler/forget/crates/estree/src/lib.rs

@@ -3,11 +3,13 @@ mod generated;
 mod generated_extensions;
 mod js_value;
 mod range;
+mod visit;
 
 pub use binding::{Binding, BindingId};
 pub use generated::*;
 pub use js_value::JsValue;
 pub use range::SourceRange;
+pub use visit::*;
 
 #[cfg(test)]
 mod tests {

compiler/forget/crates/estree/src/visit.rs

@@ -0,0 +1,296 @@
+use crate::{
+    AssignmentTarget, Expression, ExpressionOrSpread, ForInInit, ForInit, Function, FunctionBody,
+    Identifier, ImportDeclarationSpecifier, ImportOrExportDeclaration, Literal, ModuleItem,
+    Pattern, Program, Statement, SwitchCase, VariableDeclarator,
+};
+
+/// Trait for visiting an estree
+pub trait Visitor<'ast> {
+    fn visit_lvalue<F>(&mut self, f: F)
+    where
+        F: FnOnce(&mut Self) -> (),
+    {
+        f(self);
+    }
+
+    fn visit_program(&mut self, program: &'ast Program) {
+        for item in &program.body {
+            self.visit_module_item(item);
+        }
+    }
+
+    fn visit_function(&mut self, function: &'ast Function) {
+        self.default_visit_function(function);
+    }
+
+    fn default_visit_function(&mut self, function: &'ast Function) {
+        // todo: params and such
+        match &function.body {
+            Some(FunctionBody::BlockStatement(body)) => {
+                for stmt in &body.body {
+                    self.visit_statement(stmt)
+                }
+            }
+            Some(FunctionBody::Expression(body)) => self.visit_expression(body),
+            None => {}
+        }
+    }
+
+    fn visit_module_item(&mut self, item: &'ast ModuleItem) {
+        match item {
+            ModuleItem::Statement(item) => self.visit_statement(item),
+            ModuleItem::ImportOrExportDeclaration(item) => {
+                self.visit_import_or_export_declaration(item)
+            }
+        }
+    }
+
+    fn visit_import_or_export_declaration(&mut self, declaration: &'ast ImportOrExportDeclaration) {
+        match declaration {
+            ImportOrExportDeclaration::ImportDeclaration(declaration) => {
+                self.visit_lvalue(|visitor| {
+                    for specifier in &declaration.specifiers {
+                        visitor.visit_import_declaration_specifier(specifier, &declaration.source)
+                    }
+                });
+                self.visit_import_source(&declaration.source);
+            }
+        }
+    }
+
+    fn visit_import_declaration_specifier(
+        &mut self,
+        specifier: &'ast ImportDeclarationSpecifier,
+        _source: &'ast Literal,
+    ) {
+        match specifier {
+            ImportDeclarationSpecifier::ImportSpecifier(specifier) => {
+                self.visit_identifier(&specifier.local);
+            }
+            ImportDeclarationSpecifier::ImportDefaultSpecifier(specifier) => {
+                self.visit_identifier(&specifier.local);
+            }
+            ImportDeclarationSpecifier::ImportNamespaceSpecifier(specifier) => {
+                self.visit_identifier(&specifier.local);
+            }
+        }
+    }
+
+    fn visit_statement(&mut self, stmt: &'ast Statement) {
+        self.default_visit_statement(stmt);
+    }
+
+    fn default_visit_statement(&mut self, stmt: &'ast Statement) {
+        match stmt {
+            Statement::BlockStatement(stmt) => {
+                for stmt in &stmt.body {
+                    self.visit_statement(stmt)
+                }
+            }
+            Statement::BreakStatement(_stmt) => {
+                // todo
+            }
+            Statement::ContinueStatement(_stmt) => {
+                // todo
+            }
+            Statement::DebuggerStatement(_stmt) => {
+                // todo
+            }
+            Statement::DoWhileStatement(stmt) => {
+                self.visit_statement(&stmt.body);
+                self.visit_expression(&stmt.test);
+            }
+            Statement::EmptyStatement(_stmt) => {
+                // nothing to do
+            }
+            Statement::ExpressionStatement(stmt) => {
+                self.visit_expression(&stmt.expression);
+            }
+            Statement::ForInStatement(stmt) => {
+                self.visit_for_in_init(&stmt.left);
+                self.visit_expression(&stmt.right);
+                self.visit_statement(&stmt.body);
+            }
+            Statement::ForOfStatement(stmt) => {
+                self.visit_for_in_init(&stmt.left);
+                self.visit_expression(&stmt.right);
+                self.visit_statement(&stmt.body);
+            }
+            Statement::ForStatement(stmt) => {
+                if let Some(init) = &stmt.init {
+                    self.visit_for_init(init);
+                }
+                if let Some(test) = &stmt.test {
+                    self.visit_expression(test);
+                }
+                if let Some(update) = &stmt.update {
+                    self.visit_expression(update);
+                }
+                self.visit_statement(&stmt.body);
+            }
+            Statement::FunctionDeclaration(stmt) => {
+                self.visit_function(&stmt.function);
+            }
+            Statement::IfStatement(stmt) => {
+                self.visit_expression(&stmt.test);
+                self.visit_statement(&stmt.consequent);
+                if let Some(alternate) = &stmt.alternate {
+                    self.visit_statement(alternate);
+                }
+            }
+            Statement::LabeledStatement(stmt) => {
+                self.visit_statement(&stmt.body);
+            }
+            Statement::ReturnStatement(stmt) => {
+                if let Some(argument) = &stmt.argument {
+                    self.visit_expression(argument);
+                }
+            }
+            Statement::SwitchStatement(stmt) => {
+                self.visit_expression(&stmt.discriminant);
+                for case_ in &stmt.cases {
+                    self.visit_case(case_);
+                }
+            }
+            Statement::ThrowStatement(stmt) => {
+                self.visit_expression(&stmt.argument);
+            }
+            Statement::TryStatement(stmt) => {
+                for item in &stmt.block.body {
+                    self.visit_statement(item);
+                }
+                if let Some(handler) = &stmt.handler {
+                    self.visit_lvalue(|visitor| visitor.visit_pattern(&handler.param));
+                    for item in &handler.body.body {
+                        self.visit_statement(item);
+                    }
+                }
+                if let Some(finalizer) = &stmt.finalizer {
+                    for item in &finalizer.body {
+                        self.visit_statement(item);
+                    }
+                }
+            }
+            Statement::VariableDeclaration(stmt) => {
+                for decl in &stmt.declarations {
+                    self.visit_variable_declarator(decl);
+                }
+            }
+            Statement::WhileStatement(stmt) => {
+                self.visit_expression(&stmt.test);
+                self.visit_statement(&stmt.body);
+            }
+            Statement::WithStatement(stmt) => {
+                self.visit_expression(&stmt.object);
+                self.visit_statement(&stmt.body);
+            }
+        }
+    }
+
+    fn visit_case(&mut self, case_: &'ast SwitchCase) {
+        if let Some(test) = &case_.test {
+            self.visit_expression(test);
+        }
+        for stmt in &case_.consequent {
+            self.visit_statement(stmt)
+        }
+    }
+
+    fn visit_for_init(&mut self, init: &'ast ForInit) {
+        match init {
+            ForInit::Expression(init) => {
+                self.visit_expression(init);
+            }
+            ForInit::VariableDeclaration(init) => {
+                for decl in &init.declarations {
+                    self.visit_variable_declarator(decl);
+                }
+            }
+        }
+    }
+
+    fn visit_for_in_init(&mut self, init: &'ast ForInInit) {
+        match init {
+            ForInInit::Pattern(init) => {
+                self.visit_pattern(init);
+            }
+            ForInInit::VariableDeclaration(init) => {
+                for decl in &init.declarations {
+                    self.visit_variable_declarator(decl);
+                }
+            }
+        }
+    }
+
+    fn visit_pattern(&mut self, pattern: &'ast Pattern) {
+        match pattern {
+            Pattern::Identifier(pattern) => self.visit_identifier(pattern),
+        }
+    }
+
+    fn visit_variable_declarator(&mut self, decl: &'ast VariableDeclarator) {
+        self.visit_lvalue(|visitor| {
+            visitor.visit_pattern(&decl.id);
+        });
+        if let Some(init) = &decl.init {
+            self.visit_expression(init);
+        }
+    }
+
+    fn visit_assignment_target(&mut self, target: &'ast AssignmentTarget) {
+        match target {
+            AssignmentTarget::Expression(target) => {
+                self.visit_expression(target);
+            }
+            AssignmentTarget::Pattern(target) => self.visit_pattern(target),
+        }
+    }
+
+    fn visit_expression(&mut self, expr: &'ast Expression) {
+        self.visit_expression(expr);
+    }
+
+    fn default_visit_expression(&mut self, expr: &'ast Expression) {
+        match expr {
+            Expression::ArrayExpression(expr) => {
+                for item in &expr.elements {
+                    match item {
+                        Some(ExpressionOrSpread::SpreadElement(item)) => {
+                            self.visit_expression(&item.argument)
+                        }
+                        Some(ExpressionOrSpread::Expression(item)) => self.visit_expression(item),
+                        _ => {}
+                    }
+                }
+            }
+            Expression::AssignmentExpression(expr) => {
+                self.visit_lvalue(|visitor| visitor.visit_assignment_target(&expr.left));
+                self.visit_expression(&expr.right);
+            }
+            Expression::BinaryExpression(expr) => {
+                self.visit_expression(&expr.left);
+                self.visit_expression(&expr.right);
+            }
+            Expression::Identifier(expr) => {
+                self.visit_identifier(expr);
+            }
+            Expression::Literal(expr) => self.visit_literal(expr),
+            Expression::FunctionExpression(expr) => self.visit_function(&expr.function),
+            _ => {
+                todo!("more expression types")
+            }
+        }
+    }
+
+    fn visit_identifier(&mut self, _identifier: &'ast Identifier) {
+        // nothing to do unless overridden
+    }
+
+    fn visit_import_source(&mut self, literal: &'ast Literal) {
+        self.visit_literal(literal);
+    }
+
+    fn visit_literal(&mut self, _literal: &'ast Literal) {
+        // nothing to do unless overridden
+    }
+}

compiler/forget/crates/fixtures/tests/fixtures_test.rs

@@ -35,6 +35,11 @@ fn fixtures() {
                     }
                     match build(&environment, fun.function) {
                         Ok(mut fun) => {
+                            // println!("{fun:#?}");
+                            // let mut out = String::new();
+                            // fun.print(&fun.body, &mut out).unwrap();
+                            // println!("{out}");
+
                             enter_ssa(&environment, &mut fun).unwrap();
                             eliminate_redundant_phis(&environment, &mut fun);
                             constant_propagation(&environment, &mut fun);

compiler/forget/crates/fixtures/tests/snapshots/fixtures_test__fixtures@function-expressions.js.snap

@@ -27,7 +27,7 @@ entry bb0
 bb0 (block)
   [0] #0 = 2
   [1] #1 = StoreLocal Const unknown y$7 = unknown #0
-  [2] #2 = Function @deps[] @context[]:
+  [2] #2 = Function @deps[] @context[unknown y$7]:
       function foo(
         unknown x$8,
       )
@@ -50,7 +50,7 @@ bb0 (block)
       bb2 (block)
         predecessors: bb3
         [12] #12 = LoadLocal unknown x$8
-        [13] #13 = LoadLocal unknown y$0
+        [13] #13 = 2
         [14] #14 = Binary unknown #12 + unknown #13
         [15] #15 = 8
         [16] #16 = Binary unknown #14 + unknown #15

compiler/forget/crates/hir-ssa/src/eliminate_redundant_phis.rs

@@ -1,6 +1,6 @@
 use std::collections::{HashMap, HashSet};
 
-use hir::{BlockId, Environment, Function, Identifier, IdentifierId, InstructionValue, HIR};
+use hir::{BlockId, Environment, Function, Identifier, IdentifierId, InstructionValue};
 use utils::RetainMut;
 
 /// Pass to eliminate redundant phi nodes:

compiler/forget/crates/hir-ssa/src/enter.rs

@@ -13,7 +13,7 @@ use thiserror::Error;
 pub struct SSAError;
 
 pub fn enter_ssa<'a>(env: &Environment<'a>, fun: &mut Function<'a>) -> Result<(), SSAError> {
-    // assert_eq!(fun.context.is_empty())
+    assert!(fun.context.is_empty());
     enter_ssa_impl(env, fun, None)
 }
 
@@ -63,10 +63,11 @@ fn visit_instructions<'a, 'e, 'f>(
                 let context_defs: IndexMap<IdentifierId, Identifier> = fun
                     .lowered_function
                     .context
-                    .iter()
-                    .map(|id| {
-                        let identifier = builder.get_id_at(block.id, &id.identifier);
-                        (id.identifier.id, identifier)
+                    .iter_mut()
+                    .map(|identifier| {
+                        let old_id = identifier.identifier.id;
+                        builder.visit_load(identifier);
+                        (old_id, identifier.identifier.clone())
                     })
                     .collect();
                 enter_ssa_impl(env, &mut fun.lowered_function, Some(context_defs))?;

compiler/forget/crates/hir/src/print.rs

@@ -160,12 +160,12 @@ impl<'a> Print<'a> for InstructionValue<'a> {
                     dep.print(hir, out)?;
                 }
                 write!(out, "] @context[")?;
-                // for (ix, dep) in value.lowered_function.context.iter().enumerate() {
-                //     if ix != 0 {
-                //         write!(out, ", ")?;
-                //     }
-                //     dep.print(hir, out)?;
-                // }
+                for (ix, dep) in value.lowered_function.context.iter().enumerate() {
+                    if ix != 0 {
+                        write!(out, ", ")?;
+                    }
+                    dep.print(hir, out)?;
+                }
                 writeln!(out, "]:")?;
                 let mut inner_output = String::new();
                 value