From 01072640266b208867e8f2a683bbcab05d680a7f Mon Sep 17 00:00:00 2001 From: Michael Schimmel Date: Sun, 22 Feb 2026 00:16:04 +0100 Subject: [PATCH] feat: Implement purity inference for optimization This commit introduces purity inference to the compiler's optimizer. This allows for more aggressive constant folding and dead code elimination by tracking which functions and global variables are free of side effects. Key changes include: - Added `global_purity` field to `Optimizer` and `Environment`. - Modified `Optimizer::is_pure` to recursively determine if an AST node represents a pure computation. - Introduced `Optimizer::try_fold_pure` to replace the old `try_fold_intrinsic`, enabling folding of pure function calls with constant arguments. - Updated `Environment::register_native` and `Environment::register_constant` to optionally record purity. - Added purity flags to several built-in functions in `core.rs` and `datetime.rs`. - A new example `optimizer_purity.myc` demonstrates the new feature. --- examples/optimizer_purity.myc | 5 + src/ast/compiler/optimizer.rs | 197 +++++++++++++++++++++++++--------- src/ast/environment.rs | 16 ++- src/ast/rtl/core.rs | 36 +++---- src/ast/rtl/datetime.rs | 2 +- src/main.rs | 1 + 6 files changed, 184 insertions(+), 73 deletions(-) create mode 100644 examples/optimizer_purity.myc diff --git a/examples/optimizer_purity.myc b/examples/optimizer_purity.myc new file mode 100644 index 0000000..288882e --- /dev/null +++ b/examples/optimizer_purity.myc @@ -0,0 +1,5 @@ +(do + (def PI 3.1415) + (def area (fn [x] (* PI x))) + (area 10) +) \ No newline at end of file diff --git a/src/ast/compiler/optimizer.rs b/src/ast/compiler/optimizer.rs index d3500ec..9cf7d4e 100644 --- a/src/ast/compiler/optimizer.rs +++ b/src/ast/compiler/optimizer.rs @@ -12,11 +12,12 @@ pub struct Optimizer { pub level: u32, max_passes: usize, pub globals: Option>>>, + pub global_purity: Option>>>, } impl Optimizer { pub fn new(level: u32) -> Self { - Self { level, max_passes: 5, globals: None } + Self { level, max_passes: 5, globals: None, global_purity: None } } pub fn with_globals(mut self, globals: Rc>>) -> Self { @@ -24,6 +25,11 @@ impl Optimizer { self } + pub fn with_purity(mut self, purity: Rc>>) -> Self { + self.global_purity = Some(purity); + self + } + pub fn optimize(&self, node: TypedNode) -> TypedNode { if self.level == 0 { return node; @@ -119,7 +125,7 @@ impl Optimizer { } } - if let Some(folded) = self.try_fold_intrinsic(&callee, &args) { + if let Some(folded) = self.try_fold_pure(&callee, &args) { return folded; } } @@ -193,14 +199,14 @@ impl Optimizer { let removable = match &e.kind { BoundKind::DefLocal { slot, .. } | BoundKind::Set { addr: Address::Local(slot), .. } => { - !sub.used_slots.contains(slot) && !sub.captured_slots.contains(slot) && self.is_side_effect_free(&e) + !sub.used_slots.contains(slot) && !sub.captured_slots.contains(slot) && self.is_pure(&e) } BoundKind::DefGlobal { global_index, .. } => { // Phase 4: DCE for Globals // A global can be removed if it was only used for inlining in this script. - !sub.used_globals.contains(global_index) && self.is_side_effect_free(&e) + !sub.used_globals.contains(global_index) && self.is_pure(&e) } - _ => self.is_side_effect_free(&e), + _ => self.is_pure(&e), }; if !removable { @@ -283,6 +289,14 @@ impl Optimizer { if let BoundKind::Constant(val) = &value.kind { sub.add_global(global_index, val.clone()); } + + // Purity Inference: Globals are pure if their value is pure + if self.is_pure(&value) { + if let Some(purity_rc) = &self.global_purity { + purity_rc.borrow_mut().insert(global_index, true); + } + } + (BoundKind::DefGlobal { name, global_index, value }, node.ty) }, BoundKind::Set { addr, value } => { @@ -339,21 +353,118 @@ impl Optimizer { } } - fn is_side_effect_free(&self, node: &TypedNode) -> bool { + fn is_pure(&self, node: &TypedNode) -> bool { match &node.kind { - BoundKind::Constant(_) | BoundKind::Get { .. } | BoundKind::Parameter { .. } | BoundKind::Nop | BoundKind::Lambda { .. } => true, - BoundKind::Tuple { elements } => elements.iter().all(|e| self.is_side_effect_free(e)), - BoundKind::Record { fields } => fields.iter().all(|(k, v)| self.is_side_effect_free(k) && self.is_side_effect_free(v)), - BoundKind::If { cond, then_br, else_br } => { - self.is_side_effect_free(cond) && self.is_side_effect_free(then_br) && else_br.as_ref().is_none_or(|e| self.is_side_effect_free(e)) + BoundKind::Constant(_) | BoundKind::Parameter { .. } | BoundKind::Nop | BoundKind::Lambda { .. } => true, + BoundKind::Get { addr, .. } => { + match addr { + Address::Global(idx) => { + if let Some(purity_rc) = &self.global_purity { + *purity_rc.borrow().get(idx).unwrap_or(&false) + } else { + false + } + } + _ => true, // Locals and Upvalues are pure in terms of side-effects + } } - BoundKind::Block { exprs } => exprs.iter().all(|e| self.is_side_effect_free(e)), - BoundKind::Expansion { bound_expanded, .. } => self.is_side_effect_free(bound_expanded), - BoundKind::DefLocal { value, .. } | BoundKind::Set { value, .. } => self.is_side_effect_free(value), - _ => false, // Call and DefGlobal are considered impure + BoundKind::Tuple { elements } => elements.iter().all(|e| self.is_pure(e)), + BoundKind::Record { fields } => fields.iter().all(|(k, v)| self.is_pure(k) && self.is_pure(v)), + BoundKind::If { cond, then_br, else_br } => { + self.is_pure(cond) && self.is_pure(then_br) && else_br.as_ref().is_none_or(|e| self.is_pure(e)) + } + BoundKind::Block { exprs } => exprs.iter().all(|e| self.is_pure(e)), + BoundKind::Expansion { bound_expanded, .. } => self.is_pure(bound_expanded), + BoundKind::DefLocal { value, .. } | BoundKind::Set { addr: Address::Local(_), value } => self.is_pure(value), + + BoundKind::Call { callee, args } => { + // A call is pure if the callee is a known pure function and args are pure + let callee_pure = match &callee.kind { + BoundKind::Get { addr: Address::Global(idx), .. } => { + if let Some(purity_rc) = &self.global_purity { + *purity_rc.borrow().get(idx).unwrap_or(&false) + } else { + false + } + } + BoundKind::Constant(Value::Function(_)) => true, // Native fns are checked at registry + BoundKind::Constant(Value::Object(obj)) => { + if let Some(closure) = obj.as_any().downcast_ref::() { + self.is_pure(&closure.function_node) + } else { + false + } + } + _ => false, + }; + callee_pure && self.is_pure(args) + } + + _ => false, // DefGlobal and Set(Global) are impure } } + fn try_fold_pure(&self, callee: &TypedNode, args: &TypedNode) -> Option { + // 1. Check if the whole call is pure + // Create a temporary call node to check purity + let temp_call = Node { + identity: callee.identity.clone(), + ty: StaticType::Any, + kind: BoundKind::Call { callee: Box::new(callee.clone()), args: Box::new(args.clone()) } + }; + + if !self.is_pure(&temp_call) { + return None; + } + + // 2. Extract constant arguments + let mut arg_nodes = Vec::new(); + self.flatten_typed_tuple(args, &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; // All arguments must be constants for folding + } + } + + // 3. Resolve the function to execute + let func_val = match &callee.kind { + BoundKind::Get { addr: Address::Global(idx), .. } => { + self.globals.as_ref()?.borrow().get(*idx as usize)?.clone() + } + BoundKind::Constant(val) => val.clone(), + _ => return None, + }; + + // 4. Execute + let result = match func_val { + Value::Function(f) => f(arg_values), + Value::Object(obj) => { + if let Some(_closure) = obj.as_any().downcast_ref::() { + // For closures, we'd need a VM to execute. + // However, if it's pure and we are in the Optimizer, + // we might want to avoid full VM spin-up here if possible, + // or use a shared VM instance. + // For now, we only fold Native Functions. + return None; + } else { + return None; + } + } + _ => return None, + }; + + // 5. Return new constant node + Some(Node { + identity: callee.identity.clone(), + ty: result.static_type(), + kind: BoundKind::Constant(result), + }) + } + fn try_beta_reduce(&self, params: &TypedNode, args: &TypedNode, body: TypedNode) -> Option { if self.contains_def_local(&body) { return None; @@ -405,42 +516,6 @@ impl Optimizer { } } - fn try_fold_intrinsic(&self, callee: &TypedNode, args: &TypedNode) -> Option { - if let BoundKind::Get { name, .. } = &callee.kind - && let BoundKind::Tuple { elements } = &args.kind - && elements.len() == 2 - { - let val_a = self.get_const_int(&elements[0]); - let val_b = self.get_const_int(&elements[1]); - - if let (Some(a), Some(b)) = (val_a, val_b) { - let res = match &*name.name { - "+" => Some(Value::Int(a + b)), - "-" => Some(Value::Int(a - b)), - "*" => Some(Value::Int(a * b)), - "/" if b != 0 => Some(Value::Int(a / b)), - _ => None - }; - if let Some(val) = res { - return Some(Node { - identity: callee.identity.clone(), - ty: StaticType::Int, - kind: BoundKind::Constant(val), - }); - } - } - } - None - } - - fn get_const_int(&self, node: &TypedNode) -> Option { - match &node.kind { - BoundKind::Constant(Value::Int(i)) => Some(*i), - BoundKind::Tuple { elements } if elements.len() == 1 => self.get_const_int(&elements[0]), - _ => None - } - } - fn contains_def_local(&self, node: &TypedNode) -> bool { match &node.kind { BoundKind::DefLocal { .. } => true, @@ -571,7 +646,6 @@ impl SubstitutionMap { #[cfg(test)] mod tests { - use super::*; use crate::ast::environment::Environment; fn get_optimized_dump(source: &str, level: u32) -> String { @@ -586,6 +660,25 @@ mod tests { insta::assert_snapshot!(dump); } + #[test] + fn test_opt_folding_complex() { + // Test float folding + let dump_float = get_optimized_dump("(+ 1.5 2.5)", 2); + assert!(dump_float.contains("Constant: 4")); + + // Test text folding + let dump_text = get_optimized_dump("(+ \"hello \" \"world\")", 2); + assert!(dump_text.contains("Constant: \"hello world\"")); + + // Test date folding + let dump_date = get_optimized_dump("(date \"2023-01-01\")", 2); + assert!(dump_date.contains("Constant: #2023-01-01 00:00:00#")); + + // Test comparison folding + let dump_cmp = get_optimized_dump("(> 10 5)", 2); + assert!(dump_cmp.contains("Constant: true")); + } + #[test] fn test_opt_local_inlining() { let source = "(fn [] (do (def x 10) (+ x 5)))"; diff --git a/src/ast/environment.rs b/src/ast/environment.rs index e4558ce..53b7fa7 100644 --- a/src/ast/environment.rs +++ b/src/ast/environment.rs @@ -21,6 +21,7 @@ use crate::ast::rtl::intrinsics; pub struct Environment { pub global_names: Rc>>, pub global_types: Rc>>, + pub global_purity: Rc>>, pub global_values: Rc>>, pub function_registry: Rc>>, pub monomorph_cache: Rc>, @@ -85,6 +86,7 @@ impl Environment { let env = Self { global_names: Rc::new(RefCell::new(HashMap::new())), global_types: Rc::new(RefCell::new(HashMap::new())), + global_purity: Rc::new(RefCell::new(HashMap::new())), global_values: Rc::new(RefCell::new(Vec::new())), function_registry: Rc::new(RefCell::new(HashMap::new())), monomorph_cache: Rc::new(RefCell::new(HashMap::new())), @@ -112,15 +114,18 @@ impl Environment { &self, name: &str, ty: StaticType, + is_pure: bool, func: impl Fn(Vec) -> Value + 'static, ) { let mut names = self.global_names.borrow_mut(); let mut types = self.global_types.borrow_mut(); let mut values = self.global_values.borrow_mut(); + let mut purity = self.global_purity.borrow_mut(); let idx = values.len() as u32; names.insert(Symbol::from(name), idx); types.insert(idx, ty); + purity.insert(idx, is_pure); values.push(Value::Function(Rc::new(func))); } @@ -128,10 +133,12 @@ impl Environment { let mut names = self.global_names.borrow_mut(); let mut types = self.global_types.borrow_mut(); let mut values = self.global_values.borrow_mut(); + let mut purity = self.global_purity.borrow_mut(); let idx = values.len() as u32; names.insert(Symbol::from(name), idx); types.insert(idx, ty); + purity.insert(idx, true); // Constants are always pure values.push(val); } @@ -182,7 +189,9 @@ impl Environment { let specialized = self.specialize_node(node); // 2. Optimize (Level 1: Cracking, Level 2: Collapsing) - let optimizer = Optimizer::new(self.optimization_level).with_globals(self.global_values.clone()); + let optimizer = Optimizer::new(self.optimization_level) + .with_globals(self.global_values.clone()) + .with_purity(self.global_purity.clone()); let optimized = optimizer.optimize(specialized); // 3. TCO (Always performed, converts to ExecNode) @@ -200,6 +209,7 @@ impl Environment { let mono_cache = self.monomorph_cache.clone(); let global_values = self.global_values.clone(); let global_types = self.global_types.clone(); + let global_purity = self.global_purity.clone(); let opt_level = self.optimization_level; let compiler = Rc::new( @@ -227,7 +237,9 @@ impl Environment { let specialized_ast = sub_specializer.specialize(retyped_ast); // 3. Optimize (Phase 2: Cracking & Folding) - let optimizer = Optimizer::new(opt_level).with_globals(global_values.clone()); + let optimizer = Optimizer::new(opt_level) + .with_globals(global_values.clone()) + .with_purity(global_purity.clone()); let optimized_ast = optimizer.optimize(specialized_ast); // 4. TCO (converts to ExecNode) diff --git a/src/ast/rtl/core.rs b/src/ast/rtl/core.rs index 53b496b..3c6fff3 100644 --- a/src/ast/rtl/core.rs +++ b/src/ast/rtl/core.rs @@ -28,7 +28,7 @@ fn register_arithmetic(env: &Environment) { Signature { params: StaticType::Tuple(vec![StaticType::Text, StaticType::Text]), ret: StaticType::Text }, Signature { params: StaticType::Tuple(vec![StaticType::DateTime, StaticType::Int]), ret: StaticType::DateTime }, ]); - env.register_native("+", add_ty, |args| { + env.register_native("+", add_ty, true, |args| { if args.len() == 2 { match (&args[0], &args[1]) { (Value::Int(a), Value::Int(b)) => Value::Int(a + b), @@ -63,7 +63,7 @@ fn register_arithmetic(env: &Environment) { Signature { params: StaticType::Tuple(vec![StaticType::DateTime, StaticType::DateTime]), ret: StaticType::Int }, Signature { params: StaticType::Tuple(vec![StaticType::DateTime, StaticType::Int]), ret: StaticType::DateTime }, ]); - env.register_native("-", sub_ty, |args| { + env.register_native("-", sub_ty, true, |args| { if args.is_empty() { return Value::Void; } if args.len() == 1 { return match args[0] { @@ -101,7 +101,7 @@ fn register_arithmetic(env: &Environment) { Signature { params: StaticType::Tuple(vec![StaticType::Int, StaticType::Int]), ret: StaticType::Int }, Signature { params: StaticType::Tuple(vec![StaticType::Float, StaticType::Float]), ret: StaticType::Float }, ]); - env.register_native("*", mul_ty, |args| { + env.register_native("*", mul_ty, true, |args| { if args.len() == 2 { match (&args[0], &args[1]) { (Value::Int(a), Value::Int(b)) => Value::Int(a * b), @@ -125,7 +125,7 @@ fn register_arithmetic(env: &Environment) { Signature { params: StaticType::Tuple(vec![StaticType::Int, StaticType::Int]), ret: StaticType::Float }, Signature { params: StaticType::Tuple(vec![StaticType::Float, StaticType::Float]), ret: StaticType::Float }, ]); - env.register_native("/", div_ty, |args| { + env.register_native("/", div_ty, true, |args| { if args.len() != 2 { return Value::Void; } let a = match args[0] { Value::Int(i) => i as f64, Value::Float(f) => f, _ => return Value::Void }; let b = match args[1] { Value::Int(i) => i as f64, Value::Float(f) => f, _ => return Value::Void }; @@ -136,7 +136,7 @@ fn register_arithmetic(env: &Environment) { let int_div_ty = StaticType::Function(Box::new( Signature { params: StaticType::Tuple(vec![StaticType::Int, StaticType::Int]), ret: StaticType::Int } )); - env.register_native("//", int_div_ty, |args| { + env.register_native("//", int_div_ty, true, |args| { if args.len() != 2 { return Value::Void; } match (&args[0], &args[1]) { (Value::Int(a), Value::Int(b)) => { @@ -152,7 +152,7 @@ fn register_arithmetic(env: &Environment) { let mod_ty = StaticType::Function(Box::new( Signature { params: StaticType::Tuple(vec![StaticType::Int, StaticType::Int]), ret: StaticType::Int } )); - env.register_native("%", mod_ty, |args| { + env.register_native("%", mod_ty, true, |args| { if args.len() != 2 { return Value::Void; } match (&args[0], &args[1]) { (Value::Int(a), Value::Int(b)) => { @@ -171,7 +171,7 @@ fn register_comparison(env: &Environment) { ]); // --- Greater Than (>) --- - env.register_native(">", cmp_ty.clone(), |args| { + env.register_native(">", cmp_ty.clone(), true, |args| { if args.len() != 2 { return Value::Void; } match (&args[0], &args[1]) { (Value::Int(a), Value::Int(b)) => Value::Bool(a > b), @@ -184,7 +184,7 @@ fn register_comparison(env: &Environment) { }); // --- Less Than (<) --- - env.register_native("<", cmp_ty.clone(), |args| { + env.register_native("<", cmp_ty.clone(), true, |args| { if args.len() != 2 { return Value::Void; } match (&args[0], &args[1]) { (Value::Int(a), Value::Int(b)) => Value::Bool(a < b), @@ -197,7 +197,7 @@ fn register_comparison(env: &Environment) { }); // --- Greater Or Equal (>=) --- - env.register_native(">=", cmp_ty.clone(), |args| { + env.register_native(">=", cmp_ty.clone(), true, |args| { if args.len() != 2 { return Value::Void; } match (&args[0], &args[1]) { (Value::Int(a), Value::Int(b)) => Value::Bool(a >= b), @@ -210,7 +210,7 @@ fn register_comparison(env: &Environment) { }); // --- Less Or Equal (<=) --- - env.register_native("<=", cmp_ty.clone(), |args| { + env.register_native("<=", cmp_ty.clone(), true, |args| { if args.len() != 2 { return Value::Void; } match (&args[0], &args[1]) { (Value::Int(a), Value::Int(b)) => Value::Bool(a <= b), @@ -226,7 +226,7 @@ fn register_comparison(env: &Environment) { let eq_ty = StaticType::Function(Box::new( Signature { params: StaticType::Tuple(vec![StaticType::Any, StaticType::Any]), ret: StaticType::Bool } )); - env.register_native("=", eq_ty.clone(), |args| { + env.register_native("=", eq_ty.clone(), true, |args| { if args.len() != 2 { return Value::Void; } // Simple equality check. // Note: Floating point equality is tricky, but we follow standard behavior for now. @@ -245,7 +245,7 @@ fn register_comparison(env: &Environment) { }); // --- Not Equal (<>) --- - env.register_native("<>", eq_ty, |args| { + env.register_native("<>", eq_ty, true, |args| { if args.len() != 2 { return Value::Void; } match (&args[0], &args[1]) { (Value::Int(a), Value::Int(b)) => Value::Bool(a != b), @@ -268,7 +268,7 @@ fn register_logic(env: &Environment) { Signature { params: StaticType::Tuple(vec![StaticType::Bool]), ret: StaticType::Bool }, Signature { params: StaticType::Tuple(vec![StaticType::Int]), ret: StaticType::Int }, ]); - env.register_native("not", not_ty, |args| { + env.register_native("not", not_ty, true, |args| { if args.len() != 1 { return Value::Void; } match &args[0] { Value::Bool(b) => Value::Bool(!b), @@ -282,7 +282,7 @@ fn register_logic(env: &Environment) { Signature { params: StaticType::Tuple(vec![StaticType::Bool, StaticType::Bool]), ret: StaticType::Bool }, Signature { params: StaticType::Tuple(vec![StaticType::Int, StaticType::Int]), ret: StaticType::Int }, ]); - env.register_native("and", logic_op_ty.clone(), |args| { + env.register_native("and", logic_op_ty.clone(), true, |args| { if args.len() != 2 { return Value::Void; } match (&args[0], &args[1]) { (Value::Bool(a), Value::Bool(b)) => Value::Bool(*a && *b), @@ -292,7 +292,7 @@ fn register_logic(env: &Environment) { }); // --- Or (or) --- - env.register_native("or", logic_op_ty.clone(), |args| { + env.register_native("or", logic_op_ty.clone(), true, |args| { if args.len() != 2 { return Value::Void; } match (&args[0], &args[1]) { (Value::Bool(a), Value::Bool(b)) => Value::Bool(*a || *b), @@ -302,7 +302,7 @@ fn register_logic(env: &Environment) { }); // --- Xor (xor) --- - env.register_native("xor", logic_op_ty.clone(), |args| { + env.register_native("xor", logic_op_ty.clone(), true, |args| { if args.len() != 2 { return Value::Void; } match (&args[0], &args[1]) { (Value::Bool(a), Value::Bool(b)) => Value::Bool(*a ^ *b), @@ -315,7 +315,7 @@ fn register_logic(env: &Environment) { let shift_ty = StaticType::Function(Box::new( Signature { params: StaticType::Tuple(vec![StaticType::Int, StaticType::Int]), ret: StaticType::Int } )); - env.register_native("<<", shift_ty.clone(), |args| { + env.register_native("<<", shift_ty.clone(), true, |args| { if args.len() != 2 { return Value::Void; } match (&args[0], &args[1]) { (Value::Int(a), Value::Int(b)) => Value::Int(a << b), @@ -324,7 +324,7 @@ fn register_logic(env: &Environment) { }); // --- Shift Right (>>) --- - env.register_native(">>", shift_ty, |args| { + env.register_native(">>", shift_ty, true, |args| { if args.len() != 2 { return Value::Void; } match (&args[0], &args[1]) { (Value::Int(a), Value::Int(b)) => Value::Int(a >> b), diff --git a/src/ast/rtl/datetime.rs b/src/ast/rtl/datetime.rs index 35097db..75132e5 100644 --- a/src/ast/rtl/datetime.rs +++ b/src/ast/rtl/datetime.rs @@ -8,7 +8,7 @@ pub fn register(env: &Environment) { ret: StaticType::DateTime })); - env.register_native("date", date_ty, |args| { + env.register_native("date", date_ty, true, |args| { if let Value::Text(s) = &args[0] { // Try parse YYYY-MM-DD if let Ok(dt) = NaiveDate::parse_from_str(s, "%Y-%m-%d") { diff --git a/src/main.rs b/src/main.rs index 48b18ac..e435ee8 100644 --- a/src/main.rs +++ b/src/main.rs @@ -157,6 +157,7 @@ impl CompilerApp { fn dump_ast(&mut self) { self.env = Environment::new(); + self.env.optimization_level = 2; // Enable optimization for AST dump match self.env.dump_ast(&self.source_code) { Ok(dump) => { self.output_log = format!("--- BOUND AST DUMP ---\n\n{}", dump);