Brummel/RustAst
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Implement purity inference for optimization

Browse Source 
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.
This commit is contained in:
Michael Schimmel
2026-02-22 00:16:04 +01:00
parent ac73aaf59e
commit 0107264026
6 changed files with 184 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/optimizer_purity.myc
+5
View File
@@ -0,0 +1,5 @@
(do
(def PI 3.1415)
(def area (fn [x] (* PI x)))
(area 10)
)
src/ast/compiler/optimizer.rs
+145 -52
View File
@@ -12,11 +12,12 @@ pub struct Optimizer {
pub level: u32,
max_passes: usize,
pub globals: Option<Rc<RefCell<Vec<Value>>>>,
pub global_purity: Option<Rc<RefCell<HashMap<u32, bool>>>>,
}
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<RefCell<Vec<Value>>>) -> Self {
@@ -24,6 +25,11 @@ impl Optimizer {
self
}
pub fn with_purity(mut self, purity: Rc<RefCell<HashMap<u32, bool>>>) -> 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::<Closure>() {
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<TypedNode> {
// 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::<Closure>() {
// 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<TypedNode> {
if self.contains_def_local(&body) {
return None;
@@ -405,42 +516,6 @@ impl Optimizer {
}
}
fn try_fold_intrinsic(&self, callee: &TypedNode, args: &TypedNode) -> Option<TypedNode> {
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<i64> {
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)))";
src/ast/environment.rs
+14 -2
View File
@@ -21,6 +21,7 @@ use crate::ast::rtl::intrinsics;
pub struct Environment {
pub global_names: Rc<RefCell<HashMap<Symbol, u32>>>,
pub global_types: Rc<RefCell<HashMap<u32, StaticType>>>,
pub global_purity: Rc<RefCell<HashMap<u32, bool>>>,
pub global_values: Rc<RefCell<Vec<Value>>>,
pub function_registry: Rc<RefCell<HashMap<u32, BoundNode>>>,
pub monomorph_cache: Rc<RefCell<MonoCache>>,
@@ -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>) -> 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)
src/ast/rtl/core.rs
+18 -18
View File
@@ -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),
src/ast/rtl/datetime.rs
+1 -1
View File
@@ -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") {
src/main.rs
+1
View File
@@ -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);