From 961b0c27f4a6dfd36bd33dd9e548ef259dfb78aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=81lvaro=20Rodr=C3=ADguez?= Date: Wed, 4 Oct 2023 22:17:14 +0200 Subject: [PATCH] feat: Add aztec selectors for event structs (#2983) --- .../src/hir/def_map/aztec_library.rs | 258 ++++++++++++++++-- .../noirc_frontend/src/hir/def_map/mod.rs | 3 + compiler/noirc_frontend/src/node_interner.rs | 13 + 3 files changed, 244 insertions(+), 30 deletions(-) diff --git a/compiler/noirc_frontend/src/hir/def_map/aztec_library.rs b/compiler/noirc_frontend/src/hir/def_map/aztec_library.rs index 69adb864317..384bb65deff 100644 --- a/compiler/noirc_frontend/src/hir/def_map/aztec_library.rs +++ b/compiler/noirc_frontend/src/hir/def_map/aztec_library.rs @@ -1,8 +1,12 @@ use acvm::FieldElement; +use iter_extended::vecmap; use noirc_errors::Span; use crate::graph::CrateId; use crate::hir::def_collector::errors::DefCollectorErrorKind; +use crate::hir_def::expr::{HirExpression, HirLiteral}; +use crate::hir_def::stmt::HirStatement; +use crate::node_interner::{NodeInterner, StructId}; use crate::token::SecondaryAttribute; use crate::{ hir::Context, BlockExpression, CallExpression, CastExpression, Distinctness, Expression, @@ -11,9 +15,14 @@ use crate::{ ParsedModule, Path, PathKind, Pattern, Statement, UnresolvedType, UnresolvedTypeData, Visibility, }; -use crate::{PrefixExpression, UnaryOp}; +use crate::{ + FunctionDefinition, NoirStruct, PrefixExpression, Shared, Signedness, StructType, Type, + TypeBinding, TypeImpl, TypeVariableKind, UnaryOp, +}; use fm::FileId; +use super::ModuleDefId; + // // Helper macros for creating noir ast nodes // @@ -163,7 +172,7 @@ pub(crate) fn transform( for submodule in ast.submodules.iter_mut().filter(|submodule| submodule.is_contract) { let storage_defined = check_for_storage_definition(&submodule.contents); - if transform_module(&mut submodule.contents.functions, storage_defined) { + if transform_module(&mut submodule.contents, storage_defined) { match check_for_aztec_dependency(crate_id, context) { Ok(()) => include_relevant_imports(&mut submodule.contents), Err(file_id) => { @@ -175,6 +184,15 @@ pub(crate) fn transform( Ok(ast) } +// +// Transform Hir Nodes for Aztec +// + +/// Completes the Hir with data gathered from type resolution +pub(crate) fn transform_hir(crate_id: &CrateId, context: &mut Context) { + transform_events(crate_id, context); +} + /// Includes an import to the aztec library if it has not been included yet fn include_relevant_imports(ast: &mut ParsedModule) { // Create the aztec import path using the assumed chained_dep! macro @@ -206,34 +224,45 @@ fn check_for_storage_definition(module: &ParsedModule) -> bool { module.types.iter().any(|function| function.name.0.contents == "Storage") } -/// Determines if the function is annotated with `aztec(private)` or `aztec(public)` -/// If it is, it calls the `transform` function which will perform the required transformations. -/// Returns true if an annotated function is found, false otherwise -fn transform_module(functions: &mut [NoirFunction], storage_defined: bool) -> bool { - let mut has_annotated_functions = false; - for func in functions.iter_mut() { +/// Checks if an attribute is a custom attribute with a specific name +fn is_custom_attribute(attr: &SecondaryAttribute, attribute_name: &str) -> bool { + if let SecondaryAttribute::Custom(custom_attr) = attr { + custom_attr.as_str() == attribute_name + } else { + false + } +} + +/// Determines if ast nodes are annotated with aztec attributes. +/// For annotated functions it calls the `transform` function which will perform the required transformations. +/// Returns true if an annotated node is found, false otherwise +fn transform_module(module: &mut ParsedModule, storage_defined: bool) -> bool { + let mut has_transformed_module = false; + + for structure in module.types.iter_mut() { + if structure.attributes.iter().any(|attr| matches!(attr, SecondaryAttribute::Event)) { + module.impls.push(generate_selector_impl(structure)); + has_transformed_module = true; + } + } + + for func in module.functions.iter_mut() { for secondary_attribute in func.def.attributes.secondary.clone() { - if let SecondaryAttribute::Custom(custom_attribute) = secondary_attribute { - match custom_attribute.as_str() { - "aztec(private)" => { - transform_function("Private", func, storage_defined); - has_annotated_functions = true; - } - "aztec(public)" => { - transform_function("Public", func, storage_defined); - has_annotated_functions = true; - } - _ => continue, - } + if is_custom_attribute(&secondary_attribute, "aztec(private)") { + transform_function("Private", func, storage_defined); + has_transformed_module = true; + } else if is_custom_attribute(&secondary_attribute, "aztec(public)") { + transform_function("Public", func, storage_defined); + has_transformed_module = true; } } // Add the storage struct to the beginning of the function if it is unconstrained in an aztec contract if storage_defined && func.def.is_unconstrained { transform_unconstrained(func); - has_annotated_functions = true; + has_transformed_module = true; } } - has_annotated_functions + has_transformed_module } /// If it does, it will insert the following things: @@ -293,6 +322,141 @@ fn transform_unconstrained(func: &mut NoirFunction) { func.def.body.0.insert(0, abstract_storage("Unconstrained", true)); } +fn collect_crate_structs(crate_id: &CrateId, context: &Context) -> Vec { + context + .def_map(crate_id) + .expect("ICE: Missing crate in def_map") + .modules() + .iter() + .flat_map(|(_, module)| { + module.type_definitions().filter_map(|typ| { + if let ModuleDefId::TypeId(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, interner: &mut NodeInterner) { + let selector_id = + interner.lookup_method(struct_id, "selector").expect("Selector method not found"); + let selector_function = interner.function(&selector_id); + + let compute_selector_statement = interner.statement( + selector_function + .block(interner) + .statements() + .first() + .expect("Compute selector statement not found"), + ); + + let compute_selector_expression = match compute_selector_statement { + HirStatement::Expression(expression_id) => match interner.expression(&expression_id) { + HirExpression::Call(hir_call_expression) => Some(hir_call_expression), + _ => None, + }, + _ => None, + } + .expect("Compute selector statement is not a call expression"); + + let first_arg_id = compute_selector_expression + .arguments + .first() + .expect("Missing argument for compute selector"); + + match interner.expression(first_arg_id) { + HirExpression::Literal(HirLiteral::Str(signature)) + if signature == SIGNATURE_PLACEHOLDER => + { + let selector_literal_id = first_arg_id; + let compute_selector_call_id = compute_selector_expression.func; + + let structure = interner.get_struct(struct_id); + let signature = event_signature(&structure.borrow()); + interner.update_expression(*selector_literal_id, |expr| { + *expr = HirExpression::Literal(HirLiteral::Str(signature.clone())); + }); + + // Also update the type! It might have a different length now than the placeholder. + interner.push_expr_type( + selector_literal_id, + Type::String(Box::new(Type::Constant(signature.len() as u64))), + ); + interner.push_expr_type( + &compute_selector_call_id, + Type::Function( + vec![Type::String(Box::new(Type::TypeVariable( + Shared::new(TypeBinding::Bound(Type::Constant(signature.len() as u64))), + TypeVariableKind::Normal, + )))], + Box::new(Type::FieldElement), + Box::new(Type::Unit), + ), + ); + } + _ => unreachable!("Signature placeholder literal does not match"), + } +} + +fn transform_events(crate_id: &CrateId, context: &mut Context) { + for struct_id in collect_crate_structs(crate_id, context) { + let attributes = context.def_interner.struct_attributes(&struct_id); + if attributes.iter().any(|attr| matches!(attr, SecondaryAttribute::Event)) { + transform_event(struct_id, &mut context.def_interner); + } + } +} + +const SIGNATURE_PLACEHOLDER: &str = "SIGNATURE_PLACEHOLDER"; + +/// Generates the impl for an event selector +/// +/// Inserts the following code: +/// ```noir +/// impl SomeStruct { +/// fn selector() -> Field { +/// aztec::oracle::compute_selector::compute_selector("SIGNATURE_PLACEHOLDER") +/// } +/// } +/// ``` +/// +/// This allows developers to emit events without having to write the signature of the event every time they emit it. +/// The signature cannot be known at this point since types are not resolved yet, so we use a signature placeholder. +/// It'll get resolved after by transforming the HIR. +fn generate_selector_impl(structure: &mut NoirStruct) -> TypeImpl { + let struct_type = make_type(UnresolvedTypeData::Named(path(structure.name.clone()), vec![])); + + let selector_fun_body = BlockExpression(vec![Statement::Expression(call( + variable_path(chained_path!("aztec", "oracle", "compute_selector", "compute_selector")), + vec![expression(ExpressionKind::Literal(Literal::Str(SIGNATURE_PLACEHOLDER.to_string())))], + ))]); + + let mut selector_fn_def = FunctionDefinition::normal( + &ident("selector"), + &vec![], + &[], + &selector_fun_body, + &[], + &FunctionReturnType::Ty(make_type(UnresolvedTypeData::FieldElement)), + ); + + selector_fn_def.is_public = true; + + // Seems to be necessary on contract modules + selector_fn_def.return_visibility = Visibility::Public; + + TypeImpl { + object_type: struct_type, + type_span: structure.span, + generics: vec![], + methods: vec![NoirFunction::normal(selector_fn_def)], + } +} + /// Helper function that returns what the private context would look like in the ast /// This should make it available to be consumed within aztec private annotated functions. /// @@ -537,9 +701,9 @@ fn make_return_push_array(push_value: Expression) -> Statement { /// `context.return_values.push_array({push_value}.serialize())` fn make_struct_return_type(expression: Expression) -> Statement { let serialized_call = method_call( - expression.clone(), // variable - "serialize", // method name - vec![], // args + expression, // variable + "serialize", // method name + vec![], // args ); make_return_push_array(serialized_call) } @@ -561,7 +725,7 @@ fn make_array_return_type(expression: Expression) -> Statement { vec![inner_cast_expression], )); - create_loop_over(expression.clone(), vec![assignment]) + create_loop_over(expression, vec![assignment]) } /// Castable return type @@ -572,7 +736,7 @@ fn make_array_return_type(expression: Expression) -> Statement { /// ``` fn make_castable_return_type(expression: Expression) -> Statement { // Cast these types to a field before pushing - let cast_expression = cast(expression.clone(), UnresolvedTypeData::FieldElement); + let cast_expression = cast(expression, UnresolvedTypeData::FieldElement); make_return_push(cast_expression) } @@ -657,9 +821,9 @@ fn create_loop_over(var: Expression, loop_body: Vec) -> Statement { // `array.len()` let end_range_expression = method_call( - var.clone(), // variable - "len", // method name - vec![], // args + var, // variable + "len", // method name + vec![], // args ); // What will be looped over @@ -721,3 +885,37 @@ fn add_cast_to_hasher(identifier: &Ident) -> Statement { vec![cast_operation], // args )) } + +/// Computes the aztec signature for a resolved type. +fn signature_of_type(typ: &Type) -> String { + match typ { + Type::Integer(Signedness::Signed, bit_size) => format!("i{}", bit_size), + Type::Integer(Signedness::Unsigned, bit_size) => format!("u{}", bit_size), + Type::FieldElement => "Field".to_owned(), + Type::Bool => "bool".to_owned(), + Type::Array(len, typ) => { + if let Type::Constant(len) = **len { + format!("[{};{len}]", signature_of_type(typ)) + } else { + unimplemented!("Cannot generate signature for array with length type {:?}", typ) + } + } + Type::Struct(def, args) => { + let fields = def.borrow().get_fields(args); + let fields = vecmap(fields, |(_, typ)| signature_of_type(&typ)); + format!("({})", fields.join(",")) + } + Type::Tuple(types) => { + let fields = vecmap(types, signature_of_type); + format!("({})", fields.join(",")) + } + _ => unimplemented!("Cannot generate signature for type {:?}", typ), + } +} + +/// Computes the signature for a resolved event type. +/// It has the form 'EventName(Field,(Field),[u8;2])' +fn event_signature(event: &StructType) -> String { + let fields = vecmap(event.get_fields(&[]), |(_, typ)| signature_of_type(&typ)); + format!("{}({})", event.name.0.contents, fields.join(",")) +} diff --git a/compiler/noirc_frontend/src/hir/def_map/mod.rs b/compiler/noirc_frontend/src/hir/def_map/mod.rs index ba2f4fdaae3..3ea75c57865 100644 --- a/compiler/noirc_frontend/src/hir/def_map/mod.rs +++ b/compiler/noirc_frontend/src/hir/def_map/mod.rs @@ -110,6 +110,9 @@ impl CrateDefMap { // Now we want to populate the CrateDefMap using the DefCollector errors.extend(DefCollector::collect(def_map, context, ast, root_file_id)); + #[cfg(feature = "aztec")] + aztec_library::transform_hir(&crate_id, context); + errors.extend( parsing_errors.iter().map(|e| (e.clone().into(), root_file_id)).collect::>(), ); diff --git a/compiler/noirc_frontend/src/node_interner.rs b/compiler/noirc_frontend/src/node_interner.rs index 84752553585..94220d39402 100644 --- a/compiler/noirc_frontend/src/node_interner.rs +++ b/compiler/noirc_frontend/src/node_interner.rs @@ -512,6 +512,19 @@ impl NodeInterner { } } + /// Updates the interned expression corresponding to `expr_id` + pub fn update_expression(&mut self, expr_id: ExprId, f: impl FnOnce(&mut HirExpression)) { + let def = + self.nodes.get_mut(expr_id.0).expect("ice: all expression ids should have definitions"); + + match def { + Node::Expression(expr) => f(expr), + _ => { + panic!("ice: all expression ids should correspond to a expression in the interner") + } + } + } + /// Store the type for an interned Identifier pub fn push_definition_type(&mut self, definition_id: DefinitionId, typ: Type) { self.id_to_type.insert(definition_id.into(), typ);