fix: Quick TXE after public executor changes (#8661)

This PR https://github.com/AztecProtocol/aztec-packages/pull/8585/files
made it so the TXE uses a world state db. But a fresh one is created for
every call so all the changes uncommited to the tree were lost.

---------

Co-authored-by: ludamad <[email protected]>

.github/workflows/ci.yml

@@ -55,6 +55,7 @@ jobs:
       barretenberg-cpp: ${{ steps.filter.outputs.barretenberg-cpp }}
       noir: ${{ steps.filter.outputs.noir }}
       noir-projects: ${{ steps.filter.outputs.noir-projects }}
+      txe: ${{ steps.filter.outputs.txe }}
       l1-contracts: ${{ steps.filter.outputs.l1-contracts }}
       non-docs: ${{ steps.filter.outputs.non-docs }}
       non-misc-ci: ${{ steps.filter.outputs.non-misc-ci }}

yarn-project/txe/src/oracle/txe_oracle.ts

@@ -59,7 +59,6 @@ import {
   type PackedValuesCache,
   PublicExecutor,
   type TypedOracle,
-  WorldStateDB,
   acvm,
   createSimulationError,
   extractCallStack,
@@ -73,6 +72,7 @@ import { MerkleTreeSnapshotOperationsFacade, type MerkleTrees } from '@aztec/wor
 
 import { type TXEDatabase } from '../util/txe_database.js';
 import { TXEPublicContractDataSource } from '../util/txe_public_contract_data_source.js';
+import { TXEWorldStateDB } from '../util/txe_world_state_db.js';
 
 export class TXE implements TypedOracle {
   private blockNumber = 0;
@@ -717,7 +717,7 @@ export class TXE implements TypedOracle {
     header.globalVariables.blockNumber = new Fr(await this.getBlockNumber());
 
     const executor = new PublicExecutor(
-      new WorldStateDB(this.trees.asLatest(), new TXEPublicContractDataSource(this)),
+      new TXEWorldStateDB(this.trees.asLatest(), new TXEPublicContractDataSource(this)),
       header,
       new NoopTelemetryClient(),
     );

yarn-project/txe/src/util/txe_world_state_db.ts

@@ -0,0 +1,55 @@
+import { MerkleTreeId, type MerkleTreeOperations } from '@aztec/circuit-types';
+import {
+  type AztecAddress,
+  Fr,
+  PUBLIC_DATA_SUBTREE_HEIGHT,
+  PublicDataTreeLeaf,
+  type PublicDataTreeLeafPreimage,
+} from '@aztec/circuits.js';
+import { computePublicDataTreeLeafSlot } from '@aztec/circuits.js/hash';
+import { WorldStateDB } from '@aztec/simulator';
+import { type ContractDataSource } from '@aztec/types/contracts';
+
+export class TXEWorldStateDB extends WorldStateDB {
+  constructor(private merkleDb: MerkleTreeOperations, dataSource: ContractDataSource) {
+    super(merkleDb, dataSource);
+  }
+
+  override async storageRead(contract: AztecAddress, slot: Fr): Promise<Fr> {
+    const leafSlot = computePublicDataTreeLeafSlot(contract, slot).toBigInt();
+
+    const lowLeafResult = await this.merkleDb.getPreviousValueIndex(MerkleTreeId.PUBLIC_DATA_TREE, leafSlot);
+
+    let value = Fr.ZERO;
+    if (lowLeafResult && lowLeafResult.alreadyPresent) {
+      const preimage = (await this.merkleDb.getLeafPreimage(
+        MerkleTreeId.PUBLIC_DATA_TREE,
+        lowLeafResult.index,
+      )) as PublicDataTreeLeafPreimage;
+      value = preimage.value;
+    }
+    return value;
+  }
+
+  override async storageWrite(contract: AztecAddress, slot: Fr, newValue: Fr): Promise<bigint> {
+    await this.merkleDb.batchInsert(
+      MerkleTreeId.PUBLIC_DATA_TREE,
+      [new PublicDataTreeLeaf(computePublicDataTreeLeafSlot(contract, slot), newValue).toBuffer()],
+      PUBLIC_DATA_SUBTREE_HEIGHT,
+    );
+    return newValue.toBigInt();
+  }
+
+  override checkpoint(): Promise<void> {
+    return Promise.resolve();
+  }
+  override rollbackToCheckpoint(): Promise<void> {
+    throw new Error('Cannot rollback');
+  }
+  override commit(): Promise<void> {
+    return Promise.resolve();
+  }
+  override rollbackToCommit(): Promise<void> {
+    throw new Error('Cannot rollback');
+  }
+}