feat(avm): partially integrate with initializers

AztecProtocol · Mar 26, 2024 · 074d44b · 074d44b
1 parent 1b390fb
commit 074d44b
Show file tree

Hide file tree

Showing 6 changed files with 85 additions and 48 deletions.
diff --git a/noir-projects/aztec-nr/aztec/src/initializer.nr b/noir-projects/aztec-nr/aztec/src/initializer.nr
@@ -1,45 +1,83 @@
 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,
-    history::nullifier_inclusion::prove_nullifier_inclusion,
+    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) {
+    let init_nullifier = compute_unsiloed_contract_initialization_nullifier(context.this_address());
+    context.push_new_nullifier(init_nullifier, 0);
+}
+
+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) {
+    let init_nullifier = compute_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) {
+    // FIXME(https://github.com/AztecProtocol/aztec-packages/issues/5463): getContractInstance not supported in AVM.
 }
 
-pub fn assert_initialization_matches_address_preimage<TContext>(context: TContext) where TContext: ContextInterface {
-    let address = context.this_address(); 
+pub fn assert_initialization_matches_address_preimage_private(context: PrivateContext) {
+    assert_initialization_matches_address_preimage(context);
+}
+
+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());
     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");
+    assert(
+        (instance.deployer.is_zero()) | (instance.deployer == context.msg_sender()), "Initializer address is not the contract deployer"
+    );
 }
 
 pub fn compute_initialization_hash(init_selector: FunctionSelector, init_args_hash: Field) -> Field {
-    pedersen_hash([init_selector.to_field(), init_args_hash], GENERATOR_INDEX__CONSTRUCTOR)
-}
+    pedersen_hash(
+        [init_selector.to_field(), init_args_hash],
+        GENERATOR_INDEX__CONSTRUCTOR
+    )
+}
diff --git a/noir-projects/noir-contracts/contracts/avm_test_contract/src/main.nr b/noir-projects/noir-contracts/contracts/avm_test_contract/src/main.nr
@@ -45,12 +45,11 @@ contract AvmTest {
     /************************************************************************
      * Storage
      ************************************************************************/
-    // FIX: calls unsupported getNullifierMembershipWitness.
-    // #[aztec(public-vm)]
-    // #[aztec(initializer)]
-    // fn constructor() {
-    //     storage.immutable.initialize(42);
-    // }
+    #[aztec(public-vm)]
+    #[aztec(initializer)]
+    fn constructor() {
+        storage.immutable.initialize(42);
+    }
 
     unconstrained fn view_storage_immutable() -> pub Field {
         storage.immutable.read()

diff --git a/noir-projects/noir-contracts/contracts/stateful_test_contract/src/main.nr b/noir-projects/noir-contracts/contracts/stateful_test_contract/src/main.nr
@@ -3,9 +3,8 @@ contract StatefulTest {
     use dep::aztec::prelude::{PrivateContext, NoteHeader, Map, PublicMutable, PrivateSet, AztecAddress, FunctionSelector};
     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,
+        deploy::deploy_contract as aztec_deploy_contract, context::{PublicContext, Context},
+        oracle::get_contract_instance::get_contract_instance, initializer::assert_is_initialized_private
     };
 
     struct Storage {
@@ -46,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);

diff --git a/noir/noir-repo/aztec_macros/src/transforms/functions.rs b/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.0.insert(0, init_check);
     }
 
     // Add assertion for initialization arguments and sender
     if is_initializer {
-        func.def.body.0.insert(0, create_assert_initializer());
+        func.def.body.0.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.0.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")],
     )))
 }

diff --git a/yarn-project/end-to-end/src/e2e_avm_simulator.test.ts b/yarn-project/end-to-end/src/e2e_avm_simulator.test.ts
@@ -25,10 +25,9 @@ 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('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();

diff --git a/yarn-project/simulator/src/avm/temporary_executor_migration.ts b/yarn-project/simulator/src/avm/temporary_executor_migration.ts
@@ -95,8 +95,10 @@ export function temporaryConvertAvmResults(
   const nestedExecutions: PublicExecutionResult[] = [];
   const nullifierReadRequests: ReadRequest[] = [];
   const nullifierNonExistentReadRequests: ReadRequest[] = [];
-  const newNullifiers: SideEffectLinkedToNoteHash[] = [];
-  const unencryptedLogs = FunctionL2Logs.empty();
+  const newNullifiers: SideEffectLinkedToNoteHash[] = newWorldState.newNullifiers.map(
+    nullifier => new SideEffectLinkedToNoteHash(Fr.zero(), nullifier.toField(), Fr.zero()),
+  );
+  const unencryptedLogs = new FunctionL2Logs(newWorldState.newLogs.map(log => log.data));
   const newL2ToL1Messages = newWorldState.newL1Messages.map(() => L2ToL1Message.empty());
   // TODO keep track of side effect counters
   const startSideEffectCounter = Fr.ZERO;