feat(avm): integrate AVM with initializers (#5469)

This PR creates private/public/avm variants of some lib functions used for initializers. This is because in general the flow in private/public/avm will be substantially different. The public version will be removed once the AVM is fully enabled.

I'm adding a new AVM test contract because: having an initializer adds a check to EVERY other function in the contract. This makes the unit tests fail.

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

@@ -1,37 +1,79 @@
 use dep::protocol_types::{
-    hash::{silo_nullifier, pedersen_hash}, constants::GENERATOR_INDEX__CONSTRUCTOR,
-    abis::function_selector::FunctionSelector
+    address::AztecAddress, hash::{silo_nullifier, pedersen_hash},
+    constants::GENERATOR_INDEX__CONSTRUCTOR, abis::function_selector::FunctionSelector
 };
 
 use crate::{
-    context::{PrivateContext, PublicContext, ContextInterface},
+    context::{PrivateContext, PublicContext, AvmContext, ContextInterface},
     oracle::get_contract_instance::get_contract_instance,
+    oracle::get_contract_instance::get_contract_instance_avm,
     history::nullifier_inclusion::prove_nullifier_inclusion
 };
 
-pub fn mark_as_initialized<TContext>(context: &mut TContext) where TContext: ContextInterface {
-    let init_nullifier = compute_unsiloed_contract_initialization_nullifier(*context);
+pub fn mark_as_initialized_public(context: &mut PublicContext) {
+    mark_as_initialized(context);
+}
+
+pub fn mark_as_initialized_avm(context: &mut AvmContext) {
+    mark_as_initialized(context);
+}
+
+pub fn mark_as_initialized_private(context: &mut PrivateContext) {
+    mark_as_initialized(context);
+}
+
+fn mark_as_initialized<TContext>(context: &mut TContext) where TContext: ContextInterface {
+    let init_nullifier = compute_unsiloed_contract_initialization_nullifier((*context).this_address());
     ContextInterface::push_new_nullifier(context, init_nullifier, 0);
 }
 
-pub fn assert_is_initialized<TContext>(context: &mut TContext) where TContext: ContextInterface {
-    let init_nullifier = compute_contract_initialization_nullifier(*context);
+pub fn assert_is_initialized_public(context: &mut PublicContext) {
+    let init_nullifier = compute_contract_initialization_nullifier(context.this_address());
     prove_nullifier_inclusion(init_nullifier, *context);
 }
 
-pub fn compute_contract_initialization_nullifier<TContext>(context: TContext) -> Field where TContext: ContextInterface {
-    let address = context.this_address();
+pub fn assert_is_initialized_avm(context: &mut AvmContext) {
+    // WARNING: the AVM always expects UNSILOED nullifiers!
+    // TODO(fcarreiro@): Change current private/public to take unsiloed nullifiers and an address.
+    let init_nullifier = compute_unsiloed_contract_initialization_nullifier(context.this_address());
+    assert(context.nullifier_exists(init_nullifier));
+}
+
+pub fn assert_is_initialized_private(context: &mut PrivateContext) {
+    let init_nullifier = compute_contract_initialization_nullifier(context.this_address());
+    prove_nullifier_inclusion(init_nullifier, *context);
+}
+
+fn compute_contract_initialization_nullifier(address: AztecAddress) -> Field {
     silo_nullifier(
         address,
-        compute_unsiloed_contract_initialization_nullifier(context)
+        compute_unsiloed_contract_initialization_nullifier(address)
     )
 }
 
-pub fn compute_unsiloed_contract_initialization_nullifier<TContext>(context: TContext) -> Field where TContext: ContextInterface {
-    context.this_address().to_field()
+fn compute_unsiloed_contract_initialization_nullifier(address: AztecAddress) -> Field {
+    address.to_field()
+}
+
+pub fn assert_initialization_matches_address_preimage_public(context: PublicContext) {
+    assert_initialization_matches_address_preimage(context);
+}
+
+pub fn assert_initialization_matches_address_preimage_avm(context: AvmContext) {
+    let address = context.this_address();
+    let instance = get_contract_instance_avm(address).unwrap();
+    let expected_init = compute_initialization_hash(context.selector(), context.get_args_hash());
+    assert(instance.initialization_hash == expected_init, "Initialization hash does not match");
+    assert(
+        (instance.deployer.is_zero()) | (instance.deployer == context.msg_sender()), "Initializer address is not the contract deployer"
+    );
+}
+
+pub fn assert_initialization_matches_address_preimage_private(context: PrivateContext) {
+    assert_initialization_matches_address_preimage(context);
 }
 
-pub fn assert_initialization_matches_address_preimage<TContext>(context: TContext) where TContext: ContextInterface {
+fn assert_initialization_matches_address_preimage<TContext>(context: TContext) where TContext: ContextInterface {
     let address = context.this_address();
     let instance = get_contract_instance(address);
     let expected_init = compute_initialization_hash(context.selector(), context.get_args_hash());

noir-projects/noir-contracts/Nargo.toml

@@ -1,6 +1,7 @@
 [workspace]
 members = [
     "contracts/app_subscription_contract",
+    "contracts/avm_initializer_test_contract",
     "contracts/avm_test_contract",
     "contracts/fpc_contract",
     "contracts/benchmarking_contract",

noir-projects/noir-contracts/contracts/avm_initializer_test_contract/Nargo.toml

@@ -0,0 +1,8 @@
+[package]
+name = "avm_initializer_test_contract"
+authors = [""]
+compiler_version = ">=0.25.0"
+type = "contract"
+
+[dependencies]
+aztec = { path = "../../../aztec-nr/aztec" }

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

@@ -0,0 +1,22 @@
+contract AvmInitializerTest {
+    // Libs
+    use dep::aztec::state_vars::PublicImmutable;
+    use dep::aztec::protocol_types::address::AztecAddress;
+
+    struct Storage {
+        immutable: PublicImmutable<Field>,
+    }
+
+    /************************************************************************
+     * Storage
+     ************************************************************************/
+    #[aztec(public-vm)]
+    #[aztec(initializer)]
+    fn constructor() {
+        storage.immutable.initialize(42);
+    }
+
+    unconstrained fn view_storage_immutable() -> pub Field {
+        storage.immutable.read()
+    }
+}

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::{PublicImmutable, PublicMutable};
+    use dep::aztec::state_vars::PublicMutable;
     use dep::aztec::protocol_types::{
         address::{AztecAddress, EthAddress}, constants::L1_TO_L2_MESSAGE_LENGTH,
         contract_instance::ContractInstance
@@ -43,23 +43,11 @@ 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-projects/noir-contracts/contracts/stateful_test_contract/src/main.nr

@@ -4,7 +4,7 @@ contract StatefulTest {
     use dep::value_note::{balance_utils, utils::{increment, decrement}, value_note::{VALUE_NOTE_LEN, ValueNote}};
     use dep::aztec::{
         deploy::deploy_contract as aztec_deploy_contract, context::{PublicContext, Context},
-        oracle::get_contract_instance::get_contract_instance, initializer::assert_is_initialized
+        oracle::get_contract_instance::get_contract_instance, initializer::assert_is_initialized_private
     };
 
     struct Storage {
@@ -45,7 +45,7 @@ contract StatefulTest {
 
     #[aztec(private)]
     fn destroy_and_create(recipient: AztecAddress, amount: Field) {
-        assert_is_initialized(&mut context);
+        assert_is_initialized_private(&mut context);
         let sender = context.msg_sender();
 
         let sender_notes = storage.notes.at(sender);

noir/noir-repo/aztec_macros/src/transforms/functions.rs

@@ -45,13 +45,13 @@ pub fn transform_function(
 
     // Add initialization check
     if insert_init_check {
-        let init_check = create_init_check();
+        let init_check = create_init_check(ty);
         func.def.body.statements.insert(0, init_check);
     }
 
     // Add assertion for initialization arguments and sender
     if is_initializer {
-        func.def.body.statements.insert(0, create_assert_initializer());
+        func.def.body.statements.insert(0, create_assert_initializer(ty));
     }
 
     // Add access to the storage struct
@@ -85,7 +85,7 @@ pub fn transform_function(
 
     // Before returning mark the contract as initialized
     if is_initializer {
-        let mark_initialized = create_mark_as_initialized();
+        let mark_initialized = create_mark_as_initialized(ty);
         func.def.body.statements.push(mark_initialized);
     }
 
@@ -159,9 +159,10 @@ fn create_inputs(ty: &str) -> Param {
 /// ```noir
 /// assert_is_initialized(&mut context);
 /// ```
-fn create_init_check() -> Statement {
+fn create_init_check(ty: &str) -> Statement {
+    let fname = format!("assert_is_initialized_{}", ty.to_case(Case::Snake));
     make_statement(StatementKind::Expression(call(
-        variable_path(chained_dep!("aztec", "initializer", "assert_is_initialized")),
+        variable_path(chained_dep!("aztec", "initializer", &fname)),
         vec![mutable_reference("context")],
     )))
 }
@@ -172,9 +173,10 @@ fn create_init_check() -> Statement {
 /// ```noir
 /// mark_as_initialized(&mut context);
 /// ```
-fn create_mark_as_initialized() -> Statement {
+fn create_mark_as_initialized(ty: &str) -> Statement {
+    let fname = format!("mark_as_initialized_{}", ty.to_case(Case::Snake));
     make_statement(StatementKind::Expression(call(
-        variable_path(chained_dep!("aztec", "initializer", "mark_as_initialized")),
+        variable_path(chained_dep!("aztec", "initializer", &fname)),
         vec![mutable_reference("context")],
     )))
 }
@@ -205,13 +207,11 @@ fn create_internal_check(fname: &str) -> Statement {
 /// ```noir
 /// assert_initialization_matches_address_preimage(context);
 /// ```
-fn create_assert_initializer() -> Statement {
+fn create_assert_initializer(ty: &str) -> Statement {
+    let fname =
+        format!("assert_initialization_matches_address_preimage_{}", ty.to_case(Case::Snake));
     make_statement(StatementKind::Expression(call(
-        variable_path(chained_dep!(
-            "aztec",
-            "initializer",
-            "assert_initialization_matches_address_preimage"
-        )),
+        variable_path(chained_dep!("aztec", "initializer", &fname)),
         vec![variable("context")],
     )))
 }

yarn-project/cli/src/cmds/example_contracts.ts

@@ -4,6 +4,6 @@ import { getExampleContractArtifacts } from '../utils.js';
 
 export async function exampleContracts(log: LogFn) {
   const abisList = await getExampleContractArtifacts();
-  const names = Object.keys(abisList).filter(name => name !== 'AvmTestContractArtifact');
+  const names = Object.keys(abisList).filter(name => !name.startsWith('Avm'));
   names.forEach(name => log(name));
 }

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

@@ -0,0 +1,32 @@
+import { type Wallet } from '@aztec/aztec.js';
+import { AvmInitializerTestContract } from '@aztec/noir-contracts.js';
+
+import { jest } from '@jest/globals';
+
+import { setup } from './fixtures/utils.js';
+
+const TIMEOUT = 100_000;
+
+describe('e2e_avm_initializer', () => {
+  jest.setTimeout(TIMEOUT);
+
+  let wallet: Wallet;
+  let avmContact: AvmInitializerTestContract;
+  let teardown: () => Promise<void>;
+
+  beforeAll(async () => {
+    ({ teardown, wallet } = await setup());
+  }, 100_000);
+
+  afterAll(() => teardown());
+
+  beforeEach(async () => {
+    avmContact = await AvmInitializerTestContract.deploy(wallet).send().deployed();
+  }, 50_000);
+
+  describe('Storage', () => {
+    it('Read immutable (initialized) storage (Field)', async () => {
+      expect(await avmContact.methods.view_storage_immutable().simulate()).toEqual(42n);
+    });
+  });
+});

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

@@ -25,11 +25,6 @@ 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().simulate()).toEqual(42n);
-    // });
-
     it('Modifies storage (Field)', async () => {
       await avmContact.methods.set_storage_single(20n).send().wait();
       expect(await avmContact.methods.view_storage_single().simulate()).toEqual(20n);

yarn-project/simulator/src/public/executor.ts

@@ -256,13 +256,14 @@ export class PublicExecutor {
     const worldStateJournal = new AvmPersistableStateManager(hostStorage);
     const executionEnv = temporaryCreateAvmExecutionEnvironment(execution, globalVariables);
     // TODO(@spalladino) Load initial gas from the public execution request
-    const machineState = new AvmMachineState(1e6, 1e6, 1e6);
+    const machineState = new AvmMachineState(1e10, 1e10, 1e10);
 
     const context = new AvmContext(worldStateJournal, executionEnv, machineState);
     const simulator = new AvmSimulator(context);
 
     const result = await simulator.execute();
     const newWorldState = context.persistableState.flush();
+
     // TODO(@spalladino) Read gas left from machineState and return it
     return temporaryConvertAvmResults(execution, newWorldState, result);
   }