From cc92cbc10ac8483aae77f5c87c5ace4e0c011799 Mon Sep 17 00:00:00 2001 From: AztecBot Date: Mon, 5 Feb 2024 00:40:23 +0000 Subject: [PATCH] feat(aztec-nr): initial work for aztec public vm macro (#4400)... BEGIN_COMMIT_OVERRIDE\nfeat(aztec-nr): initial work for aztec public vm macro (#4400) chore: surpress chained macro warning (#4396) feat!: init storage macro (#4200)\nEND_COMMIT_OVERRIDE --- aztec_macros/src/lib.rs | 408 ++++++++++++++++++- compiler/noirc_frontend/src/node_interner.rs | 30 +- 2 files changed, 430 insertions(+), 8 deletions(-) diff --git a/aztec_macros/src/lib.rs b/aztec_macros/src/lib.rs index 1dbe7631388..228664c83fb 100644 --- a/aztec_macros/src/lib.rs +++ b/aztec_macros/src/lib.rs @@ -1,5 +1,7 @@ -use iter_extended::vecmap; +use std::borrow::{Borrow, BorrowMut}; +use std::vec; +use iter_extended::vecmap; use noirc_frontend::macros_api::FieldElement; use noirc_frontend::macros_api::{ BlockExpression, CallExpression, CastExpression, Distinctness, Expression, ExpressionKind, @@ -13,6 +15,8 @@ use noirc_frontend::macros_api::{ use noirc_frontend::macros_api::{CrateId, FileId}; use noirc_frontend::macros_api::{MacroError, MacroProcessor}; use noirc_frontend::macros_api::{ModuleDefId, NodeInterner, SortedModule, StructId}; +use noirc_frontend::node_interner::{TraitId, TraitImplKind}; +use noirc_frontend::Lambda; pub struct AztecMacro; @@ -45,6 +49,8 @@ pub enum AztecMacroError { ContractHasTooManyFunctions { span: Span }, ContractConstructorMissing { span: Span }, UnsupportedFunctionArgumentType { span: Span, typ: UnresolvedTypeData }, + UnsupportedStorageType { span: Option, typ: UnresolvedTypeData }, + CouldNotAssignStorageSlots { secondary_message: Option }, EventError { span: Span, message: String }, } @@ -76,6 +82,16 @@ impl From for MacroError { secondary_message: None, span: Some(span), }, + AztecMacroError::UnsupportedStorageType { span, typ } => MacroError { + primary_message: format!("Provided storage type `{typ:?}` is not directly supported in Aztec. Please provide a custom storage implementation"), + secondary_message: None, + span, + }, + AztecMacroError::CouldNotAssignStorageSlots { secondary_message } => MacroError { + primary_message: "Could not assign storage slots, please provide a custom storage implementation".to_string(), + secondary_message, + span: None, + }, AztecMacroError::EventError { span, message } => MacroError { primary_message: message, secondary_message: None, @@ -166,7 +182,28 @@ fn member_access(lhs: &str, rhs: &str) -> Expression { }))) } +fn return_type(path: Path) -> FunctionReturnType { + let ty = make_type(UnresolvedTypeData::Named(path, vec![])); + FunctionReturnType::Ty(ty) +} + +fn lambda(parameters: Vec<(Pattern, UnresolvedType)>, body: Expression) -> Expression { + expression(ExpressionKind::Lambda(Box::new(Lambda { + parameters, + return_type: UnresolvedType { + typ: UnresolvedTypeData::Unspecified, + span: Some(Span::default()), + }, + body, + }))) +} + macro_rules! chained_path { + ( $base:expr ) => { + { + ident_path($base) + } + }; ( $base:expr $(, $tail:expr)* ) => { { let mut base_path = ident_path($base); @@ -196,7 +233,7 @@ fn cast(lhs: Expression, ty: UnresolvedTypeData) -> Expression { } fn make_type(typ: UnresolvedTypeData) -> UnresolvedType { - UnresolvedType { typ, span: None } + UnresolvedType { typ, span: Some(Span::default()) } } fn index_array(array: Ident, index: &str) -> Expression { @@ -251,7 +288,8 @@ fn transform_hir( crate_id: &CrateId, context: &mut HirContext, ) -> Result<(), (AztecMacroError, FileId)> { - transform_events(crate_id, context) + transform_events(crate_id, context)?; + assign_storage_slots(crate_id, context) } /// Includes an import to the aztec library if it has not been included yet @@ -288,6 +326,16 @@ fn check_for_storage_definition(module: &SortedModule) -> bool { module.types.iter().any(|r#struct| r#struct.name.0.contents == "Storage") } +// Check to see if the user has defined a storage struct +fn check_for_storage_implementation(module: &SortedModule) -> bool { + module.impls.iter().any(|r#impl| match &r#impl.object_type.typ { + UnresolvedTypeData::Named(path, _) => { + path.segments.last().is_some_and(|segment| segment.0.contents == "Storage") + } + _ => false, + }) +} + // Check if "compute_note_hash_and_nullifier(AztecAddress,Field,Field,[Field; N]) -> [Field; 4]" is defined fn check_for_compute_note_hash_and_nullifier_definition(module: &SortedModule) -> bool { module.functions.iter().any(|func| { @@ -343,9 +391,15 @@ fn transform_module( // Check for a user defined storage struct let storage_defined = check_for_storage_definition(module); + let storage_implemented = check_for_storage_implementation(module); + + let crate_graph = &context.crate_graph[crate_id]; + + if storage_defined && !storage_implemented { + generate_storage_implementation(module).map_err(|err| (err, crate_graph.root_file_id))?; + } if storage_defined && !check_for_compute_note_hash_and_nullifier_definition(module) { - let crate_graph = &context.crate_graph[crate_id]; return Err(( AztecMacroError::ComputeNoteHashAndNullifierNotFound { span: Span::default() }, crate_graph.root_file_id, @@ -370,6 +424,10 @@ fn transform_module( transform_function("Public", func, storage_defined) .map_err(|err| (err, crate_graph.root_file_id))?; has_transformed_module = true; + } else if is_custom_attribute(&secondary_attribute, "aztec(public-vm)") { + transform_vm_function(func, storage_defined) + .map_err(|err| (err, crate_graph.root_file_id))?; + has_transformed_module = true; } } // Add the storage struct to the beginning of the function if it is unconstrained in an aztec contract @@ -403,6 +461,134 @@ fn transform_module( Ok(has_transformed_module) } +/// Auxiliary function to generate the storage constructor for a given field, using +/// the Storage definition as a reference. Supports nesting. +fn generate_storage_field_constructor( + (type_ident, unresolved_type): &(Ident, UnresolvedType), + slot: Expression, +) -> Result { + let typ = &unresolved_type.typ; + match typ { + UnresolvedTypeData::Named(path, generics) => { + let mut new_path = path.clone().to_owned(); + new_path.segments.push(ident("new")); + match path.segments.last().unwrap().0.contents.as_str() { + "Map" => Ok(call( + variable_path(new_path), + vec![ + variable("context"), + slot, + lambda( + vec![ + ( + pattern("context"), + make_type(UnresolvedTypeData::Named( + chained_path!("aztec", "context", "Context"), + vec![], + )), + ), + ( + Pattern::Identifier(ident("slot")), + make_type(UnresolvedTypeData::FieldElement), + ), + ], + generate_storage_field_constructor( + &(type_ident.clone(), generics.iter().last().unwrap().clone()), + variable("slot"), + )?, + ), + ], + )), + _ => Ok(call(variable_path(new_path), vec![variable("context"), slot])), + } + } + _ => Err(AztecMacroError::UnsupportedStorageType { + typ: typ.clone(), + span: Some(type_ident.span()), + }), + } +} + +// Generates the Storage implementation block from the Storage struct definition if it does not exist +/// From: +/// +/// struct Storage { +/// a_map: Map>, +/// a_nested_map: Map>>, +/// a_field: SomeStoragePrimitive, +/// } +/// +/// To: +/// +/// impl Storage { +/// fn init(context: Context) -> Self { +/// Storage { +/// a_map: Map::new(context, 0, |context, slot| { +/// SomeStoragePrimitive::new(context, slot) +/// }), +/// a_nested_map: Map::new(context, 0, |context, slot| { +/// Map::new(context, slot, |context, slot| { +/// SomeStoragePrimitive::new(context, slot) +/// }) +/// }), +/// a_field: SomeStoragePrimitive::new(context, 0), +/// } +/// } +/// } +/// +/// Storage slots are generated as 0 and will be populated using the information from the HIR +/// at a later stage. +fn generate_storage_implementation(module: &mut SortedModule) -> Result<(), AztecMacroError> { + let definition = + module.types.iter().find(|r#struct| r#struct.name.0.contents == "Storage").unwrap(); + + let slot_zero = expression(ExpressionKind::Literal(Literal::Integer( + FieldElement::from(i128::from(0)), + false, + ))); + + let field_constructors = definition + .fields + .iter() + .flat_map(|field| { + generate_storage_field_constructor(field, slot_zero.clone()) + .map(|expression| (field.0.clone(), expression)) + }) + .collect(); + + let storage_constructor_statement = make_statement(StatementKind::Expression(expression( + ExpressionKind::constructor((chained_path!("Storage"), field_constructors)), + ))); + + let init = NoirFunction::normal(FunctionDefinition::normal( + &ident("init"), + &vec![], + &[( + ident("context"), + make_type(UnresolvedTypeData::Named( + chained_path!("aztec", "context", "Context"), + vec![], + )), + )], + &BlockExpression(vec![storage_constructor_statement]), + &[], + &return_type(chained_path!("Self")), + )); + + let storage_impl = TypeImpl { + object_type: UnresolvedType { + typ: UnresolvedTypeData::Named(chained_path!("Storage"), vec![]), + span: Some(Span::default()), + }, + type_span: Span::default(), + generics: vec![], + methods: vec![init], + }; + module.impls.push(storage_impl); + + Ok(()) +} + /// If it does, it will insert the following things: /// - A new Input that is provided for a kernel app circuit, named: {Public/Private}ContextInputs /// - Hashes all of the function input variables @@ -454,6 +640,20 @@ fn transform_function( Ok(()) } +/// Transform a function to work with AVM bytecode +fn transform_vm_function( + func: &mut NoirFunction, + _storage_defined: bool, +) -> Result<(), AztecMacroError> { + // We want the function to be seen as a public function + func.def.is_open = true; + + // NOTE: the line below is a temporary hack to trigger external transpilation tools + // It will be removed once the transpiler is integrated into the Noir compiler + func.def.name.0.contents = format!("avm_{}", func.def.name.0.contents); + Ok(()) +} + /// Transform Unconstrained /// /// Inserts the following code at the beginning of an unconstrained function @@ -484,6 +684,23 @@ fn collect_crate_structs(crate_id: &CrateId, context: &HirContext) -> Vec Vec { + let crates = context.crates(); + crates + .flat_map(|crate_id| context.def_map(&crate_id).map(|def_map| def_map.modules())) + .flatten() + .flat_map(|(_, module)| { + module.type_definitions().filter_map(|typ| { + if let ModuleDefId::TraitId(struct_id) = typ { + Some(struct_id) + } else { + None + } + }) + }) + .collect() +} + /// Substitutes the signature literal that was introduced in the selector method previously with the actual signature. fn transform_event( struct_id: StructId, @@ -576,6 +793,185 @@ fn transform_events( Ok(()) } +/// Obtains the serialized length of a type that implements the Serialize trait. +fn get_serialized_length( + traits: &[TraitId], + typ: &Type, + interner: &NodeInterner, +) -> Result { + let (struct_name, maybe_stored_in_state) = match typ { + Type::Struct(struct_type, generics) => { + Ok((struct_type.borrow().name.0.contents.clone(), generics.get(0))) + } + _ => Err(AztecMacroError::CouldNotAssignStorageSlots { + secondary_message: Some("State storage variable must be a struct".to_string()), + }), + }?; + let stored_in_state = + maybe_stored_in_state.ok_or(AztecMacroError::CouldNotAssignStorageSlots { + secondary_message: Some("State storage variable must be generic".to_string()), + })?; + + let is_note = traits.iter().any(|&trait_id| { + let r#trait = interner.get_trait(trait_id); + r#trait.name.0.contents == "NoteInterface" + && !interner.lookup_all_trait_implementations(stored_in_state, trait_id).is_empty() + }); + + // Maps and (private) Notes always occupy a single slot. Someone could store a Note in PublicState for whatever reason though. + if struct_name == "Map" || (is_note && struct_name != "PublicState") { + return Ok(1); + } + + let serialized_trait_impl_kind = traits + .iter() + .filter_map(|&trait_id| { + let r#trait = interner.get_trait(trait_id); + if r#trait.borrow().name.0.contents == "Serialize" + && r#trait.borrow().generics.len() == 1 + { + interner + .lookup_all_trait_implementations(stored_in_state, trait_id) + .into_iter() + .next() + } else { + None + } + }) + .next() + .ok_or(AztecMacroError::CouldNotAssignStorageSlots { + secondary_message: Some("Stored data must implement Serialize trait".to_string()), + })?; + + let serialized_trait_impl_id = match serialized_trait_impl_kind { + TraitImplKind::Normal(trait_impl_id) => Ok(trait_impl_id), + _ => Err(AztecMacroError::CouldNotAssignStorageSlots { secondary_message: None }), + }?; + + let serialized_trait_impl_shared = interner.get_trait_implementation(*serialized_trait_impl_id); + let serialized_trait_impl = serialized_trait_impl_shared.borrow(); + + match serialized_trait_impl.trait_generics.get(0).unwrap() { + Type::Constant(value) => Ok(*value), + _ => Err(AztecMacroError::CouldNotAssignStorageSlots { secondary_message: None }), + } +} + +/// Assigns storage slots to the storage struct fields based on the serialized length of the types. This automatic assignment +/// will only trigger if the assigned storage slot is invalid (0 as generated by generate_storage_implementation) +fn assign_storage_slots( + crate_id: &CrateId, + context: &mut HirContext, +) -> Result<(), (AztecMacroError, FileId)> { + let traits: Vec<_> = collect_traits(context); + for struct_id in collect_crate_structs(crate_id, context) { + let interner: &mut NodeInterner = context.def_interner.borrow_mut(); + let r#struct = interner.get_struct(struct_id); + let file_id = r#struct.borrow().location.file; + if r#struct.borrow().name.0.contents == "Storage" && r#struct.borrow().id.krate().is_root() + { + let init_id = interner + .lookup_method( + &Type::Struct(interner.get_struct(struct_id), vec![]), + struct_id, + "init", + false, + ) + .ok_or(( + AztecMacroError::CouldNotAssignStorageSlots { + secondary_message: Some( + "Storage struct must have an init function".to_string(), + ), + }, + file_id, + ))?; + let init_function = interner.function(&init_id).block(interner); + let init_function_statement_id = init_function.statements().first().ok_or(( + AztecMacroError::CouldNotAssignStorageSlots { + secondary_message: Some("Init storage statement not found".to_string()), + }, + file_id, + ))?; + let storage_constructor_statement = interner.statement(init_function_statement_id); + + let storage_constructor_expression = match storage_constructor_statement { + HirStatement::Expression(expression_id) => { + match interner.expression(&expression_id) { + HirExpression::Constructor(hir_constructor_expression) => { + Ok(hir_constructor_expression) + } + _ => Err((AztecMacroError::CouldNotAssignStorageSlots { + secondary_message: Some( + "Storage constructor statement must be a constructor expression" + .to_string(), + ), + }, file_id)) + } + } + _ => Err(( + AztecMacroError::CouldNotAssignStorageSlots { + secondary_message: Some( + "Storage constructor statement must be an expression".to_string(), + ), + }, + file_id, + )), + }?; + + let mut storage_slot: u64 = 1; + for (index, (_, expr_id)) in storage_constructor_expression.fields.iter().enumerate() { + let fields = r#struct.borrow().get_fields(&[]); + let (_, field_type) = fields.get(index).unwrap(); + let new_call_expression = match interner.expression(expr_id) { + HirExpression::Call(hir_call_expression) => Ok(hir_call_expression), + _ => Err(( + AztecMacroError::CouldNotAssignStorageSlots { + secondary_message: Some( + "Storage field initialization expression is not a call expression" + .to_string(), + ), + }, + file_id, + )), + }?; + + let slot_arg_expression = interner.expression(&new_call_expression.arguments[1]); + + let current_storage_slot = match slot_arg_expression { + HirExpression::Literal(HirLiteral::Integer(slot, _)) => { + Ok(slot.borrow().to_u128()) + } + _ => Err(( + AztecMacroError::CouldNotAssignStorageSlots { + secondary_message: Some( + "Storage slot argument expression must be a literal integer" + .to_string(), + ), + }, + file_id, + )), + }?; + + if current_storage_slot != 0 { + continue; + } + + let type_serialized_len = get_serialized_length(&traits, field_type, interner) + .map_err(|err| (err, file_id))?; + interner.update_expression(new_call_expression.arguments[1], |expr| { + *expr = HirExpression::Literal(HirLiteral::Integer( + FieldElement::from(u128::from(storage_slot)), + false, + )); + }); + + storage_slot += type_serialized_len; + } + } + } + Ok(()) +} + const SIGNATURE_PLACEHOLDER: &str = "SIGNATURE_PLACEHOLDER"; /// Generates the impl for an event selector @@ -969,9 +1365,7 @@ fn make_castable_return_type(expression: Expression) -> Statement { /// } fn create_return_type(ty: &str) -> FunctionReturnType { let return_path = chained_path!("aztec", "abi", ty); - - let ty = make_type(UnresolvedTypeData::Named(return_path, vec![])); - FunctionReturnType::Ty(ty) + return_type(return_path) } /// Create Context Finish diff --git a/compiler/noirc_frontend/src/node_interner.rs b/compiler/noirc_frontend/src/node_interner.rs index b856b54f6ca..11ef12ef83e 100644 --- a/compiler/noirc_frontend/src/node_interner.rs +++ b/compiler/noirc_frontend/src/node_interner.rs @@ -1,4 +1,5 @@ use std::collections::HashMap; +use std::ops::Deref; use arena::{Arena, Index}; use fm::FileId; @@ -1042,6 +1043,34 @@ impl NodeInterner { Ok(impl_kind) } + /// Given a `ObjectType: TraitId` pair, find all implementations without taking constraints into account or + /// applying any type bindings. Useful to look for a specific trait in a type that is used in a macro. + pub fn lookup_all_trait_implementations( + &self, + object_type: &Type, + trait_id: TraitId, + ) -> Vec<&TraitImplKind> { + let trait_impl = self.trait_implementation_map.get(&trait_id); + + trait_impl + .map(|trait_impl| { + trait_impl + .iter() + .filter_map(|(typ, impl_kind)| match &typ { + Type::Forall(_, typ) => { + if typ.deref() == object_type { + Some(impl_kind) + } else { + None + } + } + _ => None, + }) + .collect() + }) + .unwrap_or(vec![]) + } + /// Similar to `lookup_trait_implementation` but does not apply any type bindings on success. pub fn try_lookup_trait_implementation( &self, @@ -1263,7 +1292,6 @@ impl NodeInterner { force_type_check: bool, ) -> Option { let methods = self.struct_methods.get(&(id, method_name.to_owned())); - // If there is only one method, just return it immediately. // It will still be typechecked later. if !force_type_check {