chore: aztec-macros refactor (#5127)

Partially addresses
#5080

Until we have full macro support in noir, this PR aims to make using
what we currently have a little bit less painful.
The crate is now divided into utils of several kind (ast manipulation,
error handling, hir manipulation...) and transforms.

A transform is an end-to-end modification to the noir code that
generates a valid compilation result from a given contract. This makes
it easier to manipulate (and hopefully soon, replace) said modifications
atomically, improving readability.

This also modifies the first macro pass to iterate over all crates, not
only the root one, which allows us to implement more stuff without
having to dig deep into the def collector or (god forbid) the interner.

The second macro pass has been renamed to `process_collected_defs`.
Tried to generalize it a bit more, but providing a mutable ref for the
whole `def_collector` turned out to be tricky without cloning
`def_maps`, so it's staying like it is now.

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(())
+}