Refactor Optimizer to use Folder

Extract common AST manipulation and folding logic into a new Folder
struct. This improves code organization and reusability. The Optimizer
now delegates these tasks to the Folder.

src/ast/compiler/optimizer/engine.rs

@@ -1,14 +1,15 @@
 use crate::ast::compiler::bound_nodes::{
-    Address, AnalyzedNode, BoundKind, GlobalIdx, LocalSlot, NodeMetrics, UpvalueIdx,
+    Address, AnalyzedNode, BoundKind, GlobalIdx, LocalSlot, UpvalueIdx,
 };
 use crate::ast::nodes::Node;
-use crate::ast::types::{Purity, StaticType, Value};
+use crate::ast::types::{Purity, Value};
 use crate::ast::vm::Closure;
 use std::cell::RefCell;
 use std::collections::{HashMap, HashSet};
 use std::rc::Rc;
 
 use super::substitution_map::{SubstitutionMap, UsageInfo};
+use super::folder::Folder;
 
 pub struct Optimizer {
     pub enabled: bool,
@@ -99,6 +100,7 @@ impl Optimizer {
         sub: &mut SubstitutionMap,
         path: &mut PathTracker,
     ) -> AnalyzedNode {
+        let folder = Folder::new(&self.globals);
         let (new_kind, metrics) = match node.kind {
             BoundKind::Get { addr, ref name } => {
                 if !sub.assigned.contains(&addr) {
@@ -106,7 +108,7 @@ impl Optimizer {
                     if let Some(val) = sub.get_value(&addr)
                         && self.is_inlinable_value(val, addr)
                     {
-                        return self.make_constant_node(val.clone(), &node);
+                        return folder.make_constant_node(val.clone(), &node);
                     }
 
                     // 2. Try inlining from AST substitution map (pure expressions)
@@ -122,7 +124,7 @@ impl Optimizer {
                         if let Some(val) = globals.get(idx.0 as usize)
                             && self.is_inlinable_value(val, addr)
                         {
-                            return self.make_constant_node(val.clone(), &node);
+                            return folder.make_constant_node(val.clone(), &node);
                         }
                     }
                 }
@@ -305,7 +307,7 @@ impl Optimizer {
                         }
                     }
 
-                    if let Some(folded) = self.try_fold_pure(&callee, &args) {
+                    if let Some(folded) = folder.try_fold_pure(&callee, &args) {
                         return folded;
                     }
                 }
@@ -343,7 +345,7 @@ impl Optimizer {
                     } else if let Some(else_node) = else_br {
                         return self.visit_node((**else_node).clone(), sub, path);
                     } else {
-                        return self.make_nop_node(&node);
+                        return folder.make_nop_node(&node);
                     }
                 }
 
@@ -414,7 +416,7 @@ impl Optimizer {
 
                 if self.enabled {
                     if new_exprs.is_empty() {
-                        return self.make_nop_node(&node);
+                        return folder.make_nop_node(&node);
                     } else if new_exprs.len() == 1 {
                         return new_exprs.pop().unwrap();
                     }
@@ -606,71 +608,6 @@ impl Optimizer {
         true // Locals/Upvalues that reach here and passed the type check are inlinable
     }
 
-    fn make_constant_node(&self, val: Value, template: &AnalyzedNode) -> AnalyzedNode {
-        let ty = val.static_type();
-        let typed_original = Rc::new(Node {
-            identity: template.identity.clone(),
-            kind: BoundKind::Constant(val.clone()),
-            ty: ty.clone(),
-        });
-        Node {
-            identity: template.identity.clone(),
-            kind: BoundKind::Constant(val),
-            ty: NodeMetrics {
-                original: typed_original,
-                purity: Purity::Pure,
-                is_recursive: false,
-            },
-        }
-    }
-
-    fn make_nop_node(&self, template: &AnalyzedNode) -> AnalyzedNode {
-        let typed_original = Rc::new(Node {
-            identity: template.identity.clone(),
-            kind: BoundKind::Nop,
-            ty: StaticType::Void,
-        });
-        Node {
-            identity: template.identity.clone(),
-            kind: BoundKind::Nop,
-            ty: NodeMetrics {
-                original: typed_original,
-                purity: Purity::Pure,
-                is_recursive: false,
-            },
-        }
-    }
-
-    fn try_fold_pure(&self, callee: &AnalyzedNode, args: &AnalyzedNode) -> Option<AnalyzedNode> {
-        if callee.ty.purity < Purity::Pure || args.ty.purity < Purity::Pure {
-            return None;
-        }
-
-        let mut arg_nodes = Vec::new();
-        self.flatten_tuple(args.clone(), &mut arg_nodes);
-        let mut arg_values = Vec::with_capacity(arg_nodes.len());
-        for node in arg_nodes {
-            if let BoundKind::Constant(val) = node.kind {
-                arg_values.push(val);
-            } else {
-                return None;
-            }
-        }
-        let func_val = match &callee.kind {
-            BoundKind::Get {
-                addr: Address::Global(idx),
-                ..
-            } => self.globals.as_ref()?.borrow().get(idx.0 as usize)?.clone(),
-            BoundKind::Constant(val) => val.clone(),
-            _ => return None,
-        };
-        let result = match func_val {
-            Value::Function(f) => (f.func)(arg_values),
-            _ => return None,
-        };
-        Some(self.make_constant_node(result, callee))
-    }
-
     fn try_beta_reduce_with_sub(
         &self,
         params: &AnalyzedNode,

src/ast/compiler/optimizer/folder.rs

@@ -0,0 +1,97 @@
+use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, BoundKind, NodeMetrics};
+use crate::ast::nodes::Node;
+use crate::ast::types::{Purity, StaticType, Value};
+use std::cell::RefCell;
+use std::rc::Rc;
+
+pub struct Folder<'a> {
+    pub globals: &'a Option<Rc<RefCell<Vec<Value>>>>,
+}
+
+impl<'a> Folder<'a> {
+    pub fn new(globals: &'a Option<Rc<RefCell<Vec<Value>>>>) -> Self {
+        Self { globals }
+    }
+
+    pub fn make_constant_node(&self, val: Value, template: &AnalyzedNode) -> AnalyzedNode {
+        let ty = val.static_type();
+        let typed_original = Rc::new(Node {
+            identity: template.identity.clone(),
+            kind: BoundKind::Constant(val.clone()),
+            ty: ty.clone(),
+        });
+        Node {
+            identity: template.identity.clone(),
+            kind: BoundKind::Constant(val),
+            ty: NodeMetrics {
+                original: typed_original,
+                purity: Purity::Pure,
+                is_recursive: false,
+            },
+        }
+    }
+
+    pub fn make_nop_node(&self, template: &AnalyzedNode) -> AnalyzedNode {
+        let typed_original = Rc::new(Node {
+            identity: template.identity.clone(),
+            kind: BoundKind::Nop,
+            ty: StaticType::Void,
+        });
+        Node {
+            identity: template.identity.clone(),
+            kind: BoundKind::Nop,
+            ty: NodeMetrics {
+                original: typed_original,
+                purity: Purity::Pure,
+                is_recursive: false,
+            },
+        }
+    }
+
+    pub fn try_fold_pure(&self, callee: &AnalyzedNode, args: &AnalyzedNode) -> Option<AnalyzedNode> {
+        if callee.ty.purity < Purity::Pure || args.ty.purity < Purity::Pure {
+            return None;
+        }
+
+        let mut arg_nodes = Vec::new();
+        self.flatten_tuple(args.clone(), &mut arg_nodes);
+        let mut arg_values = Vec::with_capacity(arg_nodes.len());
+        for node in arg_nodes {
+            if let BoundKind::Constant(val) = node.kind {
+                arg_values.push(val);
+            } else {
+                return None;
+            }
+        }
+        let func_val = match &callee.kind {
+            BoundKind::Get {
+                addr: Address::Global(idx),
+                ..
+            } => self.globals.as_ref()?.borrow().get(idx.0 as usize)?.clone(),
+            BoundKind::Constant(val) => val.clone(),
+            _ => return None,
+        };
+        let result = match func_val {
+            Value::Function(f) => (f.func)(arg_values),
+            _ => return None,
+        };
+        Some(self.make_constant_node(result, callee))
+    }
+
+    fn flatten_tuple(&self, node: AnalyzedNode, into: &mut Vec<AnalyzedNode>) {
+        match node.kind {
+            BoundKind::Tuple { elements } => {
+                for el in elements {
+                    self.flatten_tuple(el, into);
+                }
+            }
+            BoundKind::Record { fields } => {
+                for (_, v) in fields {
+                    self.flatten_tuple(v, into);
+                }
+            }
+            BoundKind::Nop => {}
+            _ => into.push(node),
+        }
+    }
+}

src/ast/compiler/optimizer/mod.rs

@@ -1,5 +1,7 @@
-pub mod optimizer;
+pub mod engine;
+pub mod folder;
 pub mod substitution_map;
 
-pub use optimizer::Optimizer;
+pub use engine::Optimizer;
+pub use folder::Folder;
 pub use substitution_map::{SubstitutionMap, UsageInfo};

src/ast/compiler/optimizer/substitution_map.rs

@@ -143,6 +143,7 @@ impl UsageInfo {
     }
 }
 
+#[derive(Default)]
 pub struct SubstitutionMap {
     pub values: HashMap<Address, Value>,
     pub ast_substitutions: HashMap<Address, AnalyzedNode>,
@@ -155,15 +156,7 @@ pub struct SubstitutionMap {
 
 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(),
-        }
+        Self::default()
     }
 
     pub fn add_ast_substitution(&mut self, addr: Address, node: AnalyzedNode) {