From 463b88413eb0a651037f99b3623d25be2f20c201 Mon Sep 17 00:00:00 2001 From: microproofs Date: Fri, 13 Dec 2024 14:15:19 +0700 Subject: [PATCH] feat: new builtins constr_index and constr_fields for alternative fast ways to take apart Data --- crates/aiken-lang/src/builtins.rs | 165 +++++++++++++++++- crates/aiken-lang/src/gen_uplc.rs | 89 +++------- crates/aiken-lang/src/gen_uplc/builder.rs | 7 +- .../src/gen_uplc/stick_break_set.rs | 18 +- crates/uplc/src/optimize/shrinker.rs | 11 ++ examples/acceptance_tests/116/aiken.toml | 2 + examples/acceptance_tests/116/lib/tests.ak | 12 ++ 7 files changed, 227 insertions(+), 77 deletions(-) create mode 100644 examples/acceptance_tests/116/aiken.toml create mode 100644 examples/acceptance_tests/116/lib/tests.ak diff --git a/crates/aiken-lang/src/builtins.rs b/crates/aiken-lang/src/builtins.rs index c3b6f3894..563d07b64 100644 --- a/crates/aiken-lang/src/builtins.rs +++ b/crates/aiken-lang/src/builtins.rs @@ -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 { 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 diff --git a/crates/aiken-lang/src/gen_uplc.rs b/crates/aiken-lang/src/gen_uplc.rs index c38fa8058..25c83a33f 100644 --- a/crates/aiken-lang/src/gen_uplc.rs +++ b/crates/aiken-lang/src/gen_uplc.rs @@ -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) } diff --git a/crates/aiken-lang/src/gen_uplc/builder.rs b/crates/aiken-lang/src/gen_uplc/builder.rs index 159ec328b..43df0668f 100644 --- a/crates/aiken-lang/src/gen_uplc/builder.rs +++ b/crates/aiken-lang/src/gen_uplc/builder.rs @@ -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, } } diff --git a/crates/aiken-lang/src/gen_uplc/stick_break_set.rs b/crates/aiken-lang/src/gen_uplc/stick_break_set.rs index 3e1568036..61e06bb55 100644 --- a/crates/aiken-lang/src/gen_uplc/stick_break_set.rs +++ b/crates/aiken-lang/src/gen_uplc/stick_break_set.rs @@ -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, - then: AirTree, - ) -> AirTree { + pub fn produce_air(self, prev_name: String, subject_tipo: Rc, 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, ) }) diff --git a/crates/uplc/src/optimize/shrinker.rs b/crates/uplc/src/optimize/shrinker.rs index 44c15412a..8dacdd069 100644 --- a/crates/uplc/src/optimize/shrinker.rs +++ b/crates/uplc/src/optimize/shrinker.rs @@ -1448,6 +1448,17 @@ impl Term { 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) + } } } } diff --git a/examples/acceptance_tests/116/aiken.toml b/examples/acceptance_tests/116/aiken.toml new file mode 100644 index 000000000..c7db180cf --- /dev/null +++ b/examples/acceptance_tests/116/aiken.toml @@ -0,0 +1,2 @@ +name = "aiken-lang/acceptance_test_116" +version = "0.0.0" diff --git a/examples/acceptance_tests/116/lib/tests.ak b/examples/acceptance_tests/116/lib/tests.ak new file mode 100644 index 000000000..29aaa94b1 --- /dev/null +++ b/examples/acceptance_tests/116/lib/tests.ak @@ -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", + } +}