Skip to content

Commit

Permalink
working storage layout and notes
Browse files Browse the repository at this point in the history
  • Loading branch information
Thunkar committed Mar 26, 2024
1 parent 025bdca commit 554d3c1
Show file tree
Hide file tree
Showing 8 changed files with 276 additions and 142 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ contract DocsExample {
use dep::aztec::prelude::{
AztecAddress, FunctionSelector, NoteHeader, NoteGetterOptions, NoteViewerOptions,
PrivateContext, Map, PublicMutable, PublicImmutable, PrivateMutable, PrivateImmutable,
PrivateSet, SharedImmutable, Storable
PrivateSet, SharedImmutable
};
use dep::aztec::{note::note_getter_options::Comparator, context::{PublicContext, Context}};
// how to import methods from other files/folders within your workspace
Expand Down
258 changes: 125 additions & 133 deletions noir/noir-repo/aztec_macros/src/transforms/storage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -266,9 +266,8 @@ pub fn assign_storage_slots(
) -> Result<(), (AztecMacroError, FileId)> {
let traits: Vec<_> = collect_traits(context);
if let Some((_, file_id)) = get_contract_module_data(context, crate_id) {
let storage_struct = collect_crate_structs(crate_id, context)
.iter()
.find_map(|&struct_id| {
let maybe_storage_struct =
collect_crate_structs(crate_id, context).iter().find_map(|&struct_id| {
let r#struct = context.def_interner.get_struct(struct_id);
let attributes = context.def_interner.struct_attributes(&struct_id);
if attributes.iter().any(|attr| is_custom_attribute(attr, "aztec(storage)"))
Expand All @@ -278,19 +277,10 @@ pub fn assign_storage_slots(
} else {
None
}
})
.ok_or((
AztecMacroError::CouldNotAssignStorageSlots {
secondary_message: Some("Storage struct not found".to_string()),
},
file_id,
))?;
});

let storage_layout = context
.def_interner
.get_all_globals()
.iter()
.find_map(|global_info| {
let maybe_storage_layout =
context.def_interner.get_all_globals().iter().find_map(|global_info| {
let statement = context.def_interner.get_global_let_statement(global_info.id);
if statement.clone().is_some_and(|stmt| {
stmt.attributes
Expand All @@ -307,43 +297,44 @@ pub fn assign_storage_slots(
} else {
None
}
})
.ok_or((
AztecMacroError::CouldNotAssignStorageSlots {
secondary_message: Some("Storage layout struct not found".to_string()),
},
file_id,
))?;
});

let init_id = context
.def_interner
.lookup_method(
&Type::Struct(context.def_interner.get_struct(storage_struct.borrow().id), vec![]),
storage_struct.borrow().id,
"init",
false,
)
.ok_or((
AztecMacroError::CouldNotAssignStorageSlots {
secondary_message: Some(
"Storage struct must have an init function".to_string(),
if let (Some(storage_struct), Some(storage_layout)) =
(maybe_storage_struct, maybe_storage_layout)
{
let init_id = context
.def_interner
.lookup_method(
&Type::Struct(
context.def_interner.get_struct(storage_struct.borrow().id),
vec![],
),
storage_struct.borrow().id,
"init",
false,
)
.ok_or((
AztecMacroError::CouldNotAssignStorageSlots {
secondary_message: Some(
"Storage struct must have an init function".to_string(),
),
},
file_id,
))?;
let init_function =
context.def_interner.function(&init_id).block(&context.def_interner);
let init_function_statement_id = init_function.statements().first().ok_or((
AztecMacroError::CouldNotAssignStorageSlots {
secondary_message: Some("Init storage statement not found".to_string()),
},
file_id,
))?;
let init_function = context.def_interner.function(&init_id).block(&context.def_interner);
let init_function_statement_id = init_function.statements().first().ok_or((
AztecMacroError::CouldNotAssignStorageSlots {
secondary_message: Some("Init storage statement not found".to_string()),
},
file_id,
))?;
let storage_constructor_statement =
context.def_interner.statement(init_function_statement_id);

let storage_constructor_expression = match storage_constructor_statement {
HirStatement::Expression(expression_id) => {
match context.def_interner.expression(&expression_id) {
let storage_constructor_statement =
context.def_interner.statement(init_function_statement_id);

let storage_constructor_expression = match storage_constructor_statement {
HirStatement::Expression(expression_id) => {
match context.def_interner.expression(&expression_id) {
HirExpression::Constructor(hir_constructor_expression) => {
Ok(hir_constructor_expression)
}
Expand All @@ -357,106 +348,107 @@ pub fn assign_storage_slots(
file_id,
)),
}
}
_ => Err((
AztecMacroError::CouldNotAssignStorageSlots {
secondary_message: Some(
"Storage constructor statement must be an expression".to_string(),
),
},
file_id,
)),
}?;

let mut storage_slot: u64 = 1;
for (index, (_, expr_id)) in storage_constructor_expression.fields.iter().enumerate() {
let fields = storage_struct.borrow().get_fields(&[]);
let (field_name, field_type) = fields.get(index).unwrap();
let new_call_expression = match context.def_interner.expression(expr_id) {
HirExpression::Call(hir_call_expression) => Ok(hir_call_expression),
}
_ => Err((
AztecMacroError::CouldNotAssignStorageSlots {
secondary_message: Some(
"Storage field initialization expression is not a call expression"
.to_string(),
"Storage constructor statement must be an expression".to_string(),
),
},
file_id,
)),
}?;

let slot_arg_expression =
context.def_interner.expression(&new_call_expression.arguments[1]);
let mut storage_slot: u64 = 1;
for (index, (_, expr_id)) in storage_constructor_expression.fields.iter().enumerate() {
let fields = storage_struct.borrow().get_fields(&[]);
let (field_name, field_type) = fields.get(index).unwrap();
let new_call_expression = match context.def_interner.expression(expr_id) {
HirExpression::Call(hir_call_expression) => Ok(hir_call_expression),
_ => Err((
AztecMacroError::CouldNotAssignStorageSlots {
secondary_message: Some(
"Storage field initialization expression is not a call expression"
.to_string(),
),
},
file_id,
)),
}?;

let slot_arg_expression =
context.def_interner.expression(&new_call_expression.arguments[1]);

let current_storage_slot = match slot_arg_expression {
HirExpression::Literal(HirLiteral::Integer(slot, _)) => Ok(slot.to_u128()),
_ => Err((
let current_storage_slot = match slot_arg_expression {
HirExpression::Literal(HirLiteral::Integer(slot, _)) => Ok(slot.to_u128()),
_ => Err((
AztecMacroError::CouldNotAssignStorageSlots {
secondary_message: Some(
"Storage slot argument expression must be a literal integer"
.to_string(),
),
},
file_id,
)),
}?;

let storage_layout_field =
storage_layout.fields.iter().find(|field| field.0 .0.contents == *field_name);

let storage_layout_slot_expr = if let Some((_, expr_id)) = storage_layout_field {
let expr = context.def_interner.expression(expr_id);
if let HirExpression::Constructor(storage_layout_field_storable_expr) = expr {
storage_layout_field_storable_expr.fields.iter().find_map(
|(field, expr_id)| {
if field.0.contents == "slot" {
Some(*expr_id)
} else {
None
}
},
)
} else {
None
}
} else {
None
}
.ok_or((
AztecMacroError::CouldNotAssignStorageSlots {
secondary_message: Some(
"Storage slot argument expression must be a literal integer"
.to_string(),
),
secondary_message: Some(format!(
"Storage layout field ({}) not found or has an incorrect type",
field_name
)),
},
file_id,
)),
}?;
))?;

let storage_layout_field = storage_layout
.fields
.iter()
.find(|field| field.0 .0.contents == field_name.to_string());

let storage_layout_slot_expr = if let Some((_, expr_id)) = storage_layout_field {
let expr = context.def_interner.expression(expr_id);
if let HirExpression::Constructor(storage_layout_field_storable_expr) = expr {
storage_layout_field_storable_expr.fields.iter().find_map(|(field, expr_id)| {
if field.0.contents == "slot" {
Some(*expr_id)
} else {
None
}
})
let new_storage_slot = if current_storage_slot == 0 {
u128::from(storage_slot)
} else {
None
}
} else {
None
current_storage_slot
};

let type_serialized_len =
get_serialized_length(&traits, field_type, &context.def_interner)
.map_err(|err| (err, file_id))?;

context.def_interner.update_expression(new_call_expression.arguments[1], |expr| {
*expr = HirExpression::Literal(HirLiteral::Integer(
FieldElement::from(new_storage_slot),
false,
));
});

context.def_interner.update_expression(storage_layout_slot_expr, |expr| {
*expr = HirExpression::Literal(HirLiteral::Integer(
FieldElement::from(new_storage_slot),
false,
));
});

storage_slot += type_serialized_len;
}
.ok_or((
AztecMacroError::CouldNotAssignStorageSlots {
secondary_message: Some(format!(
"Storage layout field ({}) not found or has an incorrect type",
field_name
)),
},
file_id,
))?;

let new_storage_slot = if current_storage_slot == 0 {
u128::from(storage_slot)
} else {
current_storage_slot
};

let type_serialized_len =
get_serialized_length(&traits, field_type, &context.def_interner)
.map_err(|err| (err, file_id))?;

context.def_interner.update_expression(new_call_expression.arguments[1], |expr| {
*expr = HirExpression::Literal(HirLiteral::Integer(
FieldElement::from(new_storage_slot),
false,
));
});

context.def_interner.update_expression(storage_layout_slot_expr, |expr| {
*expr = HirExpression::Literal(HirLiteral::Integer(
FieldElement::from(new_storage_slot),
false,
));
});

storage_slot += type_serialized_len;
}
}

Expand All @@ -478,10 +470,10 @@ pub fn generate_storage_layout(
let mut storable_fields_impl = vec![];

definition.fields.iter().enumerate().for_each(|(index, (field_ident, field_type))| {
storable_fields.push(format!("{}: Storable<N{}>", field_ident, index));
storable_fields.push(format!("{}: dep::aztec::prelude::Storable<N{}>", field_ident, index));
generic_args.push(format!("N{}", index));
storable_fields_impl.push(format!(
"{}: Storable {{ slot: 0, typ: \"{}\" }}",
"{}: dep::aztec::prelude::Storable {{ slot: 0, typ: \"{}\" }}",
field_ident,
field_type.to_string().replace("plain::", "")
));
Expand Down
2 changes: 1 addition & 1 deletion noir/noir-repo/aztec_macros/src/utils/errors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ impl From<AztecMacroError> for MacroError {
},
AztecMacroError::CouldNotExportStorageLayout { secondary_message, span } => MacroError {
primary_message: "Could not generate and export storage layout".to_string(),
secondary_message: secondary_message,
secondary_message,
span,
},
AztecMacroError::EventError { span, message } => MacroError {
Expand Down
44 changes: 43 additions & 1 deletion yarn-project/aztec.js/src/contract/contract_base.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { computePartialAddress } from '@aztec/circuits.js';
import { Fr, computePartialAddress } from '@aztec/circuits.js';
import { ContractArtifact, FunctionArtifact, FunctionSelector } from '@aztec/foundation/abi';
import { ContractInstanceWithAddress } from '@aztec/types/contracts';

Expand All @@ -16,6 +16,48 @@ export type ContractMethod = ((...args: any[]) => ContractFunctionInteraction) &
readonly selector: FunctionSelector;
};

/**
* Type representing a field layout in the storage of a contract.
*/
type FieldLayout = {
/**
* Slot in which the field is stored.
*/
slot: Fr;
/**
* Type being stored at the slot
*/
typ: string;
};

/**
* Type representing a note in use in the contract.
*/
type ContractNote = {
/**
* Note identifier
*/
id: Fr;
/**
* Type of the note
*/
typ: string;
};

/**
* Type representing the storage layout of a contract.
*/
export type ContractStorageLayout<T extends string> = {
[K in T]: FieldLayout;
};

/**
* Type representing the notes used in a contract.
*/
export type ContractNotes<T extends string> = {
[K in T]: ContractNote;
};

/**
* Abstract implementation of a contract extended by the Contract class and generated contract types.
*/
Expand Down
2 changes: 2 additions & 0 deletions yarn-project/aztec.js/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@ export {
Contract,
ContractBase,
ContractMethod,
ContractStorageLayout,
ContractNotes,
SentTx,
BatchCall,
DeployMethod,
Expand Down
Loading

0 comments on commit 554d3c1

Please sign in to comment.