refactor(avm): unify noir macros flow (AztecProtocol/aztec-packages#5461

)

Step towards integrating with initializers, etc.

However in the future we might want to separate private and public-vm, many things might end up being different enough (see IFs in code).

.aztec-sync-commit

@@ -1 +1 @@
-27bd8d318df486f6d30a01212f9d7894cafcec74
+54aee58952b2433ccad83f1b5fc3088957b10fbb

aztec_macros/src/lib.rs

@@ -4,7 +4,7 @@ mod utils;
 use transforms::{
     compute_note_hash_and_nullifier::inject_compute_note_hash_and_nullifier,
     events::{generate_selector_impl, transform_events},
-    functions::{transform_function, transform_unconstrained, transform_vm_function},
+    functions::{transform_function, transform_unconstrained},
     note_interface::generate_note_interface_impl,
     storage::{
         assign_storage_slots, check_for_storage_definition, check_for_storage_implementation,
@@ -138,19 +138,22 @@ fn transform_module(module: &mut SortedModule) -> Result<bool, AztecMacroError>
         }
 
         // Apply transformations to the function based on collected attributes
-        if is_private || is_public {
+        if is_private || is_public || is_public_vm {
             transform_function(
-                if is_private { "Private" } else { "Public" },
+                if is_private {
+                    "Private"
+                } else if is_public_vm {
+                    "Avm"
+                } else {
+                    "Public"
+                },
                 func,
                 storage_defined,
                 is_initializer,
                 insert_init_check,
                 is_internal,
             )?;
             has_transformed_module = true;
-        } else if is_public_vm {
-            transform_vm_function(func, storage_defined)?;
-            has_transformed_module = true;
         } else if storage_defined && func.def.is_unconstrained {
             transform_unconstrained(func);
             has_transformed_module = true;

aztec_macros/src/transforms/functions.rs

@@ -35,6 +35,7 @@ pub fn transform_function(
     let context_name = format!("{}Context", ty);
     let inputs_name = format!("{}ContextInputs", ty);
     let return_type_name = format!("{}CircuitPublicInputs", ty);
+    let is_avm = ty == "Avm";
 
     // Add check that msg sender equals this address and flag function as internal
     if is_internal {
@@ -60,20 +61,26 @@ pub fn transform_function(
     }
 
     // Insert the context creation as the first action
-    let create_context = create_context(&context_name, &func.def.parameters)?;
+    let create_context = if !is_avm {
+        create_context(&context_name, &func.def.parameters)?
+    } else {
+        create_context_avm()?
+    };
     func.def.body.statements.splice(0..0, (create_context).iter().cloned());
 
     // Add the inputs to the params
     let input = create_inputs(&inputs_name);
     func.def.parameters.insert(0, input);
 
     // Abstract return types such that they get added to the kernel's return_values
-    if let Some(return_values) = abstract_return_values(func) {
-        // In case we are pushing return values to the context, we remove the statement that originated it
-        // This avoids running duplicate code, since blocks like if/else can be value returning statements
-        func.def.body.statements.pop();
-        // Add the new return statement
-        func.def.body.statements.push(return_values);
+    if !is_avm {
+        if let Some(return_values) = abstract_return_values(func) {
+            // In case we are pushing return values to the context, we remove the statement that originated it
+            // This avoids running duplicate code, since blocks like if/else can be value returning statements
+            func.def.body.statements.pop();
+            // Add the new return statement
+            func.def.body.statements.push(return_values);
+        }
     }
 
     // Before returning mark the contract as initialized
@@ -83,49 +90,29 @@ pub fn transform_function(
     }
 
     // Push the finish method call to the end of the function
-    let finish_def = create_context_finish();
-    func.def.body.statements.push(finish_def);
+    if !is_avm {
+        let finish_def = create_context_finish();
+        func.def.body.statements.push(finish_def);
+    }
 
-    let return_type = create_return_type(&return_type_name);
-    func.def.return_type = return_type;
-    func.def.return_visibility = Visibility::Public;
+    // The AVM doesn't need a return type yet.
+    if !is_avm {
+        let return_type = create_return_type(&return_type_name);
+        func.def.return_type = return_type;
+        func.def.return_visibility = Visibility::Public;
+    }
 
     // Distinct return types are only required for private functions
     // Public functions should have unconstrained auto-inferred
     match ty {
         "Private" => func.def.return_distinctness = Distinctness::Distinct,
-        "Public" => func.def.is_unconstrained = true,
+        "Public" | "Avm" => func.def.is_unconstrained = true,
         _ => (),
     }
 
     Ok(())
 }
 
-/// Transform a function to work with AVM bytecode
-pub fn transform_vm_function(
-    func: &mut NoirFunction,
-    storage_defined: bool,
-) -> Result<(), AztecMacroError> {
-    // Create access to storage
-    if storage_defined {
-        let storage = abstract_storage("public_vm", true);
-        func.def.body.statements.insert(0, storage);
-    }
-
-    // Push Avm context creation to the beginning of the function
-    let create_context = create_avm_context()?;
-    func.def.body.statements.insert(0, create_context);
-
-    // Add the inputs to the params (first!)
-    let input = create_inputs("AvmContextInputs");
-    func.def.parameters.insert(0, input);
-
-    // We want the function to be seen as a public function
-    func.def.is_unconstrained = true;
-
-    Ok(())
-}
-
 /// Transform Unconstrained
 ///
 /// Inserts the following code at the beginning of an unconstrained function
@@ -339,36 +326,36 @@ fn create_context(ty: &str, params: &[Param]) -> Result<Vec<Statement>, AztecMac
     Ok(injected_expressions)
 }
 
-/// Creates an mutable avm context
+/// Creates the private context object to be accessed within the function, the parameters need to be extracted to be
+/// appended into the args hash object.
 ///
+/// The replaced code:
 /// ```noir
-/// /// Before
 /// #[aztec(public-vm)]
-/// fn foo() -> Field {
-///   let mut context = aztec::context::AVMContext::new();
-///   let timestamp = context.timestamp();
-///   // ...
-/// }
-///
-/// /// After
-/// #[aztec(private)]
-/// fn foo() -> Field {
-///     let mut timestamp = context.timestamp();
-///     // ...
+/// fn foo(inputs: AvmContextInputs, ...) -> Field {
+///     let mut context = AvmContext::new(inputs);
 /// }
-fn create_avm_context() -> Result<Statement, AztecMacroError> {
+/// ```
+fn create_context_avm() -> Result<Vec<Statement>, AztecMacroError> {
+    let mut injected_expressions: Vec<Statement> = vec![];
+
     // Create the inputs to the context
+    let ty = "AvmContext";
     let inputs_expression = variable("inputs");
+    let path_snippet = ty.to_case(Case::Snake); // e.g. private_context
 
+    // let mut context = {ty}::new(inputs, hash);
     let let_context = mutable_assignment(
         "context", // Assigned to
         call(
-            variable_path(chained_dep!("aztec", "context", "AVMContext", "new")), // Path
-            vec![inputs_expression],                                              // args
+            variable_path(chained_dep!("aztec", "context", &path_snippet, ty, "new")), // Path
+            vec![inputs_expression],                                                   // args
         ),
     );
+    injected_expressions.push(let_context);
 
-    Ok(let_context)
+    // Return all expressions that will be injected by the hasher
+    Ok(injected_expressions)
 }
 
 /// Abstract Return Type

tooling/noir_js_backend_barretenberg/package.json

@@ -42,7 +42,7 @@
     "lint": "NODE_NO_WARNINGS=1 eslint . --ext .ts --ignore-path ./.eslintignore  --max-warnings 0"
   },
   "dependencies": {
-    "@aztec/bb.js": "0.31.0",
+    "@aztec/bb.js": "portal:../../../../barretenberg/ts",
     "@noir-lang/types": "workspace:*",
     "fflate": "^0.8.0"
   },

yarn.lock

@@ -221,19 +221,18 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@aztec/bb.js@npm:0.31.0":
-  version: 0.31.0
-  resolution: "@aztec/bb.js@npm:0.31.0"
+"@aztec/bb.js@portal:../../../../barretenberg/ts::locator=%40noir-lang%2Fbackend_barretenberg%40workspace%3Atooling%2Fnoir_js_backend_barretenberg":
+  version: 0.0.0-use.local
+  resolution: "@aztec/bb.js@portal:../../../../barretenberg/ts::locator=%40noir-lang%2Fbackend_barretenberg%40workspace%3Atooling%2Fnoir_js_backend_barretenberg"
   dependencies:
     comlink: ^4.4.1
     commander: ^10.0.1
     debug: ^4.3.4
     tslib: ^2.4.0
   bin:
-    bb.js: dest/node/main.js
-  checksum: 3bcd8cef1fb0b5767332b757875719b4b4f2c34d5b4ed22984f305b920afbb74e17ccf626b3efe66cae904c3df8f74de1aa8c798b17195be13db60a65730bd98
+    bb.js: ./dest/node/main.js
   languageName: node
-  linkType: hard
+  linkType: soft
 
 "@babel/code-frame@npm:^7.0.0, @babel/code-frame@npm:^7.10.4, @babel/code-frame@npm:^7.12.11, @babel/code-frame@npm:^7.16.0, @babel/code-frame@npm:^7.22.13, @babel/code-frame@npm:^7.23.5, @babel/code-frame@npm:^7.8.3":
   version: 7.23.5
@@ -4396,7 +4395,7 @@ __metadata:
   version: 0.0.0-use.local
   resolution: "@noir-lang/backend_barretenberg@workspace:tooling/noir_js_backend_barretenberg"
   dependencies:
-    "@aztec/bb.js": 0.31.0
+    "@aztec/bb.js": "portal:../../../../barretenberg/ts"
     "@noir-lang/types": "workspace:*"
     "@types/node": ^20.6.2
     "@types/prettier": ^3