feat: Add aztec selectors for event structs (#2983)

sirasistant · web-flow · commit 982380e54bb4 · 2023-10-04T20:17:14.000Z
diff --git 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<StructId> {
+    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>) -> 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
@@ -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::<Vec<_>>(),
         );
diff --git 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);
