Refactor: Move optimizer related types to dedicated module

This commit reorganizes the optimizer code by creating a new `optimizer`
module. The `Optimizer` struct and its related logic remain in
`optimizer/optimizer.rs`, while `SubstitutionMap` and `UsageInfo` are
moved to `optimizer/substitution_map.rs`.

This change improves code organization and makes it easier to manage the
optimizer's components.

src/ast/compiler/optimizer/mod.rs

@@ -0,0 +1,5 @@
+pub mod optimizer;
+pub mod substitution_map;
+
+pub use optimizer::Optimizer;
+pub use substitution_map::{SubstitutionMap, UsageInfo};

src/ast/compiler/optimizer/optimizer.rs

@@ -8,6 +8,8 @@ use std::cell::RefCell;
 use std::collections::{HashMap, HashSet};
 use std::rc::Rc;
 
+use super::substitution_map::{SubstitutionMap, UsageInfo};
+
 pub struct Optimizer {
     pub enabled: bool,
     max_passes: usize,
@@ -44,23 +46,6 @@ impl PathTracker {
     }
 }
 
-#[derive(Default)]
-struct UsageInfo {
-    used: HashSet<Address>,
-    assigned: HashSet<Address>,
-    used_identities: HashSet<crate::ast::types::Identity>,
-}
-
-impl UsageInfo {
-    fn is_assigned(&self, addr: &Address) -> bool {
-        self.assigned.contains(addr)
-    }
-
-    fn is_used(&self, addr: &Address) -> bool {
-        self.used.contains(addr)
-    }
-}
-
 impl Optimizer {
     pub fn new(enabled: bool) -> Self {
         Self {
@@ -108,23 +93,6 @@ impl Optimizer {
         current
     }
 
-    fn collect_pattern_usage(&self, node: &AnalyzedNode, info: &mut UsageInfo) {
-        match &node.kind {
-            BoundKind::Define { addr, .. } => {
-                info.assigned.insert(*addr);
-            }
-            BoundKind::Set { addr, .. } => {
-                info.assigned.insert(*addr);
-            }
-            BoundKind::Tuple { elements } => {
-                for el in elements {
-                    self.collect_pattern_usage(el, info);
-                }
-            }
-            _ => {}
-        }
-    }
-
     fn visit_node(
         &self,
         node: AnalyzedNode,
@@ -397,7 +365,7 @@ impl Optimizer {
                 let mut info = UsageInfo::default();
                 if !exprs.is_empty() {
                     for e in exprs {
-                        self.collect_usage(e, &mut info);
+                        info.collect(e);
                     }
                 }
 
@@ -470,7 +438,7 @@ impl Optimizer {
                     };
 
                 let mut info = UsageInfo::default();
-                self.collect_usage(&node, &mut info);
+                info.collect(&node);
 
                 let mut new_upvalues = Vec::new();
                 let mut mapping = Vec::new();
@@ -526,7 +494,7 @@ impl Optimizer {
                 let pat_opt = Box::new(self.visit_node((**pattern).clone(), sub, path));
 
                 let mut info = UsageInfo::default();
-                self.collect_pattern_usage(&pat_opt, &mut info);
+                info.collect_pattern(&pat_opt);
                 for addr in info.assigned {
                     sub.remove_value(&addr);
                 }
@@ -712,7 +680,7 @@ impl Optimizer {
         path: &mut PathTracker,
     ) -> Option<AnalyzedNode> {
         let mut body_usage = UsageInfo::default();
-        self.collect_usage(&body, &mut body_usage);
+        body_usage.collect(&body);
 
         let mut arg_vals = Vec::new();
         self.flatten_tuple(args.clone(), &mut arg_vals);
@@ -852,286 +820,5 @@ impl Optimizer {
             _ => {}
         }
     }
-
-    fn collect_usage(&self, node: &AnalyzedNode, info: &mut UsageInfo) {
-        match &node.kind {
-            BoundKind::Destructure { pattern, value } => {
-                self.collect_pattern_usage(pattern, info);
-                self.collect_usage(value, info);
-            }
-            BoundKind::Constant(v) => {
-                if let Value::Object(obj) = v
-                    && let Some(closure) = obj.as_any().downcast_ref::<Closure>()
-                {
-                    info.used_identities
-                        .insert(closure.function_node.identity.clone());
-                    self.collect_usage(&closure.function_node, info);
-                }
-            }
-            BoundKind::Get { addr, .. } => {
-                info.used.insert(*addr);
-            }
-            BoundKind::Set { addr, value } => {
-                info.assigned.insert(*addr);
-                self.collect_usage(value, info);
-            }
-            BoundKind::Lambda {
-                params,
-                body,
-                upvalues,
-                ..
-            } => {
-                info.used_identities.insert(node.identity.clone());
-
-                let mut inner_info = UsageInfo::default();
-                self.collect_usage(params, &mut inner_info);
-                self.collect_usage(body, &mut inner_info);
-
-                // Propagate globals and identities
-                for addr in &inner_info.used {
-                    if let Address::Global(_) = addr {
-                        info.used.insert(*addr);
-                    }
-                }
-                for addr in &inner_info.assigned {
-                    if let Address::Global(_) = addr {
-                        info.assigned.insert(*addr);
-                    }
-                }
-                info.used_identities.extend(inner_info.used_identities);
-
-                // Map used upvalues to parent scope
-                for addr in &inner_info.used {
-                    if let Address::Upvalue(idx) = addr
-                        && let Some(parent_addr) = upvalues.get(idx.0 as usize)
-                    {
-                        info.used.insert(*parent_addr);
-                    }
-                }
-                // Map assigned upvalues to parent scope
-                for addr in &inner_info.assigned {
-                    if let Address::Upvalue(idx) = addr
-                        && let Some(parent_addr) = upvalues.get(idx.0 as usize)
-                    {
-                        info.assigned.insert(*parent_addr);
-                    }
-                }
-            }
-            BoundKind::Block { exprs } => {
-                for e in exprs {
-                    self.collect_usage(e, info);
-                }
-            }
-            BoundKind::If {
-                cond,
-                then_br,
-                else_br,
-            } => {
-                self.collect_usage(cond, info);
-                self.collect_usage(then_br, info);
-                if let Some(e) = else_br {
-                    self.collect_usage(e, info);
-                }
-            }
-            BoundKind::Call { callee, args } => {
-                self.collect_usage(callee, info);
-                self.collect_usage(args, info);
-            }
-            BoundKind::Tuple { elements } => {
-                for e in elements {
-                    self.collect_usage(e, info);
-                }
-            }
-            BoundKind::Record { fields } => {
-                for (k, v) in fields {
-                    self.collect_usage(k, info);
-                    self.collect_usage(v, info);
-                }
-            }
-            BoundKind::Define { value, .. } => {
-                self.collect_usage(value, info);
-            }
-            BoundKind::Expansion { bound_expanded, .. } => {
-                self.collect_usage(bound_expanded, info);
-            }
-            _ => {}
-        }
-    }
 }
 
-struct SubstitutionMap {
-    values: HashMap<Address, Value>,
-    ast_substitutions: HashMap<Address, AnalyzedNode>,
-    slot_mapping: HashMap<LocalSlot, LocalSlot>,
-    assigned: HashSet<Address>,
-    next_slot: u32,
-    used: HashSet<Address>,
-    captured_slots: HashSet<LocalSlot>,
-}
-
-impl SubstitutionMap {
-    fn new() -> Self {
-        Self {
-            values: HashMap::new(),
-            ast_substitutions: HashMap::new(),
-            slot_mapping: HashMap::new(),
-            assigned: HashSet::new(),
-            next_slot: 0,
-            used: HashSet::new(),
-            captured_slots: HashSet::new(),
-        }
-    }
-
-    fn add_ast_substitution(&mut self, addr: Address, node: AnalyzedNode) {
-        self.ast_substitutions.insert(addr, node);
-    }
-
-    fn map_slot(&mut self, old_slot: LocalSlot) -> LocalSlot {
-        if let Some(&new_slot) = self.slot_mapping.get(&old_slot) {
-            return new_slot;
-        }
-        let new_slot = LocalSlot(self.next_slot);
-        self.slot_mapping.insert(old_slot, new_slot);
-        self.next_slot += 1;
-        new_slot
-    }
-
-    fn map_address(&mut self, addr: Address) -> Address {
-        match addr {
-            Address::Local(slot) => Address::Local(self.map_slot(slot)),
-            other => other,
-        }
-    }
-
-    fn add_value(&mut self, addr: Address, val: Value) {
-        self.values.insert(addr, val);
-    }
-
-    fn get_value(&self, addr: &Address) -> Option<&Value> {
-        self.values.get(addr)
-    }
-
-    fn remove_value(&mut self, addr: &Address) {
-        self.values.remove(addr);
-    }
-
-    fn reindex_addr(&self, addr: Address, mapping: &[Option<u32>]) -> Address {
-        if let Address::Upvalue(idx) = addr
-            && let Some(res) = mapping.get(idx.0 as usize)
-            && let Some(new_idx) = res
-        {
-            Address::Upvalue(UpvalueIdx(*new_idx))
-        } else {
-            addr
-        }
-    }
-
-    fn reindex_upvalues(&self, node: AnalyzedNode, mapping: &[Option<u32>]) -> AnalyzedNode {
-        let (new_kind, metrics) = match node.kind {
-            BoundKind::Get { addr, name } => (
-                BoundKind::Get {
-                    addr: self.reindex_addr(addr, mapping),
-                    name,
-                },
-                node.ty.clone(),
-            ),
-            BoundKind::Lambda {
-                params,
-                upvalues,
-                body,
-                positional_count,
-            } => {
-                let mut next_upvalues = Vec::new();
-                for addr in upvalues {
-                    next_upvalues.push(self.reindex_addr(addr, mapping));
-                }
-                (
-                    BoundKind::Lambda {
-                        params,
-                        upvalues: next_upvalues,
-                        body,
-                        positional_count,
-                    },
-                    node.ty.clone(),
-                )
-            }
-            BoundKind::If {
-                cond,
-                then_br,
-                else_br,
-            } => {
-                let cond = Box::new(self.reindex_upvalues(*cond, mapping));
-                let then_br = Box::new(self.reindex_upvalues(*then_br, mapping));
-                let else_br = else_br.map(|e| Box::new(self.reindex_upvalues(*e, mapping)));
-                (
-                    BoundKind::If {
-                        cond,
-                        then_br,
-                        else_br,
-                    },
-                    node.ty.clone(),
-                )
-            }
-            BoundKind::Block { exprs } => {
-                let exprs = exprs
-                    .into_iter()
-                    .map(|e| self.reindex_upvalues(e, mapping))
-                    .collect();
-                (BoundKind::Block { exprs }, node.ty.clone())
-            }
-            BoundKind::Call { callee, args } => {
-                let callee = Box::new(self.reindex_upvalues(*callee, mapping));
-                let args = Box::new(self.reindex_upvalues(*args, mapping));
-                (BoundKind::Call { callee, args }, node.ty.clone())
-            }
-            BoundKind::Define {
-                name,
-                addr,
-                kind,
-                value,
-                captured_by,
-            } => {
-                let value = Box::new(self.reindex_upvalues(*value, mapping));
-                (
-                    BoundKind::Define {
-                        name,
-                        addr,
-                        kind,
-                        value,
-                        captured_by,
-                    },
-                    node.ty.clone(),
-                )
-            }
-            BoundKind::Set { addr, value } => {
-                let value = Box::new(self.reindex_upvalues(*value, mapping));
-                (BoundKind::Set { addr, value }, node.ty.clone())
-            }
-            BoundKind::Tuple { elements } => {
-                let elements = elements
-                    .into_iter()
-                    .map(|e| self.reindex_upvalues(e, mapping))
-                    .collect();
-                (BoundKind::Tuple { elements }, node.ty.clone())
-            }
-            BoundKind::Record { fields } => {
-                let fields = fields
-                    .into_iter()
-                    .map(|(k, v)| {
-                        (
-                            self.reindex_upvalues(k, mapping),
-                            self.reindex_upvalues(v, mapping),
-                        )
-                    })
-                    .collect();
-                (BoundKind::Record { fields }, node.ty.clone())
-            }
-            k => (k, node.ty.clone()),
-        };
-        Node {
-            identity: node.identity.clone(),
-            kind: new_kind,
-            ty: metrics,
-        }
-    }
-}

src/ast/compiler/optimizer/substitution_map.rs

@@ -0,0 +1,321 @@
+use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, BoundKind, LocalSlot, UpvalueIdx};
+use crate::ast::nodes::Node;
+use crate::ast::types::{Identity, Value};
+use crate::ast::vm::Closure;
+use std::collections::{HashMap, HashSet};
+
+#[derive(Default)]
+pub struct UsageInfo {
+    pub used: HashSet<Address>,
+    pub assigned: HashSet<Address>,
+    pub used_identities: HashSet<Identity>,
+}
+
+impl UsageInfo {
+    pub fn is_assigned(&self, addr: &Address) -> bool {
+        self.assigned.contains(addr)
+    }
+
+    pub fn is_used(&self, addr: &Address) -> bool {
+        self.used.contains(addr)
+    }
+
+    pub fn collect(&mut self, node: &AnalyzedNode) {
+        match &node.kind {
+            BoundKind::Destructure { pattern, value } => {
+                self.collect_pattern(pattern);
+                self.collect(value);
+            }
+            BoundKind::Constant(v) => {
+                if let Value::Object(obj) = v
+                    && let Some(closure) = obj.as_any().downcast_ref::<Closure>()
+                {
+                    self.used_identities
+                        .insert(closure.function_node.identity.clone());
+                    self.collect(&closure.function_node);
+                }
+            }
+            BoundKind::Get { addr, .. } => {
+                self.used.insert(*addr);
+            }
+            BoundKind::Set { addr, value } => {
+                self.assigned.insert(*addr);
+                self.collect(value);
+            }
+            BoundKind::Lambda {
+                params,
+                body,
+                upvalues,
+                ..
+            } => {
+                self.used_identities.insert(node.identity.clone());
+
+                let mut inner_info = UsageInfo::default();
+                inner_info.collect(params);
+                inner_info.collect(body);
+
+                // Propagate globals and identities
+                for addr in &inner_info.used {
+                    if let Address::Global(_) = addr {
+                        self.used.insert(*addr);
+                    }
+                }
+                for addr in &inner_info.assigned {
+                    if let Address::Global(_) = addr {
+                        self.assigned.insert(*addr);
+                    }
+                }
+                self.used_identities.extend(inner_info.used_identities);
+
+                // Map used upvalues to parent scope
+                for addr in &inner_info.used {
+                    if let Address::Upvalue(idx) = addr
+                        && let Some(parent_addr) = upvalues.get(idx.0 as usize)
+                    {
+                        self.used.insert(*parent_addr);
+                    }
+                }
+                // Map assigned upvalues to parent scope
+                for addr in &inner_info.assigned {
+                    if let Address::Upvalue(idx) = addr
+                        && let Some(parent_addr) = upvalues.get(idx.0 as usize)
+                    {
+                        self.assigned.insert(*parent_addr);
+                    }
+                }
+            }
+            BoundKind::Block { exprs } => {
+                for e in exprs {
+                    self.collect(e);
+                }
+            }
+            BoundKind::If {
+                cond,
+                then_br,
+                else_br,
+            } => {
+                self.collect(cond);
+                self.collect(then_br);
+                if let Some(e) = else_br {
+                    self.collect(e);
+                }
+            }
+            BoundKind::Call { callee, args } => {
+                self.collect(callee);
+                self.collect(args);
+            }
+            BoundKind::Tuple { elements } => {
+                for e in elements {
+                    self.collect(e);
+                }
+            }
+            BoundKind::Record { fields } => {
+                for (k, v) in fields {
+                    self.collect(k);
+                    self.collect(v);
+                }
+            }
+            BoundKind::Define { value, .. } => {
+                self.collect(value);
+            }
+            BoundKind::Expansion { bound_expanded, .. } => {
+                self.collect(bound_expanded);
+            }
+            _ => {}
+        }
+    }
+
+    pub fn collect_pattern(&mut self, node: &AnalyzedNode) {
+        match &node.kind {
+            BoundKind::Define { addr, .. } => {
+                self.assigned.insert(*addr);
+            }
+            BoundKind::Set { addr, .. } => {
+                self.assigned.insert(*addr);
+            }
+            BoundKind::Tuple { elements } => {
+                for el in elements {
+                    self.collect_pattern(el);
+                }
+            }
+            _ => {}
+        }
+    }
+}
+
+pub struct SubstitutionMap {
+    pub values: HashMap<Address, Value>,
+    pub ast_substitutions: HashMap<Address, AnalyzedNode>,
+    pub slot_mapping: HashMap<LocalSlot, LocalSlot>,
+    pub assigned: HashSet<Address>,
+    pub next_slot: u32,
+    pub used: HashSet<Address>,
+    pub captured_slots: HashSet<LocalSlot>,
+}
+
+impl SubstitutionMap {
+    pub fn new() -> Self {
+        Self {
+            values: HashMap::new(),
+            ast_substitutions: HashMap::new(),
+            slot_mapping: HashMap::new(),
+            assigned: HashSet::new(),
+            next_slot: 0,
+            used: HashSet::new(),
+            captured_slots: HashSet::new(),
+        }
+    }
+
+    pub fn add_ast_substitution(&mut self, addr: Address, node: AnalyzedNode) {
+        self.ast_substitutions.insert(addr, node);
+    }
+
+    pub fn map_slot(&mut self, old_slot: LocalSlot) -> LocalSlot {
+        if let Some(&new_slot) = self.slot_mapping.get(&old_slot) {
+            return new_slot;
+        }
+        let new_slot = LocalSlot(self.next_slot);
+        self.slot_mapping.insert(old_slot, new_slot);
+        self.next_slot += 1;
+        new_slot
+    }
+
+    pub fn map_address(&mut self, addr: Address) -> Address {
+        match addr {
+            Address::Local(slot) => Address::Local(self.map_slot(slot)),
+            other => other,
+        }
+    }
+
+    pub fn add_value(&mut self, addr: Address, val: Value) {
+        self.values.insert(addr, val);
+    }
+
+    pub fn get_value(&self, addr: &Address) -> Option<&Value> {
+        self.values.get(addr)
+    }
+
+    pub fn remove_value(&mut self, addr: &Address) {
+        self.values.remove(addr);
+    }
+
+    fn reindex_addr(&self, addr: Address, mapping: &[Option<u32>]) -> Address {
+        if let Address::Upvalue(idx) = addr
+            && let Some(res) = mapping.get(idx.0 as usize)
+            && let Some(new_idx) = res
+        {
+            Address::Upvalue(UpvalueIdx(*new_idx))
+        } else {
+            addr
+        }
+    }
+
+    pub fn reindex_upvalues(&self, node: AnalyzedNode, mapping: &[Option<u32>]) -> AnalyzedNode {
+        let (new_kind, metrics) = match node.kind {
+            BoundKind::Get { addr, name } => (
+                BoundKind::Get {
+                    addr: self.reindex_addr(addr, mapping),
+                    name,
+                },
+                node.ty.clone(),
+            ),
+            BoundKind::Lambda {
+                params,
+                upvalues,
+                body,
+                positional_count,
+            } => {
+                let mut next_upvalues = Vec::new();
+                for addr in upvalues {
+                    next_upvalues.push(self.reindex_addr(addr, mapping));
+                }
+                (
+                    BoundKind::Lambda {
+                        params,
+                        upvalues: next_upvalues,
+                        body,
+                        positional_count,
+                    },
+                    node.ty.clone(),
+                )
+            }
+            BoundKind::If {
+                cond,
+                then_br,
+                else_br,
+            } => {
+                let cond = Box::new(self.reindex_upvalues(*cond, mapping));
+                let then_br = Box::new(self.reindex_upvalues(*then_br, mapping));
+                let else_br = else_br.map(|e| Box::new(self.reindex_upvalues(*e, mapping)));
+                (
+                    BoundKind::If {
+                        cond,
+                        then_br,
+                        else_br,
+                    },
+                    node.ty.clone(),
+                )
+            }
+            BoundKind::Block { exprs } => {
+                let exprs = exprs
+                    .into_iter()
+                    .map(|e| self.reindex_upvalues(e, mapping))
+                    .collect();
+                (BoundKind::Block { exprs }, node.ty.clone())
+            }
+            BoundKind::Call { callee, args } => {
+                let callee = Box::new(self.reindex_upvalues(*callee, mapping));
+                let args = Box::new(self.reindex_upvalues(*args, mapping));
+                (BoundKind::Call { callee, args }, node.ty.clone())
+            }
+            BoundKind::Define {
+                name,
+                addr,
+                kind,
+                value,
+                captured_by,
+            } => {
+                let value = Box::new(self.reindex_upvalues(*value, mapping));
+                (
+                    BoundKind::Define {
+                        name,
+                        addr,
+                        kind,
+                        value,
+                        captured_by,
+                    },
+                    node.ty.clone(),
+                )
+            }
+            BoundKind::Set { addr, value } => {
+                let value = Box::new(self.reindex_upvalues(*value, mapping));
+                (BoundKind::Set { addr, value }, node.ty.clone())
+            }
+            BoundKind::Tuple { elements } => {
+                let elements = elements
+                    .into_iter()
+                    .map(|e| self.reindex_upvalues(e, mapping))
+                    .collect();
+                (BoundKind::Tuple { elements }, node.ty.clone())
+            }
+            BoundKind::Record { fields } => {
+                let fields = fields
+                    .into_iter()
+                    .map(|(k, v)| {
+                        (
+                            self.reindex_upvalues(k, mapping),
+                            self.reindex_upvalues(v, mapping),
+                        )
+                    })
+                    .collect();
+                (BoundKind::Record { fields }, node.ty.clone())
+            }
+            k => (k, node.ty.clone()),
+        };
+        Node {
+            identity: node.identity.clone(),
+            kind: new_kind,
+            ty: metrics,
+        }
+    }
+}