Skip to content

Commit

Permalink
Minor fix to optimization to actually detect vars that are just force…
Browse files Browse the repository at this point in the history
…d builtins
  • Loading branch information
MicroProofs committed Jan 17, 2025
1 parent 91b6e6d commit 1075be1
Showing 1 changed file with 23 additions and 13 deletions.
36 changes: 23 additions & 13 deletions crates/uplc/src/optimize/shrinker.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ use indexmap::IndexMap;
use itertools::{FoldWhile, Itertools};
use pallas_primitives::conway::{BigInt, PlutusData};
use std::{cmp::Ordering, iter, ops::Neg, rc::Rc};
use strum::IntoEnumIterator;

#[derive(Eq, Hash, PartialEq, Clone, Debug, PartialOrd)]
pub enum ScopePath {
Expand Down Expand Up @@ -348,6 +349,12 @@ impl DefaultFunction {
format!("__{}_wrapped", self.aiken_name())
}
}
pub fn forceable_wrapped_names() -> Vec<String> {
DefaultFunction::iter()
.filter(|df| df.force_count() > 0)
.map(|df| df.wrapped_name())
.collect_vec()
}

#[derive(PartialEq, Clone, Debug)]
pub enum BuiltinArgs {
Expand Down Expand Up @@ -1316,7 +1323,7 @@ impl Term<Name> {
// So it costs more size to have them hoisted
Term::Delay(e) if matches!(e.as_ref(), Term::Error) => true,
// If it wraps a builtin with consts or arguments passed in then inline
Term::Lambda { .. } => arg_term.is_a_builtin_wrapper(context),
Term::Lambda { .. } => arg_term.is_a_builtin_wrapper(),
// Inline smaller terms too
Term::Constant(_) | Term::Var(_) | Term::Builtin(_) => true,

Expand Down Expand Up @@ -1430,15 +1437,15 @@ impl Term<Name> {

arg.split_body_lambda();

arg_stack.push(std::mem::replace(arg, Term::Error.force()));
arg_stack.push(Args::Apply(0, std::mem::replace(arg, Term::Error.force())));
}
Term::Lambda {
parameter_name,
body,
} => {
current_term = Rc::make_mut(body);

if let Some(arg) = arg_stack.pop() {
if let Some(Args::Apply(_, arg)) = arg_stack.pop() {
let names = arg.get_var_names();

let func = (parameter_name.clone(), arg);
Expand Down Expand Up @@ -1467,7 +1474,12 @@ impl Term<Name> {
unsat_lams.push(parameter_name.clone());
}
}
Term::Delay(term) | Term::Force(term) => {
Term::Force(term) => {
current_term = Rc::make_mut(term);

arg_stack.push(Args::Force(0));
}
Term::Delay(term) => {
Rc::make_mut(term).split_body_lambda();
break;
}
Expand Down Expand Up @@ -1522,7 +1534,10 @@ impl Term<Name> {
// Replace args that weren't consumed
let term = arg_stack
.into_iter()
.rfold(term_to_build_on, |term, arg| term.apply(arg));
.rfold(term_to_build_on, |term, arg| match arg {
Args::Force(_) => term.force(),
Args::Apply(_, arg) => term.apply(arg),
});

let term = function_groups.into_iter().rfold(term, |term, group| {
let term = group.iter().rfold(term, |term, (name, _)| Term::Lambda {
Expand Down Expand Up @@ -2212,7 +2227,7 @@ impl Term<Name> {
std::mem::replace(term, Term::Error.force())
}

fn is_a_builtin_wrapper(&self, context: &Context) -> bool {
fn is_a_builtin_wrapper(&self) -> bool {
let (names, term) = self.pop_lambdas_and_get_names();

let mut arg_names = vec![];
Expand All @@ -2233,12 +2248,7 @@ impl Term<Name> {
}

let func_is_builtin = match term {
Term::Var(name) => context
.builtins_map
.keys()
.map(|func| func.wrapped_name())
.any(|func| func == name.text),

Term::Var(name) => forceable_wrapped_names().contains(&name.text),
Term::Builtin(_) => true,
_ => false,
};
Expand Down Expand Up @@ -2508,7 +2518,7 @@ impl Program<Name> {

id_vec
} else {
// Brand new buitlin so we add it to the list
// Brand new builtin so we add it to the list
let curried_builtin = builtin_args.clone().args_to_curried_args(*func);

let Some(id_vec) = curried_builtin.get_id_args(builtin_args) else {
Expand Down

0 comments on commit 1075be1

Please sign in to comment.