Skip to content

Commit

Permalink
Final update to bls-constants may revist later
Browse files Browse the repository at this point in the history
  • Loading branch information
MicroProofs committed Oct 1, 2024
1 parent f169e73 commit 3246362
Show file tree
Hide file tree
Showing 3 changed files with 233 additions and 34 deletions.
172 changes: 140 additions & 32 deletions crates/aiken-lang/src/gen_uplc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,6 @@ const DELAY_ERROR: fn() -> AirTree =

#[derive(Clone)]
pub struct CodeGenerator<'a> {
#[allow(dead_code)]
plutus_version: PlutusVersion,
/// immutable index maps
functions: IndexMap<&'a FunctionAccessKey, &'a TypedFunction>,
Expand All @@ -81,6 +80,7 @@ pub struct CodeGenerator<'a> {
cyclic_functions:
IndexMap<(FunctionAccessKey, Variant), (CycleFunctionNames, usize, FunctionAccessKey)>,
monomorphized_consts: IndexMap<(FunctionAccessKey, String), AirTree>,
uplc_resolved_consts: IndexMap<(FunctionAccessKey, String), Term<Name>>,
/// mutable and reset as well
interner: AirInterner,
id_gen: IdGenerator,
Expand Down Expand Up @@ -113,6 +113,7 @@ impl<'a> CodeGenerator<'a> {
code_gen_functions: IndexMap::new(),
cyclic_functions: IndexMap::new(),
monomorphized_consts: IndexMap::new(),
uplc_resolved_consts: IndexMap::new(),
interner: AirInterner::new(),
id_gen: IdGenerator::new(),
}
Expand All @@ -123,6 +124,7 @@ impl<'a> CodeGenerator<'a> {
self.defined_functions = IndexMap::new();
self.cyclic_functions = IndexMap::new();
self.monomorphized_consts = IndexMap::new();
self.uplc_resolved_consts = IndexMap::new();
self.interner = AirInterner::new();
self.id_gen = IdGenerator::new();
if reset_special_functions {
Expand Down Expand Up @@ -154,6 +156,43 @@ impl<'a> CodeGenerator<'a> {

let full_vec = full_tree.to_vec();

self.uplc_resolved_consts = self
.monomorphized_consts
.clone()
.into_iter()
.map(|item| {
let (key, value) = item;

let value = self.hoist_functions_to_validator(value);

let const_term = self
.uplc_code_gen(value.to_vec())
.constr_fields_exposer()
.constr_index_exposer();

let mut program =
self.new_program(self.special_functions.apply_used_functions(const_term));

let mut interner = CodeGenInterner::new();

interner.program(&mut program);

let eval_program: Program<NamedDeBruijn> =
program.remove_no_inlines().try_into().unwrap();

let const_body = eval_program
.eval(ExBudget::max())
.result()
.unwrap_or_else(|e| panic!("Failed to evaluate constant: {e:#?}"))
.try_into()
.unwrap();

(key.clone(), const_body)
})
.collect();

println!("TOOOO {:#?}", self.uplc_resolved_consts);

let term = self.uplc_code_gen(full_vec);

let term = cast_validator_args(term, &validator.params, &self.interner);
Expand Down Expand Up @@ -189,6 +228,43 @@ impl<'a> CodeGenerator<'a> {
// optimizations on air tree
let full_vec = full_tree.to_vec();

self.uplc_resolved_consts = self
.monomorphized_consts
.clone()
.into_iter()
.map(|item| {
let (key, value) = item;

let value = self.hoist_functions_to_validator(value);

let const_term = self
.uplc_code_gen(value.to_vec())
.constr_fields_exposer()
.constr_index_exposer();

let mut program =
self.new_program(self.special_functions.apply_used_functions(const_term));

let mut interner = CodeGenInterner::new();

interner.program(&mut program);

let eval_program: Program<NamedDeBruijn> =
program.remove_no_inlines().try_into().unwrap();

let const_body = eval_program
.eval(ExBudget::max())
.result()
.unwrap_or_else(|e| panic!("Failed to evaluate constant: {e:#?}"))
.try_into()
.unwrap();

(key.clone(), const_body)
})
.collect();

println!("TOOOO {:#?}", self.uplc_resolved_consts);

let mut term = self.uplc_code_gen(full_vec);

term = if args.is_empty() {
Expand Down Expand Up @@ -218,6 +294,8 @@ impl<'a> CodeGenerator<'a> {
fn finalize(&mut self, mut term: Term<Name>) -> Program<Name> {
term = self.special_functions.apply_used_functions(term);

println!("TERM IS {}", term.to_pretty());

let program = aiken_optimize_and_intern(self.new_program(term));

// This is very important to call here.
Expand Down Expand Up @@ -4527,19 +4605,31 @@ impl<'a> CodeGenerator<'a> {
.into(),
)),
ValueConstructorVariant::ModuleConstant { module, name, .. } => {
let name = if !module.is_empty() {
let uplc_name = if !module.is_empty() {
format!("{module}_{name}{variant_name}")
} else {
format!("{name}{variant_name}")
};

Some(Term::Var(
Name {
text: name,
unique: 0.into(),
}
.into(),
))
let existing_term = self.uplc_resolved_consts.get(&(
FunctionAccessKey {
module_name: module.clone(),
function_name: name.clone(),
},
variant_name.clone(),
));

match existing_term {
Some(constant @ Term::Constant(_)) => Some(constant.clone()),

_ => Some(Term::Var(
Name {
text: uplc_name,
unique: 0.into(),
}
.into(),
)),
}
}
ValueConstructorVariant::ModuleFn {
name: func_name,
Expand Down Expand Up @@ -5150,37 +5240,55 @@ impl<'a> CodeGenerator<'a> {
function_name: func_name.clone(),
};

let mut value = self
.monomorphized_consts
.get(&(access_key, variant_name))
.cloned()
.unwrap();
let existing_term = self
.uplc_resolved_consts
.get(&(access_key.clone(), variant_name.clone()));

value = self.hoist_functions_to_validator(value);
match existing_term {
Some(Term::Constant(_)) => Some(term),

let const_term = self
.uplc_code_gen(value.to_vec())
.constr_fields_exposer()
.constr_index_exposer();
_ => {
let mut value = self
.monomorphized_consts
.get(&(access_key.clone(), variant_name.clone()))
.cloned()
.unwrap();

let mut program = self
.new_program(self.special_functions.apply_used_functions(const_term));
value = self.hoist_functions_to_validator(value);

let mut interner = CodeGenInterner::new();
let const_term = self
.uplc_code_gen(value.to_vec())
.constr_fields_exposer()
.constr_index_exposer();

interner.program(&mut program);
let mut program = self.new_program(
self.special_functions.apply_used_functions(const_term),
);

let eval_program: Program<NamedDeBruijn> =
program.remove_no_inlines().try_into().unwrap();
let mut interner = CodeGenInterner::new();

let const_body = eval_program
.eval(ExBudget::max())
.result()
.unwrap_or_else(|e| panic!("Failed to evaluate constant: {e:#?}"))
.try_into()
.unwrap();
interner.program(&mut program);

let eval_program: Program<NamedDeBruijn> =
program.remove_no_inlines().try_into().unwrap();

Some(term.lambda(func_uplc_name).apply(const_body))
let const_body: Term<Name> = eval_program
.eval(ExBudget::max())
.result()
.unwrap_or_else(|e| {
panic!("Failed to evaluate constant: {e:#?}")
})
.try_into()
.unwrap();

self.uplc_resolved_consts.insert(
(access_key.clone(), variant_name.clone()),
const_body.clone(),
);

Some(term.lambda(func_uplc_name).apply(const_body))
}
}
}
air::FunctionVariants::Cyclic(contained_functions) => {
let mut cyclic_functions = vec![];
Expand Down
93 changes: 92 additions & 1 deletion crates/aiken-project/src/tests/gen_uplc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3032,7 +3032,9 @@ fn acceptance_test_29_union_pair() {
inner: Pairs<key, value>,
}
const empty_list: AssocList<key, value> = AssocList { inner: [] }
const empty_list: AssocList<key, value> = {fn(k: key, v: value){
[(k,v)]
}(1, 2)}
pub fn from_list(xs: Pairs<key, value>) -> AssocList<key, value> {
AssocList { inner: do_from_list(xs) }
Expand Down Expand Up @@ -6051,6 +6053,95 @@ fn bls12_381_elements_from_data_conversion() {
)
}

#[test]
fn bls12_381_elements_constant_hoisting() {
let src = r#"
pub const generator_g1: G1Element =
#<Bls12_381, G1>"97f1d3a73197d7942695638c4fa9ac0fc3688c4f9774b905a14e3a3f171bac586c55e83ff97a1aeffb3af00adb22c6bb"
pub const other_generator_g1: G1Element =
#<Bls12_381, G1>"b28cb29bc282be68df977b35eb9d8e98b3a0a3fc7c372990bddc50419ca86693e491755338fed4fb42231a7c081252ce"
type Mew{
One(G1Element)
Two(G1Element)
}
type Choo{
Foo(Mew)
Bar
}
test thing() {
let x = Foo(One(other_generator_g1))
when x is {
Foo(y) -> {
when y is {
One(other_g1) -> {
let g1 = generator_g1
g1 != other_g1
}
Two(other_g) -> {
let ga = generator_g1
ga == other_g && True
}
}
}
Bar -> False
}
}
"#;

let bytes = vec![
0xb2, 0x8c, 0xb2, 0x9b, 0xc2, 0x82, 0xbe, 0x68, 0xdf, 0x97, 0x7b, 0x35, 0xeb, 0x9d, 0x8e,
0x98, 0xb3, 0xa0, 0xa3, 0xfc, 0x7c, 0x37, 0x29, 0x90, 0xbd, 0xdc, 0x50, 0x41, 0x9c, 0xa8,
0x66, 0x93, 0xe4, 0x91, 0x75, 0x53, 0x38, 0xfe, 0xd4, 0xfb, 0x42, 0x23, 0x1a, 0x7c, 0x08,
0x12, 0x52, 0xce,
];

let g1 = Term::Constant(
Constant::Bls12_381G1Element(blst::blst_p1::uncompress(&bytes).unwrap().into()).into(),
);

let constant = Term::Constant(
Constant::Data(Data::constr(
0,
vec![
Data::bytestring(bytes),
Data::bytestring(vec![
0xb9, 0x21, 0x5e, 0x5b, 0xc4, 0x81, 0xba, 0x65, 0x52, 0x38, 0x4c, 0x89, 0xc2,
0x3d, 0x45, 0xbd, 0x65, 0x0b, 0x69, 0x46, 0x28, 0x68, 0x24, 0x8b, 0xfb, 0xb8,
0x3a, 0xee, 0x70, 0x60, 0x57, 0x94, 0x04, 0xdb, 0xa4, 0x1c, 0x78, 0x1d, 0xec,
0x7c, 0x2b, 0xec, 0x5f, 0xcc, 0xec, 0x06, 0x84, 0x2e, 0x0e, 0x66, 0xad, 0x6d,
0x86, 0xc7, 0xc7, 0x6c, 0x46, 0x8a, 0x32, 0xc9, 0xc0, 0x08, 0x0e, 0xea, 0x02,
0x19, 0xd0, 0x95, 0x3b, 0x44, 0xb1, 0xc4, 0xf5, 0x60, 0x5a, 0xfb, 0x1e, 0x5a,
0x31, 0x93, 0x26, 0x4f, 0xf7, 0x30, 0x22, 0x2e, 0x94, 0xf5, 0x52, 0x07, 0x62,
0x82, 0x35, 0xf3, 0xb4, 0x23,
]),
],
))
.into(),
);

assert_uplc(
src,
Term::bls12_381_g1_equal()
.apply(Term::bls12_381_g1_uncompress().apply(
Term::un_b_data().apply(
Term::head_list().apply(
Term::snd_pair().apply(Term::unconstr_data().apply(constant.clone())),
),
),
))
.apply(g1),
false,
true,
)
}

#[test]
fn qualified_prelude_functions() {
let src = r#"
Expand Down
2 changes: 1 addition & 1 deletion crates/uplc/src/optimize.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ pub mod shrinker;
pub fn aiken_optimize_and_intern(program: Program<Name>) -> Program<Name> {
program
.inline_constr_ops()
.bls381_compressor()
.bls381_compressor_hoister()
.builtin_force_reducer()
.lambda_reducer()
.inline_reducer()
Expand Down

0 comments on commit 3246362

Please sign in to comment.