Skip to content

Commit

Permalink
feat: new builtins constr_index and constr_fields for alternative fas…
Browse files Browse the repository at this point in the history 
…t ways to take apart Data
  • Loading branch information
@MicroProofs
MicroProofs committed Dec 13, 2024
1 parent ebc7d89 commit 463b884
Show file tree
Hide file tree
Showing 7 changed files with 227 additions and 77 deletions.
165 changes: 164 additions & 1 deletion crates/aiken-lang/src/builtins.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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
}

Expand Down Expand Up @@ -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
Expand Down
89 changes: 25 additions & 64 deletions crates/aiken-lang/src/gen_uplc.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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, ""),
},
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
}
}
}
Expand Down Expand Up @@ -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));
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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())
Expand Down Expand Up @@ -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 {
Expand All @@ -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)
}
Expand Down Expand Up @@ -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)
}
Expand Down
7 changes: 6 additions & 1 deletion crates/aiken-lang/src/gen_uplc/builder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
}
Expand Down
Loading

0 comments on commit 463b884

Please sign in to comment.