feat: Rollback public state changes on failure (#3393)

This PR rolls back public state changes when public execution reverts.

# Checklist:
Remove the checklist to signal you've completed it. Enable auto-merge if
the PR is ready to merge.
- [ ] If the pull request requires a cryptography review (e.g.
cryptographic algorithm implementations) I have added the 'crypto' tag.
- [ ] I have reviewed my diff in github, line by line and removed
unexpected formatting changes, testing logs, or commented-out code.
- [ ] Every change is related to the PR description.
- [ ] I have
[linked](https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue)
this pull request to relevant issues (if any exist).

yarn-project/acir-simulator/src/public/db.ts

@@ -24,6 +24,18 @@ export interface PublicStateDB {
    * @returns Nothing.
    */
   storageWrite(contract: AztecAddress, slot: Fr, newValue: Fr): Promise<void>;
+
+  /**
+   * Commit the pending changes to the DB.
+   * @returns Nothing.
+   */
+  commit(): Promise<void>;
+
+  /**
+   * Rollback the pending changes.
+   * @returns Nothing.
+   */
+  rollback(): Promise<void>;
 }
 
 /**

yarn-project/sequencer-client/src/sequencer/public_processor.test.ts

@@ -26,15 +26,24 @@ import {
   makeSelector,
 } from '@aztec/circuits.js/factories';
 import { padArrayEnd } from '@aztec/foundation/collection';
-import { ExtendedContractData, FunctionCall, FunctionL2Logs, SiblingPath, Tx, TxL2Logs, mockTx } from '@aztec/types';
+import {
+  ExtendedContractData,
+  FunctionCall,
+  FunctionL2Logs,
+  SiblingPath,
+  SimulationError,
+  Tx,
+  TxL2Logs,
+  mockTx,
+} from '@aztec/types';
 import { MerkleTreeOperations, TreeInfo } from '@aztec/world-state';
 
 import { MockProxy, mock } from 'jest-mock-extended';
 import times from 'lodash.times';
 
 import { PublicProver } from '../prover/index.js';
 import { PublicKernelCircuitSimulator } from '../simulator/index.js';
-import { ContractsDataSourcePublicDB } from '../simulator/public_executor.js';
+import { ContractsDataSourcePublicDB, WorldStatePublicDB } from '../simulator/public_executor.js';
 import { WasmPublicKernelCircuitSimulator } from '../simulator/public_kernel.js';
 import { PublicProcessor } from './public_processor.js';
 
@@ -43,6 +52,7 @@ describe('public_processor', () => {
   let publicExecutor: MockProxy<PublicExecutor>;
   let publicProver: MockProxy<PublicProver>;
   let publicContractsDB: MockProxy<ContractsDataSourcePublicDB>;
+  let publicWorldStateDB: MockProxy<WorldStatePublicDB>;
 
   let proof: Proof;
   let root: Buffer;
@@ -54,6 +64,7 @@ describe('public_processor', () => {
     publicExecutor = mock<PublicExecutor>();
     publicProver = mock<PublicProver>();
     publicContractsDB = mock<ContractsDataSourcePublicDB>();
+    publicWorldStateDB = mock<WorldStatePublicDB>();
 
     proof = makeEmptyProof();
     root = Buffer.alloc(32, 5);
@@ -76,6 +87,7 @@ describe('public_processor', () => {
         GlobalVariables.empty(),
         HistoricBlockData.empty(),
         publicContractsDB,
+        publicWorldStateDB,
       );
     });
 
@@ -110,6 +122,8 @@ describe('public_processor', () => {
 
       expect(processed).toEqual([]);
       expect(failed[0].tx).toEqual(tx);
+      expect(publicWorldStateDB.commit).toHaveBeenCalledTimes(0);
+      expect(publicWorldStateDB.rollback).toHaveBeenCalledTimes(1);
     });
   });
 
@@ -128,6 +142,7 @@ describe('public_processor', () => {
         GlobalVariables.empty(),
         HistoricBlockData.empty(),
         publicContractsDB,
+        publicWorldStateDB,
       );
     });
 
@@ -165,6 +180,8 @@ describe('public_processor', () => {
       expect(processed).toEqual([await expectedTxByHash(tx)]);
       expect(failed).toHaveLength(0);
       expect(publicExecutor.simulate).toHaveBeenCalledTimes(2);
+      expect(publicWorldStateDB.commit).toHaveBeenCalledTimes(1);
+      expect(publicWorldStateDB.rollback).toHaveBeenCalledTimes(0);
     });
 
     it('runs a tx with an enqueued public call with nested execution', async function () {
@@ -201,6 +218,45 @@ describe('public_processor', () => {
       expect(processed).toEqual([await expectedTxByHash(tx)]);
       expect(failed).toHaveLength(0);
       expect(publicExecutor.simulate).toHaveBeenCalledTimes(1);
+      expect(publicWorldStateDB.commit).toHaveBeenCalledTimes(1);
+      expect(publicWorldStateDB.rollback).toHaveBeenCalledTimes(0);
+    });
+
+    it('rolls back db updates on failed public execution', async function () {
+      const callRequest: PublicCallRequest = makePublicCallRequest(0x100);
+      const callStackItem = callRequest.toPublicCallStackItem();
+      const callStackHash = computeCallStackItemHash(callStackItem);
+
+      const kernelOutput = makePrivateKernelPublicInputsFinal(0x10);
+      kernelOutput.end.publicCallStack = padArrayEnd([callStackHash], Fr.ZERO, MAX_PUBLIC_CALL_STACK_LENGTH_PER_TX);
+      kernelOutput.end.privateCallStack = padArrayEnd([], Fr.ZERO, MAX_PRIVATE_CALL_STACK_LENGTH_PER_TX);
+
+      const tx = new Tx(
+        kernelOutput,
+        proof,
+        TxL2Logs.random(2, 3),
+        TxL2Logs.random(3, 2),
+        [callRequest],
+        [ExtendedContractData.random()],
+      );
+
+      const publicExecutionResult = makePublicExecutionResultFromRequest(callRequest);
+      publicExecutionResult.nestedExecutions = [
+        makePublicExecutionResult(publicExecutionResult.execution.contractAddress, {
+          to: makeAztecAddress(30),
+          functionData: new FunctionData(makeSelector(5), false, false, false),
+          args: new Array(ARGS_LENGTH).fill(Fr.ZERO),
+        }),
+      ];
+      publicExecutor.simulate.mockRejectedValueOnce(new SimulationError('Simulation Failed', []));
+
+      const [processed, failed] = await processor.process([tx]);
+
+      expect(failed).toHaveLength(1);
+      expect(processed).toHaveLength(0);
+      expect(publicExecutor.simulate).toHaveBeenCalledTimes(1);
+      expect(publicWorldStateDB.rollback).toHaveBeenCalledTimes(1);
+      expect(publicWorldStateDB.commit).toHaveBeenCalledTimes(0);
     });
   });
 });

yarn-project/sequencer-client/src/sequencer/public_processor.ts

@@ -2,6 +2,7 @@ import {
   PublicExecution,
   PublicExecutionResult,
   PublicExecutor,
+  PublicStateDB,
   collectPublicDataReads,
   collectPublicDataUpdateRequests,
   isPublicExecutionResult,
@@ -47,7 +48,7 @@ import { getVerificationKeys } from '../index.js';
 import { EmptyPublicProver } from '../prover/empty.js';
 import { PublicProver } from '../prover/index.js';
 import { PublicKernelCircuitSimulator } from '../simulator/index.js';
-import { ContractsDataSourcePublicDB, getPublicExecutor } from '../simulator/public_executor.js';
+import { ContractsDataSourcePublicDB, WorldStateDB, WorldStatePublicDB } from '../simulator/public_executor.js';
 import { WasmPublicKernelCircuitSimulator } from '../simulator/public_kernel.js';
 import { FailedTx, ProcessedTx, makeEmptyProcessedTx, makeProcessedTx } from './processed_tx.js';
 import { getHistoricBlockData } from './utils.js';
@@ -75,14 +76,18 @@ export class PublicProcessorFactory {
   ): Promise<PublicProcessor> {
     const blockData = await getHistoricBlockData(this.merkleTree, prevGlobalVariables);
     const publicContractsDB = new ContractsDataSourcePublicDB(this.contractDataSource);
+    const worldStatePublicDB = new WorldStatePublicDB(this.merkleTree);
+    const worldStateDB = new WorldStateDB(this.merkleTree, this.l1Tol2MessagesDataSource);
+    const publicExecutor = new PublicExecutor(worldStatePublicDB, publicContractsDB, worldStateDB, blockData);
     return new PublicProcessor(
       this.merkleTree,
-      getPublicExecutor(this.merkleTree, publicContractsDB, this.l1Tol2MessagesDataSource, blockData),
+      publicExecutor,
       new WasmPublicKernelCircuitSimulator(),
       new EmptyPublicProver(),
       globalVariables,
       blockData,
       publicContractsDB,
+      worldStatePublicDB,
     );
   }
 }
@@ -100,6 +105,7 @@ export class PublicProcessor {
     protected globalVariables: GlobalVariables,
     protected blockData: HistoricBlockData,
     protected publicContractsDB: ContractsDataSourcePublicDB,
+    protected publicStateDB: PublicStateDB,
 
     private log = createDebugLogger('aztec:sequencer:public-processor'),
   ) {}
@@ -121,6 +127,8 @@ export class PublicProcessor {
         // add new contracts to the contracts db so that their functions may be found and called
         await this.publicContractsDB.addNewContracts(tx);
         result.push(await this.processTx(tx));
+        // commit the state updates from this transaction
+        await this.publicStateDB.commit();
       } catch (err) {
         this.log.warn(`Error processing tx ${await tx.getTxHash()}: ${err}`);
         failed.push({
@@ -129,6 +137,8 @@ export class PublicProcessor {
         });
         // remove contracts on failure
         await this.publicContractsDB.removeNewContracts(tx);
+        // rollback any state updates from this failed transaction
+        await this.publicStateDB.rollback();
       }
     }
 

yarn-project/sequencer-client/src/simulator/public_executor.ts

@@ -1,35 +1,9 @@
-import {
-  CommitmentsDB,
-  MessageLoadOracleInputs,
-  PublicContractsDB,
-  PublicExecutor,
-  PublicStateDB,
-} from '@aztec/acir-simulator';
-import { AztecAddress, EthAddress, Fr, FunctionSelector, HistoricBlockData } from '@aztec/circuits.js';
+import { CommitmentsDB, MessageLoadOracleInputs, PublicContractsDB, PublicStateDB } from '@aztec/acir-simulator';
+import { AztecAddress, EthAddress, Fr, FunctionSelector } from '@aztec/circuits.js';
 import { computePublicDataTreeIndex } from '@aztec/circuits.js/abis';
 import { ContractDataSource, ExtendedContractData, L1ToL2MessageSource, MerkleTreeId, Tx } from '@aztec/types';
 import { MerkleTreeOperations } from '@aztec/world-state';
 
-/**
- * Returns a new PublicExecutor simulator backed by the supplied merkle tree db and contract data source.
- * @param merkleTree - A merkle tree database.
- * @param contractDataSource - A contract data source.
- * @returns A new instance of a PublicExecutor.
- */
-export function getPublicExecutor(
-  merkleTree: MerkleTreeOperations,
-  publicContractsDB: PublicContractsDB,
-  l1toL2MessageSource: L1ToL2MessageSource,
-  blockData: HistoricBlockData,
-) {
-  return new PublicExecutor(
-    new WorldStatePublicDB(merkleTree),
-    publicContractsDB,
-    new WorldStateDB(merkleTree, l1toL2MessageSource),
-    blockData,
-  );
-}
-
 /**
  * Implements the PublicContractsDB using a ContractDataSource.
  * Progressively records contracts in transaction as they are processed in a block.
@@ -95,8 +69,9 @@ export class ContractsDataSourcePublicDB implements PublicContractsDB {
 /**
  * Implements the PublicStateDB using a world-state database.
  */
-class WorldStatePublicDB implements PublicStateDB {
-  private writeCache: Map<bigint, Fr> = new Map();
+export class WorldStatePublicDB implements PublicStateDB {
+  private commitedWriteCache: Map<bigint, Fr> = new Map();
+  private uncommitedWriteCache: Map<bigint, Fr> = new Map();
 
   constructor(private db: MerkleTreeOperations) {}
 
@@ -108,9 +83,13 @@ class WorldStatePublicDB implements PublicStateDB {
    */
   public async storageRead(contract: AztecAddress, slot: Fr): Promise<Fr> {
     const index = computePublicDataTreeIndex(contract, slot).value;
-    const cached = this.writeCache.get(index);
-    if (cached !== undefined) {
-      return cached;
+    const uncommited = this.uncommitedWriteCache.get(index);
+    if (uncommited !== undefined) {
+      return uncommited;
+    }
+    const commited = this.commitedWriteCache.get(index);
+    if (commited !== undefined) {
+      return commited;
     }
     const value = await this.db.getLeafValue(MerkleTreeId.PUBLIC_DATA_TREE, index);
     return value ? Fr.fromBuffer(value) : Fr.ZERO;
@@ -124,7 +103,27 @@ class WorldStatePublicDB implements PublicStateDB {
    */
   public storageWrite(contract: AztecAddress, slot: Fr, newValue: Fr): Promise<void> {
     const index = computePublicDataTreeIndex(contract, slot).value;
-    this.writeCache.set(index, newValue);
+    this.uncommitedWriteCache.set(index, newValue);
+    return Promise.resolve();
+  }
+
+  /**
+   * Commit the pending changes to the DB.
+   * @returns Nothing.
+   */
+  commit(): Promise<void> {
+    for (const [k, v] of this.uncommitedWriteCache) {
+      this.commitedWriteCache.set(k, v);
+    }
+    return this.rollback();
+  }
+
+  /**
+   * Rollback the pending changes.
+   * @returns Nothing.
+   */
+  rollback(): Promise<void> {
+    this.uncommitedWriteCache = new Map<bigint, Fr>();
     return Promise.resolve();
   }
 }