feat!: merge the txe session state transition oracles, redo the internal wiring, introduce context handlers (#16659)

This one leaves things in a bit of a messy state, as doing _the entire_
transition is unfeasiable, but the things that are ready _do_ look good,
and it's a good starting place for future changes.

The session now calls the translator (service) with a handler for state
changes (the session itself) and a handler for all other calls. On
session state transitions, a new handler is installed and the old one
closed and committed to. There's now a handler for `public_context`
which does the avm opcodes, and consequently the `txe_oracle` no longer
knows how to process these (because it doesn't need to).

I also cleaned up the state transition oracles to take all params in one
go instead of relying on mutating global session state, which helps make
this smoother and also helps with the cleanup and eventual removal of
the `txe_oracle`.

Author: AztecBot

noir-projects/aztec-nr/aztec/src/test/helpers/test_environment.nr

@@ -149,13 +149,11 @@ impl TestEnvironment {
         _self: Self,
         f: fn[Env](&mut PublicContext) -> T,
     ) -> T {
-        txe_oracles::set_public_txe_context();
+        txe_oracles::set_public_txe_context(Option::none());
 
         let mut context = PublicContext::new(|| 0);
         let ret_value = f(&mut context);
 
-        // merge calls below into 'exit public txe context' -> this mines a block
-        txe_oracles::advance_blocks_by(1);
         txe_oracles::set_top_level_txe_context();
 
         ret_value
@@ -164,17 +162,18 @@ impl TestEnvironment {
     /// Variant of `public_context` which allows specifying the contract address in which the public context will
     /// execute, which will affect note and nullifier siloing, storage access, etc.
     pub unconstrained fn public_context_at<Env, T>(
-        self,
+        _self: Self,
         addr: AztecAddress,
         f: fn[Env](&mut PublicContext) -> T,
     ) -> T {
-        // temporary hack until we reimplement the public context oracle and have it take this as a param
-        let pre = txe_oracles::get_contract_address();
-        txe_oracles::set_contract_address(addr);
-        let ret = self.public_context(f);
-        txe_oracles::set_contract_address(pre);
+        txe_oracles::set_public_txe_context(Option::some(addr));
 
-        ret
+        let mut context = PublicContext::new(|| 0);
+        let ret_value = f(&mut context);
+
+        txe_oracles::set_top_level_txe_context();
+
+        ret_value
     }
 
     /// Creates a `PrivateContext`, which allows using aztec-nr features as if inside a private contract function.
@@ -210,76 +209,58 @@ impl TestEnvironment {
     ///   state_var.get_note()
     /// });
     /// ```
-    pub unconstrained fn private_context<Env, T>(self, f: fn[Env](&mut PrivateContext) -> T) -> T {
-        // merge the two oracles below - the creation of a context implies a transition to a private txe context. right
-        // now we have them in this order because creating the private context requires reading the latest block number,
-        // which we only allow for top level and utility
+    pub unconstrained fn private_context<Env, T>(
+        _self: Self,
+        f: fn[Env](&mut PrivateContext) -> T,
+    ) -> T {
         let mut context = PrivateContext::new(
-            txe_oracles::get_private_context_inputs(Option::some(self.last_block_number())),
+            txe_oracles::set_private_txe_context(Option::none(), Option::none()),
             0,
         );
-        txe_oracles::set_private_txe_context();
 
         let ret_value = f(&mut context);
-        txe_oracles::set_top_level_txe_context();
 
-        // todo: should commit the context to mine a block with the side effects of the context. we should have an
-        // oracle that receives the context we produced probably
-        self.mine_block();
+        txe_oracles::set_top_level_txe_context();
 
         ret_value
     }
 
     /// Variant of `private_context` which allows specifying the contract address in which the private context will
     /// execute, which will affect note and nullifier siloing, storage access, etc.
     pub unconstrained fn private_context_at<Env, T>(
-        self,
+        _self: Self,
         addr: AztecAddress,
         f: fn[Env](&mut PrivateContext) -> T,
     ) -> T {
-        // temporary hack until we reimplement the public context oracle and have it take this as a param
-        let pre = txe_oracles::get_contract_address();
-        txe_oracles::set_contract_address(addr);
-        let ret = self.private_context(f);
-        txe_oracles::set_contract_address(pre);
+        let mut context = PrivateContext::new(
+            txe_oracles::set_private_txe_context(Option::some(addr), Option::none()),
+            0,
+        );
 
-        ret
+        let ret_value = f(&mut context);
+
+        txe_oracles::set_top_level_txe_context();
+
+        ret_value
     }
 
     /// Variant of `private_context` which allows specifying multiple configuration values via `PrivateContextOptions`.
     pub unconstrained fn private_context_opts<Env, T>(
-        self,
+        _self: Self,
         opts: PrivateContextOptions,
         f: fn[Env](&mut PrivateContext) -> T,
     ) -> T {
-        // temporary hack until we reimplement the public context oracle and have it take this as a param
-
-        let pre_contract_address = txe_oracles::get_contract_address();
-
-        if opts.contract_address.is_some() {
-            txe_oracles::set_contract_address(opts.contract_address.unwrap());
-        }
-
-        // merge the two oracles below - the creation of a context implies a transition to a private txe context. right
-        // now we have them in this order because creating the private context requires reading the latest block number,
-        // which we only allow for top level and utility
         let mut context = PrivateContext::new(
-            txe_oracles::get_private_context_inputs(opts.historical_block_number.or_else(|| {
-                Option::some(self.last_block_number())
-            })),
+            txe_oracles::set_private_txe_context(
+                opts.contract_address,
+                opts.historical_block_number,
+            ),
             0,
         );
 
-        txe_oracles::set_private_txe_context();
-
         let ret_value = f(&mut context);
-        txe_oracles::set_top_level_txe_context();
-
-        // todo: should commit the context to mine a block with the side effects of the context. we should have an
-        // oracle that receives the context we produced probably
-        self.mine_block();
 
-        txe_oracles::set_contract_address(pre_contract_address);
+        txe_oracles::set_top_level_txe_context();
 
         ret_value
     }
@@ -314,7 +295,7 @@ impl TestEnvironment {
         _self: Self,
         f: fn[Env](UtilityContext) -> T,
     ) -> T {
-        txe_oracles::set_utility_txe_context();
+        txe_oracles::set_utility_txe_context(Option::none());
         let context = UtilityContext::new();
         let ret_value = f(context);
         txe_oracles::set_top_level_txe_context();
@@ -325,17 +306,16 @@ impl TestEnvironment {
     /// Variant of `utility_context` which allows specifying the contract address in which the utility context will
     /// execute, which will affect note and storage access.
     pub unconstrained fn utility_context_at<Env, T>(
-        self,
+        _self: Self,
         addr: AztecAddress,
         f: fn[Env](UtilityContext) -> T,
     ) -> T {
-        // temporary hack until we reimplement the utility context oracle and have it take this as a param
-        let pre = txe_oracles::get_contract_address();
-        txe_oracles::set_contract_address(addr);
-        let ret = self.utility_context(f);
-        txe_oracles::set_contract_address(pre);
+        txe_oracles::set_utility_txe_context(Option::some(addr));
+        let context = UtilityContext::new();
+        let ret_value = f(context);
+        txe_oracles::set_top_level_txe_context();
 
-        ret
+        ret_value
     }
 
     /// Returns the number of the next block to be mined.

noir-projects/aztec-nr/aztec/src/test/helpers/txe_oracles.nr

@@ -65,23 +65,12 @@ pub unconstrained fn public_call_new_flow(
 #[oracle(txeGetLastBlockTimestamp)]
 pub unconstrained fn get_last_block_timestamp() -> u64 {}
 
-#[oracle(txeSetContractAddress)]
-pub unconstrained fn set_contract_address(address: AztecAddress) {}
-
-#[oracle(utilityGetContractAddress)]
-pub unconstrained fn get_contract_address() -> AztecAddress {}
-
 #[oracle(txeAdvanceBlocksBy)]
 pub unconstrained fn advance_blocks_by(blocks: u32) {}
 
 #[oracle(txeAdvanceTimestampBy)]
 pub unconstrained fn advance_timestamp_by(duration: u64) {}
 
-#[oracle(txeGetPrivateContextInputs)]
-pub unconstrained fn get_private_context_inputs(
-    historical_block_number: Option<u32>,
-) -> PrivateContextInputs {}
-
 #[oracle(txeDeploy)]
 pub unconstrained fn deploy_oracle<let N: u32, let P: u32>(
     path: str<N>,
@@ -121,10 +110,13 @@ unconstrained fn public_call_new_flow_oracle(
 pub unconstrained fn set_top_level_txe_context() {}
 
 #[oracle(txeSetPrivateTXEContext)]
-pub unconstrained fn set_private_txe_context() {}
+pub unconstrained fn set_private_txe_context(
+    contract_address: Option<AztecAddress>,
+    historical_block_number: Option<u32>,
+) -> PrivateContextInputs {}
 
 #[oracle(txeSetPublicTXEContext)]
-pub unconstrained fn set_public_txe_context() {}
+pub unconstrained fn set_public_txe_context(contract_address: Option<AztecAddress>) {}
 
 #[oracle(txeSetUtilityTXEContext)]
-pub unconstrained fn set_utility_txe_context() {}
+pub unconstrained fn set_utility_txe_context(contract_address: Option<AztecAddress>) {}

noir-projects/aztec-nr/aztec/src/utils/with_hash.nr

@@ -120,7 +120,6 @@ where
 
 mod test {
     use crate::{
-        oracle::random::random,
         test::{helpers::test_environment::TestEnvironment, mocks::mock_struct::MockStruct},
         utils::with_hash::WithHash,
     };
@@ -211,7 +210,7 @@ mod test {
             let value = MockStruct { a: 5, b: 3 };
             context.storage_write(STORAGE_SLOT, value);
 
-            let incorrect_hash = random();
+            let incorrect_hash = 13;
             let hash_storage_slot = STORAGE_SLOT + (value.pack().len() as Field);
             context.storage_write(hash_storage_slot, [incorrect_hash]);
         });

yarn-project/txe/src/index.ts

@@ -55,6 +55,7 @@ type TXEForeignCallInput = {
   inputs: ForeignCallArgs;
 };
 
+// TODO: why does the zod validation in the txe dispatcher not reject invalid function names?
 const TXEForeignCallInputSchema = z.object({
   // eslint-disable-next-line camelcase
   session_id: z.number().int().nonnegative(),