diff --git a/src/ast/compiler/optimizer/mod.rs b/src/ast/compiler/optimizer/mod.rs new file mode 100644 index 0000000..78f47eb --- /dev/null +++ b/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}; diff --git a/src/ast/compiler/optimizer.rs b/src/ast/compiler/optimizer/optimizer.rs similarity index 73% rename from src/ast/compiler/optimizer.rs rename to src/ast/compiler/optimizer/optimizer.rs index c890a03..ed082d3 100644 --- a/src/ast/compiler/optimizer.rs +++ b/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
, - assigned: HashSet
, - used_identities: HashSet, -} - -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 { 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::() - { - 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, - ast_substitutions: HashMap, - slot_mapping: HashMap, - assigned: HashSet
, - next_slot: u32, - used: HashSet
, - captured_slots: HashSet, -} - -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]) -> 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]) -> 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, - } - } -} diff --git a/src/ast/compiler/optimizer/substitution_map.rs b/src/ast/compiler/optimizer/substitution_map.rs new file mode 100644 index 0000000..62ff71f --- /dev/null +++ b/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
, + pub assigned: HashSet
, + pub used_identities: HashSet, +} + +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::() + { + 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, + pub ast_substitutions: HashMap, + pub slot_mapping: HashMap, + pub assigned: HashSet
, + pub next_slot: u32, + pub used: HashSet
, + pub captured_slots: HashSet, +} + +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]) -> 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]) -> 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, + } + } +}