diff --git a/Cargo.lock b/Cargo.lock index a7f53bf8..729f4617 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2480,6 +2480,7 @@ dependencies = [ "cryptoxide", "flat-rs", "hex", + "indexmap", "itertools", "pallas-addresses", "pallas-codec", diff --git a/crates/aiken-lang/src/uplc.rs b/crates/aiken-lang/src/uplc.rs index 5b7bd34d..5f7b90a0 100644 --- a/crates/aiken-lang/src/uplc.rs +++ b/crates/aiken-lang/src/uplc.rs @@ -1587,6 +1587,7 @@ impl<'a> CodeGenerator<'a> { ); } } + // TODO: Check constr for assert on all cases constr @ Pattern::Constructor { .. } => { if matches!(assignment_properties.kind, AssignmentKind::Assert) && assignment_properties.value_is_data diff --git a/crates/uplc/Cargo.toml b/crates/uplc/Cargo.toml index 2f33cdc2..53c89842 100644 --- a/crates/uplc/Cargo.toml +++ b/crates/uplc/Cargo.toml @@ -30,6 +30,7 @@ serde_json = "1.0.85" strum = "0.24.1" strum_macros = "0.24.3" itertools = "0.10.5" +indexmap = "1.9.2" [dev-dependencies] hex = "0.4.3" diff --git a/crates/uplc/src/ast.rs b/crates/uplc/src/ast.rs index 0d18be2b..20087b4f 100644 --- a/crates/uplc/src/ast.rs +++ b/crates/uplc/src/ast.rs @@ -1,5 +1,6 @@ use std::{ fmt::{self, Display}, + hash::{self, Hash}, rc::Rc, }; @@ -276,12 +277,19 @@ impl Display for Type { /// A Name containing it's parsed textual representation /// and a unique id from string interning. The Name's text is /// interned during parsing. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Eq)] pub struct Name { pub text: String, pub unique: Unique, } +impl hash::Hash for Name { + fn hash(&self, state: &mut H) { + self.text.hash(state); + self.unique.hash(state); + } +} + impl PartialEq for Name { fn eq(&self, other: &Self) -> bool { self.unique == other.unique @@ -326,7 +334,7 @@ impl Display for Unique { /// Similar to `Name` but for Debruijn indices. /// `Name` is replaced by `NamedDebruijn` when converting /// program to it's debruijn form. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Eq)] pub struct NamedDeBruijn { pub text: String, pub index: DeBruijn, diff --git a/crates/uplc/src/lib.rs b/crates/uplc/src/lib.rs index 6e240070..1a4623d8 100644 --- a/crates/uplc/src/lib.rs +++ b/crates/uplc/src/lib.rs @@ -7,6 +7,7 @@ pub mod parser; mod pretty; pub mod program_builder; pub mod tx; +pub mod optimize; pub use pallas_codec::utils::KeyValuePairs; pub use pallas_crypto::hash::Hash; diff --git a/crates/uplc/src/optimize.rs b/crates/uplc/src/optimize.rs new file mode 100644 index 00000000..3dd876e9 --- /dev/null +++ b/crates/uplc/src/optimize.rs @@ -0,0 +1,7 @@ +use crate::ast::{Name, Program}; + +pub mod shrinker; + +pub fn aiken_optimize(term: Program) -> Program { + term.builtin_force_reduce().lambda_reduce().inline_reduce() +} diff --git a/crates/uplc/src/optimize/shrinker.rs b/crates/uplc/src/optimize/shrinker.rs new file mode 100644 index 00000000..dd4818aa --- /dev/null +++ b/crates/uplc/src/optimize/shrinker.rs @@ -0,0 +1,265 @@ +use std::rc::Rc; + +use indexmap::IndexMap; +use itertools::Itertools; + +use crate::{ + ast::{builder::apply_wrap, Name, Program, Term}, + builtins::DefaultFunction, +}; +// use crate::builtins::{DefaultFunction}; + +#[derive(Eq, Hash, PartialEq, Clone)] +pub struct Occurrence { + name: Rc, + lambda_count: usize, +} + +impl Program { + pub fn lambda_reduce(self) -> Program { + let mut term = self.term.clone(); + lambda_reduce(&mut term); + Program { + version: self.version, + term, + } + } + + pub fn builtin_force_reduce(self) -> Program { + let mut term = self.term.clone(); + let mut builtin_map = IndexMap::new(); + builtin_force_reduce(&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(); + + term = apply_wrap( + Term::Lambda { + parameter_name: Name { + text: format!("__{}_wrapped", default_func.aiken_name()), + unique: 0.into(), + } + .into(), + body: term.into(), + }, + if default_func.force_count() == 1 { + Term::Builtin(default_func).force_wrap() + } else { + Term::Builtin(default_func).force_wrap().force_wrap() + }, + ); + } + + Program { + version: self.version, + term, + } + } + + pub fn inline_reduce(self) -> Program { + let mut term = self.term.clone(); + inline_reduce(&mut term); + Program { + version: self.version, + term, + } + } +} + +fn builtin_force_reduce(term: &mut Term, builtin_map: &mut IndexMap) { + match term { + Term::Force(f) => { + let f = Rc::make_mut(f); + + match f { + Term::Force(inner_f) => { + if let Term::Builtin(func) = inner_f.as_ref() { + builtin_map.insert(*func as u8, ()); + *term = Term::Var( + Name { + text: format!("__{}_wrapped", func.aiken_name()), + unique: 0.into(), + } + .into(), + ); + return; + } + } + Term::Builtin(func) => { + builtin_map.insert(*func as u8, ()); + *term = Term::Var( + Name { + text: format!("__{}", func.aiken_name()), + unique: 0.into(), + } + .into(), + ); + + return; + } + _ => {} + } + builtin_force_reduce(f, builtin_map); + } + Term::Delay(d) => { + let d = Rc::make_mut(d); + builtin_force_reduce(d, builtin_map); + } + Term::Lambda { body, .. } => { + let body = Rc::make_mut(body); + builtin_force_reduce(body, builtin_map); + } + Term::Apply { function, argument } => { + let func = Rc::make_mut(function); + builtin_force_reduce(func, builtin_map); + + let arg = Rc::make_mut(argument); + builtin_force_reduce(arg, builtin_map); + } + _ => {} + } +} + +fn inline_reduce(term: &mut Term) { + match term { + Term::Delay(d) => { + let d = Rc::make_mut(d); + inline_reduce(d); + } + Term::Lambda { body, .. } => { + let body = Rc::make_mut(body); + inline_reduce(body); + } + Term::Apply { function, argument } => { + let arg = Rc::make_mut(argument); + inline_reduce(arg); + + let func = Rc::make_mut(function); + inline_reduce(func); + + if let Term::Lambda { + parameter_name, + body, + } = func + { + let mut occurrences = 0; + var_occurrences(body, parameter_name.clone(), &mut occurrences); + if occurrences <= 1 { + *term = substitute_term(body.as_ref(), parameter_name.clone(), argument); + } + } + } + Term::Force(f) => { + let f = Rc::make_mut(f); + inline_reduce(f); + } + _ => {} + } +} + +fn var_occurrences(term: &Term, search_for: Rc, occurrences: &mut usize) { + match term { + Term::Var(name) => { + if name.clone() == search_for { + *occurrences += 1; + } + } + Term::Delay(body) => { + var_occurrences(body.as_ref(), search_for, occurrences); + } + Term::Lambda { + parameter_name, + body, + } => { + if parameter_name.clone() != search_for { + var_occurrences(body.as_ref(), search_for, occurrences); + } + } + Term::Apply { function, argument } => { + var_occurrences(function.as_ref(), search_for.clone(), occurrences); + var_occurrences(argument.as_ref(), search_for, occurrences); + } + Term::Force(x) => { + var_occurrences(x.as_ref(), search_for, occurrences); + } + _ => {} + } +} + +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(_)) = argument.as_ref() { + let body = Rc::make_mut(body); + lambda_reduce(body); + *body = 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); + } + _ => {} + } +} + +fn substitute_term(term: &Term, original: Rc, replace_with: &Term) -> Term { + match term { + Term::Var(name) => { + if name.clone() == original { + replace_with.clone() + } else { + Term::Var(name.clone()) + } + } + Term::Delay(body) => { + Term::Delay(substitute_term(body.as_ref(), original, replace_with).into()) + } + Term::Lambda { + parameter_name, + body, + } => { + if parameter_name.clone() != original { + Term::Lambda { + parameter_name: parameter_name.clone(), + body: Rc::new(substitute_term(body.as_ref(), original, replace_with)), + } + } else { + Term::Lambda { + parameter_name: parameter_name.clone(), + body: body.clone(), + } + } + } + Term::Apply { function, argument } => Term::Apply { + function: Rc::new(substitute_term( + function.as_ref(), + original.clone(), + replace_with, + )), + argument: Rc::new(substitute_term(argument.as_ref(), original, replace_with)), + }, + Term::Force(x) => Term::Force(Rc::new(substitute_term(x.as_ref(), original, replace_with))), + x => x.clone(), + } +}