chore: aztec-macros refactor

noir/noir-repo/aztec_macros/src/transforms/compute_note_hash_and_nullifier.rs

@@ -0,0 +1,195 @@
+use noirc_errors::{Location, Span};
+use noirc_frontend::{
+    graph::CrateId,
+    hir::{
+        def_collector::dc_crate::{UnresolvedFunctions, UnresolvedTraitImpl},
+        def_map::{LocalModuleId, ModuleId},
+    },
+    macros_api::{FileId, HirContext, MacroError},
+    node_interner::FuncId,
+    parse_program, FunctionReturnType, NoirFunction, UnresolvedTypeData,
+};
+
+use crate::utils::hir_utils::fetch_struct_trait_impls;
+
+// Check if "compute_note_hash_and_nullifier(AztecAddress,Field,Field,Field,[Field; N]) -> [Field; 4]" is defined
+fn check_for_compute_note_hash_and_nullifier_definition(
+    functions_data: &[(LocalModuleId, FuncId, NoirFunction)],
+    module_id: LocalModuleId,
+) -> bool {
+    functions_data.iter().filter(|func_data| func_data.0 == module_id).any(|func_data| {
+        func_data.2.def.name.0.contents == "compute_note_hash_and_nullifier"
+                && func_data.2.def.parameters.len() == 5
+                && match &func_data.2.def.parameters[0].typ.typ {
+                    UnresolvedTypeData::Named(path, _, _) => path.segments.last().unwrap().0.contents == "AztecAddress",
+                    _ => false,
+                }
+                && func_data.2.def.parameters[1].typ.typ == UnresolvedTypeData::FieldElement
+                && func_data.2.def.parameters[2].typ.typ == UnresolvedTypeData::FieldElement
+                && func_data.2.def.parameters[3].typ.typ == UnresolvedTypeData::FieldElement
+                // checks if the 5th parameter is an array and the Box<UnresolvedType> in
+                // Array(Option<UnresolvedTypeExpression>, Box<UnresolvedType>) contains only fields
+                && match &func_data.2.def.parameters[4].typ.typ {
+                    UnresolvedTypeData::Array(_, inner_type) => {
+                        matches!(inner_type.typ, UnresolvedTypeData::FieldElement)
+                    },
+                    _ => false,
+                }
+                // We check the return type the same way as we did the 5th parameter
+                && match &func_data.2.def.return_type {
+                    FunctionReturnType::Default(_) => false,
+                    FunctionReturnType::Ty(unresolved_type) => {
+                        match &unresolved_type.typ {
+                            UnresolvedTypeData::Array(_, inner_type) => {
+                                matches!(inner_type.typ, UnresolvedTypeData::FieldElement)
+                            },
+                            _ => false,
+                        }
+                    }
+                }
+    })
+}
+
+pub fn inject_compute_note_hash_and_nullifier(
+    crate_id: &CrateId,
+    context: &mut HirContext,
+    unresolved_traits_impls: &[UnresolvedTraitImpl],
+    collected_functions: &mut [UnresolvedFunctions],
+) -> Result<(), (MacroError, FileId)> {
+    // We first fetch modules in this crate which correspond to contracts, along with their file id.
+    let contract_module_file_ids: Vec<(LocalModuleId, FileId)> = context
+        .def_map(crate_id)
+        .expect("ICE: Missing crate in def_map")
+        .modules()
+        .iter()
+        .filter(|(_, module)| module.is_contract)
+        .map(|(idx, module)| (LocalModuleId(idx), module.location.file))
+        .collect();
+
+    // If the current crate does not contain a contract module we simply skip it.
+    if contract_module_file_ids.is_empty() {
+        return Ok(());
+    } else if contract_module_file_ids.len() != 1 {
+        panic!("Found multiple contracts in the same crate");
+    }
+
+    let (module_id, file_id) = contract_module_file_ids[0];
+
+    // If compute_note_hash_and_nullifier is already defined by the user, we skip auto-generation in order to provide an
+    // escape hatch for this mechanism.
+    // TODO(#4647): improve this diagnosis and error messaging.
+    if collected_functions.iter().any(|coll_funcs_data| {
+        check_for_compute_note_hash_and_nullifier_definition(&coll_funcs_data.functions, module_id)
+    }) {
+        return Ok(());
+    }
+
+    // In order to implement compute_note_hash_and_nullifier, we need to know all of the different note types the
+    // contract might use. These are the types that implement the NoteInterface trait, which provides the
+    // get_note_type_id function.
+    let note_types = fetch_struct_trait_impls(context, unresolved_traits_impls, "NoteInterface");
+
+    // We can now generate a version of compute_note_hash_and_nullifier tailored for the contract in this crate.
+    let func = generate_compute_note_hash_and_nullifier(&note_types);
+
+    // And inject the newly created function into the contract.
+
+    // TODO(#4373): We don't have a reasonable location for the source code of this autogenerated function, so we simply
+    // pass an empty span. This function should not produce errors anyway so this should not matter.
+    let location = Location::new(Span::empty(0), file_id);
+
+    // These are the same things the ModCollector does when collecting functions: we push the function to the
+    // NodeInterner, declare it in the module (which checks for duplicate definitions), and finally add it to the list
+    // on collected but unresolved functions.
+
+    let func_id = context.def_interner.push_empty_fn();
+    context.def_interner.push_function(
+        func_id,
+        &func.def,
+        ModuleId { krate: *crate_id, local_id: module_id },
+        location,
+    );
+
+    context.def_map_mut(crate_id).unwrap()
+        .modules_mut()[module_id.0]
+        .declare_function(
+            func.name_ident().clone(), func_id
+        ).expect(
+            "Failed to declare the autogenerated compute_note_hash_and_nullifier function, likely due to a duplicate definition. See https://github.com/AztecProtocol/aztec-packages/issues/4647."
+        );
+
+    collected_functions
+        .iter_mut()
+        .find(|fns| fns.file_id == file_id)
+        .expect("ICE: no functions found in contract file")
+        .push_fn(module_id, func_id, func.clone());
+
+    Ok(())
+}
+
+fn generate_compute_note_hash_and_nullifier(note_types: &Vec<String>) -> NoirFunction {
+    let function_source = generate_compute_note_hash_and_nullifier_source(note_types);
+
+    let (function_ast, errors) = parse_program(&function_source);
+    if !errors.is_empty() {
+        dbg!(errors.clone());
+    }
+    assert_eq!(errors.len(), 0, "Failed to parse Noir macro code. This is either a bug in the compiler or the Noir macro code");
+
+    let mut function_ast = function_ast.into_sorted();
+    function_ast.functions.remove(0)
+}
+
+fn generate_compute_note_hash_and_nullifier_source(note_types: &Vec<String>) -> String {
+    // TODO(#4649): The serialized_note parameter is a fixed-size array, but we don't know what length it should have.
+    // For now we hardcode it to 20, which is the same as MAX_NOTE_FIELDS_LENGTH.
+
+    if note_types.is_empty() {
+        // Even if the contract does not include any notes, other parts of the stack expect for this function to exist,
+        // so we include a dummy version.
+        "
+        unconstrained fn compute_note_hash_and_nullifier(
+            contract_address: AztecAddress,
+            nonce: Field,
+            storage_slot: Field,
+            note_type_id: Field,
+            serialized_note: [Field; 20]
+        ) -> pub [Field; 4] {
+            assert(false, \"This contract does not use private notes\");
+            [0, 0, 0, 0]
+        }"
+        .to_string()
+    } else {
+        // For contracts that include notes we do a simple if-else chain comparing note_type_id with the different
+        // get_note_type_id of each of the note types.
+
+        let if_statements: Vec<String> = note_types.iter().map(|note_type| format!(
+            "if (note_type_id == {0}::get_note_type_id()) {{
+                dep::aztec::note::utils::compute_note_hash_and_nullifier({0}::deserialize_content, note_header, serialized_note)
+            }}"
+        , note_type)).collect();
+
+        let full_if_statement = if_statements.join(" else ")
+            + "
+            else {
+                assert(false, \"Unknown note type ID\");
+                [0, 0, 0, 0]
+            }";
+
+        format!(
+            "
+            unconstrained fn compute_note_hash_and_nullifier(
+                contract_address: AztecAddress,
+                nonce: Field,
+                storage_slot: Field,
+                note_type_id: Field,
+                serialized_note: [Field; 20]
+            ) -> pub [Field; 4] {{
+                let note_header = dep::aztec::prelude::NoteHeader::new(contract_address, nonce, storage_slot);
+
+                {}
+            }}",
+            full_if_statement
+        )
+    }
+}

noir/noir-repo/aztec_macros/src/transforms/events.rs

@@ -0,0 +1,178 @@
+use iter_extended::vecmap;
+use noirc_errors::Span;
+use noirc_frontend::{
+    graph::CrateId,
+    macros_api::{
+        BlockExpression, FileId, HirContext, HirExpression, HirLiteral, HirStatement, NodeInterner,
+        NoirStruct, PathKind, StatementKind, StructId, StructType, Type, TypeImpl,
+        UnresolvedTypeData,
+    },
+    token::SecondaryAttribute,
+    ExpressionKind, FunctionDefinition, FunctionReturnType, FunctionVisibility, Literal,
+    NoirFunction, Visibility,
+};
+
+use crate::{
+    chained_dep,
+    utils::{
+        ast_utils::{
+            call, expression, ident, ident_path, make_statement, make_type, path, variable_path,
+        },
+        constants::SIGNATURE_PLACEHOLDER,
+        errors::AztecMacroError,
+        hir_utils::{collect_crate_structs, signature_of_type},
+    },
+};
+
+/// Generates the impl for an event selector
+///
+/// Inserts the following code:
+/// ```noir
+/// impl SomeStruct {
+///    fn selector() -> FunctionSelector {
+///       aztec::protocol_types::abis::function_selector::FunctionSelector::from_signature("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.
+pub fn generate_selector_impl(structure: &NoirStruct) -> TypeImpl {
+    let struct_type =
+        make_type(UnresolvedTypeData::Named(path(structure.name.clone()), vec![], true));
+
+    let selector_path =
+        chained_dep!("aztec", "protocol_types", "abis", "function_selector", "FunctionSelector");
+    let mut from_signature_path = selector_path.clone();
+    from_signature_path.segments.push(ident("from_signature"));
+
+    let selector_fun_body = BlockExpression(vec![make_statement(StatementKind::Expression(call(
+        variable_path(from_signature_path),
+        vec![expression(ExpressionKind::Literal(Literal::Str(SIGNATURE_PLACEHOLDER.to_string())))],
+    )))]);
+
+    // Define `FunctionSelector` return type
+    let return_type =
+        FunctionReturnType::Ty(make_type(UnresolvedTypeData::Named(selector_path, vec![], true)));
+
+    let mut selector_fn_def = FunctionDefinition::normal(
+        &ident("selector"),
+        &vec![],
+        &[],
+        &selector_fun_body,
+        &[],
+        &return_type,
+    );
+
+    selector_fn_def.visibility = FunctionVisibility::Public;
+
+    // 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), Span::default())],
+    }
+}
+
+/// 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(","))
+}
+
+/// 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,
+) -> Result<(), (AztecMacroError, FileId)> {
+    let struct_type = interner.get_struct(struct_id);
+    let selector_id = interner
+        .lookup_method(&Type::Struct(struct_type.clone(), vec![]), struct_id, "selector", false)
+        .ok_or_else(|| {
+            let error = AztecMacroError::EventError {
+                span: struct_type.borrow().location.span,
+                message: "Selector method not found".to_owned(),
+            };
+            (error, struct_type.borrow().location.file)
+        })?;
+    let selector_function = interner.function(&selector_id);
+
+    let compute_selector_statement = interner.statement(
+        selector_function.block(interner).statements().first().ok_or_else(|| {
+            let error = AztecMacroError::EventError {
+                span: struct_type.borrow().location.span,
+                message: "Compute selector statement not found".to_owned(),
+            };
+            (error, struct_type.borrow().location.file)
+        })?,
+    );
+
+    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,
+    }
+    .ok_or_else(|| {
+        let error = AztecMacroError::EventError {
+            span: struct_type.borrow().location.span,
+            message: "Compute selector statement is not a call expression".to_owned(),
+        };
+        (error, struct_type.borrow().location.file)
+    })?;
+
+    let first_arg_id = compute_selector_expression.arguments.first().ok_or_else(|| {
+        let error = AztecMacroError::EventError {
+            span: struct_type.borrow().location.span,
+            message: "Compute selector statement is not a call expression".to_owned(),
+        };
+        (error, struct_type.borrow().location.file)
+    })?;
+
+    match interner.expression(first_arg_id) {
+        HirExpression::Literal(HirLiteral::Str(signature))
+            if signature == SIGNATURE_PLACEHOLDER =>
+        {
+            let selector_literal_id = *first_arg_id;
+
+            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))),
+            );
+            Ok(())
+        }
+        _ => Err((
+            AztecMacroError::EventError {
+                span: struct_type.borrow().location.span,
+                message: "Signature placeholder literal does not match".to_owned(),
+            },
+            struct_type.borrow().location.file,
+        )),
+    }
+}
+
+pub fn transform_events(
+    crate_id: &CrateId,
+    context: &mut HirContext,
+) -> Result<(), (AztecMacroError, FileId)> {
+    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)?;
+        }
+    }
+    Ok(())
+}