feat: new builtins constr_index and constr_fields for alternative fast ways to take apart Data

This commit is contained in:
microproofs 2024-12-13 14:15:19 +07:00
parent ebc7d89d5d
commit 463b88413e
No known key found for this signature in database
GPG Key ID: 14F93C84DE6AFD17
7 changed files with 227 additions and 77 deletions

View File

@ -15,7 +15,10 @@ use crate::{
use indexmap::IndexMap;
use std::{collections::HashMap, rc::Rc};
use strum::IntoEnumIterator;
use uplc::builtins::DefaultFunction;
use uplc::{
builder::{CONSTR_FIELDS_EXPOSER, CONSTR_INDEX_EXPOSER},
builtins::DefaultFunction,
};
pub const PRELUDE: &str = "aiken";
pub const BUILTIN: &str = "aiken/builtin";
@ -514,6 +517,38 @@ pub fn plutus(id_gen: &IdGenerator) -> TypeInfo {
}
}
let index_tipo = Type::function(vec![Type::data()], Type::int());
plutus.values.insert(
"constr_index".to_string(),
ValueConstructor::public(
index_tipo,
ValueConstructorVariant::ModuleFn {
name: "constr_index".to_string(),
field_map: None,
module: "aiken/builtin".to_string(),
arity: 1,
location: Span::empty(),
builtin: None,
},
),
);
let fields_tipo = Type::function(vec![Type::data()], Type::list(Type::data()));
plutus.values.insert(
"constr_fields".to_string(),
ValueConstructor::public(
fields_tipo,
ValueConstructorVariant::ModuleFn {
name: "constr_fields".to_string(),
field_map: None,
module: "aiken/builtin".to_string(),
arity: 1,
location: Span::empty(),
builtin: None,
},
),
);
plutus
}
@ -1008,6 +1043,134 @@ pub fn prelude_functions(
) -> IndexMap<FunctionAccessKey, TypedFunction> {
let mut functions = IndexMap::new();
let constr_index_body = TypedExpr::Call {
location: Span::empty(),
tipo: Type::int(),
fun: TypedExpr::local_var(
CONSTR_INDEX_EXPOSER,
Type::function(vec![Type::data()], Type::int()),
Span::empty(),
)
.into(),
args: vec![CallArg {
label: None,
location: Span::empty(),
value: TypedExpr::Var {
location: Span::empty(),
constructor: ValueConstructor {
public: true,
tipo: Type::data(),
variant: ValueConstructorVariant::LocalVariable {
location: Span::empty(),
},
},
name: "constr".to_string(),
},
}],
};
let constr_index_func = Function {
arguments: vec![TypedArg {
arg_name: ArgName::Named {
name: "constr".to_string(),
label: "constr".to_string(),
location: Span::empty(),
},
is_validator_param: false,
doc: None,
location: Span::empty(),
annotation: None,
tipo: Type::data(),
}],
on_test_failure: OnTestFailure::FailImmediately,
doc: Some(
indoc::indoc! {
r#"
/// Access the index of a constr typed as Data. Fails if the Data object is not a constr.
"#
}.to_string()
),
location: Span::empty(),
name: "constr_index".to_string(),
public: true,
return_annotation: None,
return_type: Type::int(),
end_position: 0,
body: constr_index_body,
};
functions.insert(
FunctionAccessKey {
module_name: "aiken/builtin".to_string(),
function_name: "constr_index".to_string(),
},
constr_index_func,
);
let constr_fields_body = TypedExpr::Call {
location: Span::empty(),
tipo: Type::list(Type::data()),
fun: TypedExpr::local_var(
CONSTR_FIELDS_EXPOSER,
Type::function(vec![Type::data()], Type::list(Type::data())),
Span::empty(),
)
.into(),
args: vec![CallArg {
label: None,
location: Span::empty(),
value: TypedExpr::Var {
location: Span::empty(),
constructor: ValueConstructor {
public: true,
tipo: Type::data(),
variant: ValueConstructorVariant::LocalVariable {
location: Span::empty(),
},
},
name: "constr".to_string(),
},
}],
};
let constr_fields_func = Function {
arguments: vec![TypedArg {
arg_name: ArgName::Named {
name: "constr".to_string(),
label: "constr".to_string(),
location: Span::empty(),
},
is_validator_param: false,
doc: None,
location: Span::empty(),
annotation: None,
tipo: Type::data(),
}],
on_test_failure: OnTestFailure::FailImmediately,
doc: Some(
indoc::indoc! {
r#"
/// Access the fields of a constr typed as Data. Fails if the Data object is not a constr.
"#
}.to_string()
),
location: Span::empty(),
name: "constr_fields".to_string(),
public: true,
return_annotation: None,
return_type: Type::list(Type::data()),
end_position: 0,
body: constr_fields_body,
};
functions.insert(
FunctionAccessKey {
module_name: "aiken/builtin".to_string(),
function_name: "constr_fields".to_string(),
},
constr_fields_func,
);
// /// Negate the argument. Useful for map/fold and pipelines.
// pub fn not(self: Bool) -> Bool {
// !self

View File

@ -324,7 +324,15 @@ impl<'a> CodeGenerator<'a> {
constructor, name, ..
} => match constructor.variant {
ValueConstructorVariant::LocalVariable { .. } => {
AirTree::var(constructor.clone(), self.interner.lookup_interned(name), "")
if name != CONSTR_INDEX_EXPOSER && name != CONSTR_FIELDS_EXPOSER {
AirTree::var(
constructor.clone(),
self.interner.lookup_interned(name),
"",
)
} else {
AirTree::var(constructor.clone(), name, "")
}
}
_ => AirTree::var(constructor.clone(), name, ""),
},
@ -2492,16 +2500,7 @@ impl<'a> CodeGenerator<'a> {
clauses,
);
builtins_to_add.produce_air(
// The only reason I pass this in is to ensure I signal
// whether or not constr_fields_exposer was used. I could
// probably optimize this part out to simplify codegen in
// the future
&mut self.special_functions,
prev_subject_name,
prev_tipo,
when_air_clauses,
)
builtins_to_add.produce_air(prev_subject_name, prev_tipo, when_air_clauses)
}
DecisionTree::ListSwitch {
path,
@ -2685,16 +2684,7 @@ impl<'a> CodeGenerator<'a> {
list_clauses.1,
);
builtins_to_add.produce_air(
// The only reason I pass this in is to ensure I signal
// whether or not constr_fields_exposer was used. I could
// probably optimize this part out to simplify codegen in
// the future
&mut self.special_functions,
prev_subject_name,
prev_tipo,
when_list_cases,
)
builtins_to_add.produce_air(prev_subject_name, prev_tipo, when_list_cases)
}
DecisionTree::HoistedLeaf(name, args) => {
let air_args = args
@ -2806,12 +2796,7 @@ impl<'a> CodeGenerator<'a> {
self.handle_assigns(subject_name, subject_tipo, rest, stick_set, then),
);
builtins_to_add.produce_air(
&mut self.special_functions,
prev_subject_name,
prev_tipo,
assignment,
)
builtins_to_add.produce_air(prev_subject_name, prev_tipo, assignment)
}
}
}
@ -3768,10 +3753,7 @@ impl<'a> CodeGenerator<'a> {
value = self.hoist_functions_to_validator(value);
let term = self
.uplc_code_gen(value.to_vec())
.constr_fields_exposer()
.constr_index_exposer();
let term = self.uplc_code_gen(value.to_vec());
let mut program =
self.new_program(self.special_functions.apply_used_functions(term));
@ -4511,11 +4493,7 @@ impl<'a> CodeGenerator<'a> {
Some(UplcType::Data) => subject,
None => Term::var(
self.special_functions
.use_function_uplc(CONSTR_INDEX_EXPOSER.to_string()),
)
.apply(subject),
None => Term::var(CONSTR_INDEX_EXPOSER).apply(subject),
};
let mut term = arg_stack.pop().unwrap();
@ -4741,13 +4719,9 @@ impl<'a> CodeGenerator<'a> {
Some(UplcType::Bls12_381G2Element) => Term::bls12_381_g2_equal()
.apply(checker)
.apply(Term::var(subject_name)),
None => Term::equals_integer().apply(checker).apply(
Term::var(
self.special_functions
.use_function_uplc(CONSTR_INDEX_EXPOSER.to_string()),
)
.apply(Term::var(subject_name)),
),
None => Term::equals_integer()
.apply(checker)
.apply(Term::var(CONSTR_INDEX_EXPOSER).apply(Term::var(subject_name))),
};
Some(condition.if_then_else(then.delay(), term).force())
@ -4934,13 +4908,7 @@ impl<'a> CodeGenerator<'a> {
otherwise,
);
term = term.apply(
Term::var(
self.special_functions
.use_function_uplc(CONSTR_FIELDS_EXPOSER.to_string()),
)
.apply(value),
);
term = term.apply(Term::var(CONSTR_FIELDS_EXPOSER).apply(value));
Some(term)
} else {
@ -4953,13 +4921,10 @@ impl<'a> CodeGenerator<'a> {
let mut term = arg_stack.pop().unwrap();
let otherwise = arg_stack.pop().unwrap();
term = Term::var(
self.special_functions
.use_function_uplc(CONSTR_FIELDS_EXPOSER.to_string()),
)
.apply(value)
.choose_list(term.delay(), otherwise)
.force();
term = Term::var(CONSTR_FIELDS_EXPOSER)
.apply(value)
.choose_list(term.delay(), otherwise)
.force();
Some(term)
}
@ -5133,13 +5098,9 @@ impl<'a> CodeGenerator<'a> {
}
}
term = term.lambda(format!("{tail_name_prefix}_0")).apply(
Term::var(
self.special_functions
.use_function_uplc(CONSTR_FIELDS_EXPOSER.to_string()),
)
.apply(record),
);
term = term
.lambda(format!("{tail_name_prefix}_0"))
.apply(Term::var(CONSTR_FIELDS_EXPOSER).apply(record));
Some(term)
}

View File

@ -100,7 +100,12 @@ impl CodeGenSpecialFuncs {
);
CodeGenSpecialFuncs {
used_funcs: vec![],
// Always use these functions since they are filtered out automatically by
// the optimization code later on
used_funcs: vec![
CONSTR_FIELDS_EXPOSER.to_string(),
CONSTR_INDEX_EXPOSER.to_string(),
],
key_to_func,
}
}

View File

@ -6,7 +6,6 @@ use uplc::{builder::CONSTR_FIELDS_EXPOSER, builtins::DefaultFunction};
use crate::expr::Type;
use super::{
builder::CodeGenSpecialFuncs,
decision_tree::{get_tipo_by_path, CaseTest, Path},
tree::AirTree,
};
@ -38,7 +37,7 @@ impl PartialEq for Builtin {
impl Eq for Builtin {}
impl Builtin {
fn produce_air(self, special_funcs: &mut CodeGenSpecialFuncs, arg: AirTree) -> AirTree {
fn produce_air(self, arg: AirTree) -> AirTree {
match self {
Builtin::HeadList(t) => AirTree::builtin(DefaultFunction::HeadList, t, vec![arg]),
Builtin::ExtractField(t) => AirTree::extract_field(t, arg),
@ -48,7 +47,10 @@ impl Builtin {
vec![arg],
),
Builtin::UnConstrFields => AirTree::call(
special_funcs.use_function_tree(CONSTR_FIELDS_EXPOSER.to_string()),
AirTree::local_var(
CONSTR_FIELDS_EXPOSER,
Type::function(vec![Type::data()], Type::list(Type::data())),
),
Type::list(Type::data()),
vec![arg],
),
@ -204,13 +206,7 @@ impl Builtins {
self
}
pub fn produce_air(
self,
special_funcs: &mut CodeGenSpecialFuncs,
prev_name: String,
subject_tipo: Rc<Type>,
then: AirTree,
) -> AirTree {
pub fn produce_air(self, prev_name: String, subject_tipo: Rc<Type>, then: AirTree) -> AirTree {
let (_, _, name_builtins) = self.vec.into_iter().fold(
(prev_name, subject_tipo, vec![]),
|(prev_name, prev_tipo, mut acc), item| {
@ -228,7 +224,7 @@ impl Builtins {
.rfold(then, |then, (prev_name, prev_tipo, next_name, builtin)| {
AirTree::let_assignment(
next_name,
builtin.produce_air(special_funcs, AirTree::local_var(prev_name, prev_tipo)),
builtin.produce_air(AirTree::local_var(prev_name, prev_tipo)),
then,
)
})

View File

@ -1448,6 +1448,17 @@ impl Term<Name> {
Term::unconstr_data().apply(std::mem::replace(arg, Term::Error.force())),
)
}
} else if let Term::Lambda {
parameter_name,
body,
} = Rc::make_mut(function)
{
if parameter_name.text == CONSTR_INDEX_EXPOSER
|| parameter_name.text == CONSTR_FIELDS_EXPOSER
{
let body = Rc::make_mut(body);
*self = std::mem::replace(body, Term::Error)
}
}
}
}

View File

@ -0,0 +1,2 @@
name = "aiken-lang/acceptance_test_116"
version = "0.0.0"

View File

@ -0,0 +1,12 @@
use aiken/builtin.{constr_index}
test baz() {
let x = Some("bar")
expect [bar] = x |> builtin.constr_fields
and {
constr_index(x) == 0,
builtin.un_b_data(bar) == "bar",
}
}