diff --git a/compiler/noirc_frontend/src/hir/def_map/aztec_library.rs b/compiler/noirc_frontend/src/hir/def_map/aztec_library.rs index 9d7dfc458eb..e3a9735f936 100644 --- a/compiler/noirc_frontend/src/hir/def_map/aztec_library.rs +++ b/compiler/noirc_frontend/src/hir/def_map/aztec_library.rs @@ -10,6 +10,7 @@ use crate::{ ParsedModule, Path, PathKind, Pattern, Statement, UnresolvedType, UnresolvedTypeData, Visibility, }; +use crate::{PrefixExpression, UnaryOp}; use noirc_errors::FileDiagnostic; // @@ -55,8 +56,12 @@ fn call(func: Expression, arguments: Vec) -> Expression { expression(ExpressionKind::Call(Box::new(CallExpression { func: Box::new(func), arguments }))) } -fn mutable(pattern: &str) -> Pattern { - Pattern::Mutable(Box::new(Pattern::Identifier(ident(pattern))), Span::default()) +fn pattern(name: &str) -> Pattern { + Pattern::Identifier(ident(name)) +} + +fn mutable(name: &str) -> Pattern { + Pattern::Mutable(Box::new(pattern(name)), Span::default()) } fn mutable_assignment(name: &str, assigned_to: Expression) -> Statement { @@ -67,6 +72,21 @@ fn mutable_assignment(name: &str, assigned_to: Expression) -> Statement { }) } +fn mutable_reference(variable_name: &str) -> Expression { + expression(ExpressionKind::Prefix(Box::new(PrefixExpression { + operator: UnaryOp::MutableReference, + rhs: variable(variable_name), + }))) +} + +fn assignment(name: &str, assigned_to: Expression) -> Statement { + Statement::Let(LetStatement { + pattern: pattern(name), + r#type: make_type(UnresolvedTypeData::Unspecified), + expression: assigned_to, + }) +} + fn member_access(lhs: &str, rhs: &str) -> Expression { expression(ExpressionKind::MemberAccess(Box::new(MemberAccessExpression { lhs: variable(lhs), @@ -141,7 +161,9 @@ pub(crate) fn transform( // Covers all functions in the ast for submodule in ast.submodules.iter_mut().filter(|submodule| submodule.is_contract) { - if transform_module(&mut submodule.contents.functions) { + let storage_defined = check_for_storage_definition(&submodule.contents); + + if transform_module(&mut submodule.contents.functions, storage_defined) { check_for_aztec_dependency(crate_id, context, errors); include_relevant_imports(&mut submodule.contents); } @@ -183,27 +205,37 @@ fn check_for_aztec_dependency( } } +// Check to see if the user has defined a storage struct +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]) -> bool { +fn transform_module(functions: &mut [NoirFunction], storage_defined: bool) -> bool { let mut has_annotated_functions = false; for func in 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); + transform_function("Private", func, storage_defined); has_annotated_functions = true; } "aztec(public)" => { - transform_function("Public", func); + transform_function("Public", func, storage_defined); has_annotated_functions = true; } _ => continue, } } } + // 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_annotated_functions } @@ -212,11 +244,17 @@ fn transform_module(functions: &mut [NoirFunction]) -> bool { /// - A new Input that is provided for a kernel app circuit, named: {Public/Private}ContextInputs /// - Hashes all of the function input variables /// - This instantiates a helper function -fn transform_function(ty: &str, func: &mut NoirFunction) { +fn transform_function(ty: &str, func: &mut NoirFunction, storage_defined: bool) { let context_name = format!("{}Context", ty); let inputs_name = format!("{}ContextInputs", ty); let return_type_name = format!("{}CircuitPublicInputs", ty); + // Add access to the storage struct + if storage_defined { + let storage_def = abstract_storage(&ty.to_lowercase(), false); + func.def.body.0.insert(0, storage_def); + } + // Insert the context creation as the first action let create_context = create_context(&context_name, &func.def.parameters); func.def.body.0.splice(0..0, (create_context).iter().cloned()); @@ -247,6 +285,18 @@ fn transform_function(ty: &str, func: &mut NoirFunction) { } } +/// Transform Unconstrained +/// +/// Inserts the following code at the beginning of an unconstrained function +/// ```noir +/// let storage = Storage::init(Context::none()); +/// ``` +/// +/// This will allow developers to access their contract' storage struct in unconstrained functions +fn transform_unconstrained(func: &mut NoirFunction) { + func.def.body.0.insert(0, abstract_storage("Unconstrained", true)); +} + /// 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. /// @@ -413,6 +463,51 @@ fn abstract_return_values(func: &NoirFunction) -> Option { } } +/// Abstract storage +/// +/// For private functions: +/// ```noir +/// #[aztec(private)] +/// fn lol() { +/// let storage = Storage::init(Context::private(context)); +/// } +/// ``` +/// +/// For public functions: +/// ```noir +/// #[aztec(public)] +/// fn lol() { +/// let storage = Storage::init(Context::public(context)); +/// } +/// ``` +/// +/// For unconstrained functions: +/// ```noir +/// unconstrained fn lol() { +/// let storage = Storage::init(Context::none()); +/// } +fn abstract_storage(typ: &str, unconstrained: bool) -> Statement { + let init_context_call = if unconstrained { + call( + variable_path(chained_path!("aztec", "context", "Context", "none")), // Path + vec![], // args + ) + } else { + call( + variable_path(chained_path!("aztec", "context", "Context", typ)), // Path + vec![mutable_reference("context")], // args + ) + }; + + assignment( + "storage", // Assigned to + call( + variable_path(chained_path!("Storage", "init")), // Path + vec![init_context_call], // args + ), + ) +} + /// Context Return Values /// /// Creates an instance to the context return values