diff --git a/crates/uplc/src/optimize.rs b/crates/uplc/src/optimize.rs index 02501b5a..f34daea0 100644 --- a/crates/uplc/src/optimize.rs +++ b/crates/uplc/src/optimize.rs @@ -6,7 +6,7 @@ use crate::{ pub mod shrinker; pub fn aiken_optimize_and_intern(program: Program) -> Program { - let mut program = program.builtin_force_reduce(); + let mut program = program.builtin_force_reducer(); let mut interner = Interner::new(); @@ -18,14 +18,14 @@ pub fn aiken_optimize_and_intern(program: Program) -> Program { let program: Program = program_named.try_into().unwrap(); program - .lambda_reduce() - .inline_reduce() - .lambda_reduce() - .inline_reduce() - .force_delay_reduce() - .wrap_data_reduce() - .lambda_reduce() - .inline_reduce() - .lambda_reduce() - .inline_reduce() + .lambda_reducer() + .inline_reducer() + .lambda_reducer() + .inline_reducer() + .force_delay_reducer() + .cast_data_reducer() + .lambda_reducer() + .inline_reducer() + .lambda_reducer() + .inline_reducer() } diff --git a/crates/uplc/src/optimize/shrinker.rs b/crates/uplc/src/optimize/shrinker.rs index 2cb7a161..dcb0fe44 100644 --- a/crates/uplc/src/optimize/shrinker.rs +++ b/crates/uplc/src/optimize/shrinker.rs @@ -16,19 +16,19 @@ pub struct Occurrence { } impl Program { - pub fn lambda_reduce(self) -> Program { + pub fn lambda_reducer(self) -> Program { let mut term = self.term; - lambda_reduce(&mut term); + lambda_reducer(&mut term); Program { version: self.version, term, } } - pub fn builtin_force_reduce(self) -> Program { + pub fn builtin_force_reducer(self) -> Program { let mut term = self.term; let mut builtin_map = IndexMap::new(); - builtin_force_reduce(&mut term, &mut builtin_map); + builtin_force_reducer(&mut term, &mut builtin_map); for default_func_index in builtin_map.keys().sorted().cloned() { let default_func: DefaultFunction = default_func_index.try_into().unwrap(); @@ -48,11 +48,11 @@ impl Program { } } - pub fn inline_reduce(self) -> Program { + pub fn inline_reducer(self) -> Program { let mut term = self.term; - inline_basic_reduce(&mut term); - inline_direct_reduce(&mut term); - inline_identity_reduce(&mut term); + inline_single_occurrence_reducer(&mut term); + inline_direct_reducer(&mut term); + inline_identity_reducer(&mut term); Program { version: self.version, @@ -60,18 +60,18 @@ impl Program { } } - pub fn force_delay_reduce(self) -> Program { + pub fn force_delay_reducer(self) -> Program { let mut term = self.term; - force_delay_reduce(&mut term); + force_delay_reducer(&mut term); Program { version: self.version, term, } } - pub fn wrap_data_reduce(self) -> Program { + pub fn cast_data_reducer(self) -> Program { let mut term = self.term; - wrap_data_reduce(&mut term); + cast_data_reducer(&mut term); Program { version: self.version, term, @@ -79,7 +79,7 @@ impl Program { } } -fn builtin_force_reduce(term: &mut Term, builtin_map: &mut IndexMap) { +fn builtin_force_reducer(term: &mut Term, builtin_map: &mut IndexMap) { match term { Term::Force(f) => { let f = Rc::make_mut(f); @@ -112,22 +112,22 @@ fn builtin_force_reduce(term: &mut Term, builtin_map: &mut IndexMap {} } - builtin_force_reduce(f, builtin_map); + builtin_force_reducer(f, builtin_map); } Term::Delay(d) => { let d = Rc::make_mut(d); - builtin_force_reduce(d, builtin_map); + builtin_force_reducer(d, builtin_map); } Term::Lambda { body, .. } => { let body = Rc::make_mut(body); - builtin_force_reduce(body, builtin_map); + builtin_force_reducer(body, builtin_map); } Term::Apply { function, argument } => { let func = Rc::make_mut(function); - builtin_force_reduce(func, builtin_map); + builtin_force_reducer(func, builtin_map); let arg = Rc::make_mut(argument); - builtin_force_reduce(arg, builtin_map); + builtin_force_reducer(arg, builtin_map); } Term::Case { .. } => todo!(), Term::Constr { .. } => todo!(), @@ -135,32 +135,32 @@ fn builtin_force_reduce(term: &mut Term, builtin_map: &mut IndexMap) { +fn force_delay_reducer(term: &mut Term) { match term { Term::Force(f) => { let f = Rc::make_mut(f); if let Term::Delay(body) = f { *term = body.as_ref().clone(); - force_delay_reduce(term); + force_delay_reducer(term); } else { - force_delay_reduce(f); + force_delay_reducer(f); } } Term::Delay(d) => { let d = Rc::make_mut(d); - force_delay_reduce(d); + force_delay_reducer(d); } Term::Lambda { body, .. } => { let body = Rc::make_mut(body); - force_delay_reduce(body); + force_delay_reducer(body); } Term::Apply { function, argument } => { let func = Rc::make_mut(function); - force_delay_reduce(func); + force_delay_reducer(func); let arg = Rc::make_mut(argument); - force_delay_reduce(arg); + force_delay_reducer(arg); } Term::Case { .. } => todo!(), Term::Constr { .. } => todo!(), @@ -168,22 +168,62 @@ fn force_delay_reduce(term: &mut Term) { } } -fn inline_direct_reduce(term: &mut Term) { +fn lambda_reducer(term: &mut Term) { match term { + // TODO: change this to handle any amount of consecutive applies and lambdas + Term::Apply { function, argument } => { + let func = Rc::make_mut(function); + lambda_reducer(func); + + let arg = Rc::make_mut(argument); + lambda_reducer(arg); + + if let Term::Lambda { + parameter_name, + body, + } = func + { + if let replace_term @ (Term::Var(_) | Term::Constant(_) | Term::Builtin(_)) = arg { + let body = Rc::make_mut(body); + *term = substitute_term(body, parameter_name.clone(), replace_term); + } + } + } Term::Delay(d) => { let d = Rc::make_mut(d); - inline_direct_reduce(d); + lambda_reducer(d); } Term::Lambda { body, .. } => { let body = Rc::make_mut(body); - inline_direct_reduce(body); + lambda_reducer(body); } + Term::Force(f) => { + let f = Rc::make_mut(f); + lambda_reducer(f); + } + Term::Case { .. } => todo!(), + Term::Constr { .. } => todo!(), + _ => {} + } +} + +fn inline_direct_reducer(term: &mut Term) { + match term { + Term::Delay(d) => { + let d = Rc::make_mut(d); + inline_direct_reducer(d); + } + Term::Lambda { body, .. } => { + let body = Rc::make_mut(body); + inline_direct_reducer(body); + } + Term::Apply { function, argument } => { let func = Rc::make_mut(function); let arg = Rc::make_mut(argument); - inline_direct_reduce(func); - inline_direct_reduce(arg); + inline_direct_reducer(func); + inline_direct_reducer(arg); let Term::Lambda { parameter_name, @@ -203,7 +243,7 @@ fn inline_direct_reduce(term: &mut Term) { } Term::Force(f) => { let f = Rc::make_mut(f); - inline_direct_reduce(f); + inline_direct_reducer(f); } Term::Case { .. } => todo!(), Term::Constr { .. } => todo!(), @@ -211,22 +251,22 @@ fn inline_direct_reduce(term: &mut Term) { } } -fn inline_identity_reduce(term: &mut Term) { +fn inline_identity_reducer(term: &mut Term) { match term { Term::Delay(d) => { let d = Rc::make_mut(d); - inline_identity_reduce(d); + inline_identity_reducer(d); } Term::Lambda { body, .. } => { let body = Rc::make_mut(body); - inline_identity_reduce(body); + inline_identity_reducer(body); } Term::Apply { function, argument } => { let func = Rc::make_mut(function); let arg = Rc::make_mut(argument); - inline_identity_reduce(func); - inline_identity_reduce(arg); + inline_identity_reducer(func); + inline_identity_reducer(arg); let Term::Lambda { parameter_name, @@ -260,7 +300,7 @@ fn inline_identity_reduce(term: &mut Term) { } Term::Force(f) => { let f = Rc::make_mut(f); - inline_identity_reduce(f); + inline_identity_reducer(f); } Term::Case { .. } => todo!(), Term::Constr { .. } => todo!(), @@ -268,22 +308,23 @@ fn inline_identity_reduce(term: &mut Term) { } } -fn inline_basic_reduce(term: &mut Term) { +fn inline_single_occurrence_reducer(term: &mut Term) { match term { Term::Delay(d) => { let d = Rc::make_mut(d); - inline_basic_reduce(d); + inline_single_occurrence_reducer(d); } Term::Lambda { body, .. } => { let body = Rc::make_mut(body); - inline_basic_reduce(body); + inline_single_occurrence_reducer(body); } + // TODO: change this to handle any amount of consecutive applies and lambdas Term::Apply { function, argument } => { let func = Rc::make_mut(function); let arg = Rc::make_mut(argument); - inline_basic_reduce(func); - inline_basic_reduce(arg); + inline_single_occurrence_reducer(func); + inline_single_occurrence_reducer(arg); if let Term::Lambda { parameter_name, @@ -291,9 +332,11 @@ fn inline_basic_reduce(term: &mut Term) { } = func { let occurrences = var_occurrences(body, parameter_name.clone()); - // TODO: Once case and constr is live we need to also check if case is delaying values + let delays = delayed_execution(body.as_ref()); - if occurrences == 1 { + let recursive = is_recursive(arg); + // We don't want to inline recursive functions + if occurrences == 1 && !recursive { if delays == 0 { *term = substitute_term(body.as_ref(), parameter_name.clone(), arg); } else if let Term::Var(_) @@ -320,7 +363,7 @@ fn inline_basic_reduce(term: &mut Term) { } Term::Force(f) => { let f = Rc::make_mut(f); - inline_basic_reduce(f); + inline_single_occurrence_reducer(f); } Term::Case { .. } => todo!(), Term::Constr { .. } => todo!(), @@ -328,18 +371,18 @@ fn inline_basic_reduce(term: &mut Term) { } } -fn wrap_data_reduce(term: &mut Term) { +fn cast_data_reducer(term: &mut Term) { match term { Term::Delay(d) => { - wrap_data_reduce(Rc::make_mut(d)); + cast_data_reducer(Rc::make_mut(d)); } Term::Lambda { body, .. } => { - wrap_data_reduce(Rc::make_mut(body)); + cast_data_reducer(Rc::make_mut(body)); } Term::Apply { function, argument } => { let Term::Builtin(first_action) = function.as_ref() else { - wrap_data_reduce(Rc::make_mut(function)); - wrap_data_reduce(Rc::make_mut(argument)); + cast_data_reducer(Rc::make_mut(function)); + cast_data_reducer(Rc::make_mut(argument)); return; }; @@ -349,7 +392,7 @@ fn wrap_data_reduce(term: &mut Term) { } = Rc::make_mut(argument) { let Term::Builtin(second_action) = inner_func.as_ref() else { - wrap_data_reduce(Rc::make_mut(argument)); + cast_data_reducer(Rc::make_mut(argument)); return; }; @@ -362,11 +405,11 @@ fn wrap_data_reduce(term: &mut Term) { | (DefaultFunction::UnListData, DefaultFunction::ListData) | (DefaultFunction::MapData, DefaultFunction::UnMapData) | (DefaultFunction::UnMapData, DefaultFunction::MapData) => { - wrap_data_reduce(Rc::make_mut(inner_arg)); + cast_data_reducer(Rc::make_mut(inner_arg)); *term = inner_arg.as_ref().clone(); } _ => { - wrap_data_reduce(Rc::make_mut(argument)); + cast_data_reducer(Rc::make_mut(argument)); } } } else if let Term::Constant(c) = Rc::make_mut(argument) { @@ -435,15 +478,15 @@ fn wrap_data_reduce(term: &mut Term) { ); } _ => { - wrap_data_reduce(Rc::make_mut(argument)); + cast_data_reducer(Rc::make_mut(argument)); } } } else { - wrap_data_reduce(Rc::make_mut(argument)); + cast_data_reducer(Rc::make_mut(argument)); }; } Term::Force(f) => { - wrap_data_reduce(Rc::make_mut(f)); + cast_data_reducer(Rc::make_mut(f)); } Term::Case { .. } => todo!(), Term::Constr { .. } => todo!(), @@ -482,6 +525,24 @@ fn var_occurrences(term: &Term, search_for: Rc) -> usize { } } +fn is_recursive(term: &Term) -> bool { + match term { + Term::Delay(body) => is_recursive(body.as_ref()), + Term::Lambda { body, .. } => is_recursive(body.as_ref()), + Term::Apply { function, argument } => { + if let (Term::Var(a), Term::Var(b)) = (function.as_ref(), argument.as_ref()) { + a.as_ref() == b.as_ref() + } else { + is_recursive(function.as_ref()) || is_recursive(argument.as_ref()) + } + } + Term::Force(x) => is_recursive(x.as_ref()), + Term::Case { .. } => todo!(), + Term::Constr { .. } => todo!(), + _ => false, + } +} + fn delayed_execution(term: &Term) -> usize { match term { Term::Delay(body) => 1 + delayed_execution(body.as_ref()), @@ -503,44 +564,6 @@ fn delayed_execution(term: &Term) -> usize { } } -fn lambda_reduce(term: &mut Term) { - match term { - Term::Apply { function, argument } => { - let func = Rc::make_mut(function); - lambda_reduce(func); - - let arg = Rc::make_mut(argument); - lambda_reduce(arg); - - if let Term::Lambda { - parameter_name, - body, - } = func - { - if let replace_term @ (Term::Var(_) | Term::Constant(_) | Term::Builtin(_)) = arg { - let body = Rc::make_mut(body); - *term = substitute_term(body, parameter_name.clone(), replace_term); - } - } - } - Term::Delay(d) => { - let d = Rc::make_mut(d); - lambda_reduce(d); - } - Term::Lambda { body, .. } => { - let body = Rc::make_mut(body); - lambda_reduce(body); - } - Term::Force(f) => { - let f = Rc::make_mut(f); - lambda_reduce(f); - } - Term::Case { .. } => todo!(), - Term::Constr { .. } => todo!(), - _ => {} - } -} - fn substitute_term(term: &Term, original: Rc, replace_with: &Term) -> Term { match term { Term::Var(name) => { @@ -678,7 +701,7 @@ mod tests { interner.program(&mut expected); let expected: Program = expected.try_into().unwrap(); - let actual = program.lambda_reduce(); + let actual = program.lambda_reducer(); let actual: Program = actual.try_into().unwrap(); @@ -708,7 +731,7 @@ mod tests { interner.program(&mut expected); let expected: Program = expected.try_into().unwrap(); - let actual = program.lambda_reduce(); + let actual = program.lambda_reducer(); let actual: Program = actual.try_into().unwrap(); @@ -737,7 +760,7 @@ mod tests { let expected: Program = expected.try_into().unwrap(); - let actual = program.lambda_reduce(); + let actual = program.lambda_reducer(); let actual: Program = actual.try_into().unwrap(); @@ -788,7 +811,7 @@ mod tests { let expected: Program = expected.try_into().unwrap(); - let actual = program.lambda_reduce(); + let actual = program.lambda_reducer(); let actual: Program = actual.try_into().unwrap(); @@ -827,7 +850,7 @@ mod tests { let expected: Program = expected.try_into().unwrap(); - let actual = program.wrap_data_reduce(); + let actual = program.cast_data_reducer(); let actual: Program = actual.try_into().unwrap(); @@ -864,7 +887,7 @@ mod tests { let expected: Program = expected.try_into().unwrap(); - let actual = program.wrap_data_reduce(); + let actual = program.cast_data_reducer(); let actual: Program = actual.try_into().unwrap(); @@ -906,7 +929,7 @@ mod tests { let expected: Program = expected.try_into().unwrap(); - let mut actual = program.builtin_force_reduce(); + let mut actual = program.builtin_force_reducer(); let mut interner = Interner::new(); @@ -965,7 +988,7 @@ mod tests { let expected: Program = expected.try_into().unwrap(); - let mut actual = program.builtin_force_reduce(); + let mut actual = program.builtin_force_reducer(); let mut interner = Interner::new(); @@ -1001,7 +1024,7 @@ mod tests { let expected: Program = expected.try_into().unwrap(); - let actual = program.inline_reduce(); + let actual = program.inline_reducer(); let actual: Program = actual.try_into().unwrap(); @@ -1035,7 +1058,7 @@ mod tests { let expected: Program = expected.try_into().unwrap(); - let actual = program.inline_reduce(); + let actual = program.inline_reducer(); let actual: Program = actual.try_into().unwrap(); @@ -1087,7 +1110,7 @@ mod tests { let expected: Program = expected.try_into().unwrap(); - let actual = program.inline_reduce(); + let actual = program.inline_reducer(); let actual: Program = actual.try_into().unwrap();