refactor(avm): unify noir macros flow (#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).

noir-projects/aztec-nr/aztec/src/context.nr

@@ -9,28 +9,28 @@ mod interface;
 use interface::ContextInterface;
 use private_context::PrivateContext;
 use public_context::PublicContext;
-use avm_context::AVMContext;
+use avm_context::AvmContext;
 
 struct Context {
     private: Option<&mut PrivateContext>,
     public: Option<&mut PublicContext>,
-    public_vm: Option<&mut AVMContext>,
+    avm: Option<&mut AvmContext>,
 }
 
 impl Context {
     pub fn private(context: &mut PrivateContext) -> Context {
-        Context { private: Option::some(context), public: Option::none(), public_vm: Option::none() }
+        Context { private: Option::some(context), public: Option::none(), avm: Option::none() }
     }
 
     pub fn public(context: &mut PublicContext) -> Context {
-        Context { public: Option::some(context), private: Option::none(), public_vm: Option::none() }
+        Context { public: Option::some(context), private: Option::none(), avm: Option::none() }
     }
 
-    pub fn public_vm(context: &mut AVMContext) -> Context {
-        Context { public_vm: Option::some(context), public: Option::none(), private: Option::none() }
+    pub fn avm(context: &mut AvmContext) -> Context {
+        Context { avm: Option::some(context), public: Option::none(), private: Option::none() }
     }
 
     pub fn none() -> Context {
-        Context { public: Option::none(), private: Option::none(), public_vm: Option::none() }
+        Context { public: Option::none(), private: Option::none(), avm: Option::none() }
     }
 }

noir-projects/aztec-nr/aztec/src/context/avm_context.nr

@@ -7,13 +7,13 @@ use crate::context::inputs::avm_context_inputs::AvmContextInputs;
 use crate::context::interface::ContextInterface;
 use crate::context::interface::PublicContextInterface;
 
-struct AVMContext {
+struct AvmContext {
     inputs: AvmContextInputs,
 }
 
-impl AVMContext {
+impl AvmContext {
     pub fn new(inputs: AvmContextInputs) -> Self {
-        AVMContext { inputs }
+        AvmContext { inputs }
     }
 
     pub fn origin(self) -> AztecAddress {
@@ -73,7 +73,7 @@ impl AVMContext {
     }
 }
 
-impl PublicContextInterface for AVMContext {
+impl PublicContextInterface for AvmContext {
     fn block_number(self) -> Field {
         block_number()
     }
@@ -168,7 +168,7 @@ impl PublicContextInterface for AVMContext {
     }
 }
 
-impl ContextInterface for AVMContext {
+impl ContextInterface for AvmContext {
     fn push_new_note_hash(&mut self, note_hash: Field) {
         emit_note_hash(note_hash);
     }

noir-projects/aztec-nr/aztec/src/context/interface.nr

@@ -17,8 +17,8 @@ trait ContextInterface {
 }
 
 // TEMPORARY: This trait is to promote sharing of the current public context
-// and the upcoming AVMContext. This will be removed once the AVMContext is the default.
-// If you are adding something here, then you should also implement it in the AVMContext.
+// and the upcoming AvmContext. This will be removed once the AvmContext is the default.
+// If you are adding something here, then you should also implement it in the AvmContext.
 trait PublicContextInterface {
     fn block_number(self) -> Field;
     fn timestamp(self) -> u64;

noir-projects/noir-contracts/contracts/avm_test_contract/src/main.nr

@@ -25,7 +25,7 @@ contract AvmTest {
 
     // Libs
     use dep::aztec::prelude::Map;
-    use dep::aztec::state_vars::PublicMutable;
+    use dep::aztec::state_vars::{PublicImmutable, PublicMutable};
     use dep::aztec::protocol_types::{address::{AztecAddress, EthAddress}, constants::L1_TO_L2_MESSAGE_LENGTH};
     use dep::aztec::protocol_types::abis::function_selector::FunctionSelector;
     use dep::aztec::protocol_types::traits::ToField;
@@ -39,11 +39,23 @@ contract AvmTest {
         single: PublicMutable<Field>,
         list: PublicMutable<Note>,
         map: Map<AztecAddress, PublicMutable<u32>>,
+        immutable: PublicImmutable<Field>,
     }
 
     /************************************************************************
      * Storage
      ************************************************************************/
+    // FIX: calls unsupported getNullifierMembershipWitness.
+    // #[aztec(public-vm)]
+    // #[aztec(initializer)]
+    // fn constructor() {
+    //     storage.immutable.initialize(42);
+    // }
+
+    unconstrained fn view_storage_immutable() -> pub Field {
+        storage.immutable.read()
+    }
+
     unconstrained fn view_storage_single() -> pub Field {
         storage.single.read()
     }

noir/noir-repo/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;

noir/noir-repo/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

yarn-project/end-to-end/src/e2e_avm_simulator.test.ts

@@ -25,6 +25,11 @@ describe('e2e_avm_simulator', () => {
   }, 50_000);
 
   describe('Storage', () => {
+    // FIX: Enable once the contract function works.
+    // it('Read immutable (initialized) storage (Field)', async () => {
+    //   expect(await avmContact.methods.view_storage_immutable().view()).toEqual(42n);
+    // });
+
     it('Modifies storage (Field)', async () => {
       await avmContact.methods.set_storage_single(20n).send().wait();
       expect(await avmContact.methods.view_storage_single().view()).toEqual(20n);

yarn-project/simulator/src/avm/avm_execution_environment.ts

@@ -1,6 +1,6 @@
 import { FunctionSelector, GlobalVariables } from '@aztec/circuits.js';
+import { computeVarArgsHash } from '@aztec/circuits.js/hash';
 import { AztecAddress } from '@aztec/foundation/aztec-address';
-import { pedersenHash } from '@aztec/foundation/crypto';
 import { EthAddress } from '@aztec/foundation/eth-address';
 import { Fr } from '@aztec/foundation/fields';
 
@@ -54,7 +54,7 @@ export class AvmExecutionEnvironment {
   ) {
     // We encode some extra inputs (AvmContextInputs) in calldata.
     // This will have to go once we move away from one proof per call.
-    const inputs = new AvmContextInputs(temporaryFunctionSelector.toField(), pedersenHash(calldata));
+    const inputs = new AvmContextInputs(temporaryFunctionSelector.toField(), computeVarArgsHash(calldata));
     this.calldata = [...inputs.toFields(), ...calldata];
   }
 

yarn-project/simulator/src/avm/avm_simulator.test.ts

@@ -1,4 +1,5 @@
 import { UnencryptedL2Log } from '@aztec/circuit-types';
+import { computeVarArgsHash } from '@aztec/circuits.js/hash';
 import { EventSelector, FunctionSelector } from '@aztec/foundation/abi';
 import { AztecAddress } from '@aztec/foundation/aztec-address';
 import { keccak, pedersenHash, poseidonHash, sha256 } from '@aztec/foundation/crypto';
@@ -20,22 +21,20 @@ import {
   initL1ToL2MessageOracleInput,
   initMachineState,
 } from './fixtures/index.js';
-import { Add, CalldataCopy, Instruction, Return } from './opcodes/index.js';
+import { Add, CalldataCopy, Return } from './opcodes/index.js';
 import { encodeToBytecode } from './serialization/bytecode_serialization.js';
 
 describe('AVM simulator: injected bytecode', () => {
   let calldata: Fr[];
-  let ops: Instruction[];
   let bytecode: Buffer;
 
   beforeAll(() => {
     calldata = [new Fr(1), new Fr(2)];
-    ops = [
+    bytecode = encodeToBytecode([
       new CalldataCopy(/*indirect=*/ 0, /*cdOffset=*/ adjustCalldataIndex(0), /*copySize=*/ 2, /*dstOffset=*/ 0),
       new Add(/*indirect=*/ 0, TypeTag.FIELD, /*aOffset=*/ 0, /*bOffset=*/ 1, /*dstOffset=*/ 2),
       new Return(/*indirect=*/ 0, /*returnOffset=*/ 2, /*copySize=*/ 1),
-    ];
-    bytecode = encodeToBytecode(ops);
+    ]);
   });
 
   it('Should execute bytecode that performs basic addition', async () => {
@@ -247,7 +246,7 @@ describe('AVM simulator: transpiled Noir contracts', () => {
       const results = await new AvmSimulator(context).executeBytecode(bytecode);
 
       expect(results.reverted).toBe(false);
-      expect(results.output).toEqual([pedersenHash(calldata)]);
+      expect(results.output).toEqual([computeVarArgsHash(calldata)]);
     });
   });
 
@@ -462,6 +461,7 @@ describe('AVM simulator: transpiled Noir contracts', () => {
 
       const results = await new AvmSimulator(context).executeBytecode(callBytecode);
 
+      expect(results.revertReason).toBeUndefined();
       expect(results.reverted).toBe(false);
       expect(results.output).toEqual([new Fr(3)]);
     });