Skip to content

Commit

Permalink
macro crate cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
Thunkar committed Mar 11, 2024
1 parent 32182b1 commit 72cfa76
Show file tree
Hide file tree
Showing 19 changed files with 2,295 additions and 2,170 deletions.
2,183 changes: 41 additions & 2,142 deletions noir/noir-repo/aztec_macros/src/lib.rs

Large diffs are not rendered by default.

Original file line number Diff line number Diff line change
@@ -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
)
}
}
178 changes: 178 additions & 0 deletions noir/noir-repo/aztec_macros/src/transforms/events.rs
Original file line number Diff line number Diff line change
@@ -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(())
}
Loading

0 comments on commit 72cfa76

Please sign in to comment.