From 6e6c92468a04aaf7d2f9437a0d05bcd58fd9dc03 Mon Sep 17 00:00:00 2001
From: Alex Gherghisan <alexg@aztecprotocol.com>
Date: Thu, 28 Sep 2023 10:00:59 +0100
Subject: [PATCH] refactor: temp add new contract in current tx only

---
 .../aztec-node/src/aztec-node/server.ts       |  2 +-
 .../src/sequencer/public_processor.test.ts    |  5 ++
 .../src/sequencer/public_processor.ts         | 26 ++++------
 .../src/sequencer/sequencer.test.ts           |  2 +-
 .../src/sequencer/sequencer.ts                |  6 +--
 .../src/simulator/public_executor.ts          | 51 +++++++++++++------
 6 files changed, 53 insertions(+), 39 deletions(-)

diff --git a/yarn-project/aztec-node/src/aztec-node/server.ts b/yarn-project/aztec-node/src/aztec-node/server.ts
index 87beedf8afce..4092bacaf6bc 100644
--- a/yarn-project/aztec-node/src/aztec-node/server.ts
+++ b/yarn-project/aztec-node/src/aztec-node/server.ts
@@ -413,7 +413,7 @@ export class AztecNodeService implements AztecNode {
       this.contractDataSource,
       this.l1ToL2MessageSource,
     );
-    const processor = await publicProcessorFactory.create(prevGlobalVariables, newGlobalVariables, tx.newContracts);
+    const processor = await publicProcessorFactory.create(prevGlobalVariables, newGlobalVariables);
     const [, failedTxs] = await processor.process([tx]);
     if (failedTxs.length) {
       throw failedTxs[0].error;
diff --git a/yarn-project/sequencer-client/src/sequencer/public_processor.test.ts b/yarn-project/sequencer-client/src/sequencer/public_processor.test.ts
index b403b4610840..ec9fcdf3a9ef 100644
--- a/yarn-project/sequencer-client/src/sequencer/public_processor.test.ts
+++ b/yarn-project/sequencer-client/src/sequencer/public_processor.test.ts
@@ -35,6 +35,7 @@ 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 { WasmPublicKernelCircuitSimulator } from '../simulator/public_kernel.js';
 import { PublicProcessor } from './public_processor.js';
 
@@ -42,6 +43,7 @@ describe('public_processor', () => {
   let db: MockProxy<MerkleTreeOperations>;
   let publicExecutor: MockProxy<PublicExecutor>;
   let publicProver: MockProxy<PublicProver>;
+  let publicContractsDB: MockProxy<ContractsDataSourcePublicDB>;
 
   let proof: Proof;
   let root: Buffer;
@@ -52,6 +54,7 @@ describe('public_processor', () => {
     db = mock<MerkleTreeOperations>();
     publicExecutor = mock<PublicExecutor>();
     publicProver = mock<PublicProver>();
+    publicContractsDB = mock<ContractsDataSourcePublicDB>();
 
     proof = makeEmptyProof();
     root = Buffer.alloc(32, 5);
@@ -73,6 +76,7 @@ describe('public_processor', () => {
         publicProver,
         GlobalVariables.empty(),
         HistoricBlockData.empty(),
+        publicContractsDB,
       );
     });
 
@@ -128,6 +132,7 @@ describe('public_processor', () => {
         publicProver,
         GlobalVariables.empty(),
         HistoricBlockData.empty(),
+        publicContractsDB,
       );
     });
 
diff --git a/yarn-project/sequencer-client/src/sequencer/public_processor.ts b/yarn-project/sequencer-client/src/sequencer/public_processor.ts
index 6f33279dc8bb..e7bea78e09ee 100644
--- a/yarn-project/sequencer-client/src/sequencer/public_processor.ts
+++ b/yarn-project/sequencer-client/src/sequencer/public_processor.ts
@@ -41,21 +41,14 @@ import { computeCallStackItemHash, computeVarArgsHash } from '@aztec/circuits.js
 import { arrayNonEmptyLength, isArrayEmpty, padArrayEnd, padArrayStart } from '@aztec/foundation/collection';
 import { createDebugLogger } from '@aztec/foundation/log';
 import { Tuple, mapTuple, to2Fields } from '@aztec/foundation/serialize';
-import {
-  ContractDataSource,
-  ExtendedContractData,
-  FunctionL2Logs,
-  L1ToL2MessageSource,
-  MerkleTreeId,
-  Tx,
-} from '@aztec/types';
+import { ContractDataSource, FunctionL2Logs, L1ToL2MessageSource, MerkleTreeId, Tx } from '@aztec/types';
 import { MerkleTreeOperations } from '@aztec/world-state';
 
 import { getVerificationKeys } from '../index.js';
 import { EmptyPublicProver } from '../prover/empty.js';
 import { PublicProver } from '../prover/index.js';
 import { PublicKernelCircuitSimulator } from '../simulator/index.js';
-import { getPublicExecutor } from '../simulator/public_executor.js';
+import { ContractsDataSourcePublicDB, getPublicExecutor } 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';
@@ -80,22 +73,17 @@ export class PublicProcessorFactory {
   public async create(
     prevGlobalVariables: GlobalVariables,
     globalVariables: GlobalVariables,
-    newContracts: ExtendedContractData[] = [],
   ): Promise<PublicProcessor> {
     const blockData = await getHistoricBlockData(this.merkleTree, prevGlobalVariables);
+    const publicContractsDB = new ContractsDataSourcePublicDB(this.contractDataSource);
     return new PublicProcessor(
       this.merkleTree,
-      getPublicExecutor(
-        this.merkleTree,
-        this.contractDataSource,
-        this.l1Tol2MessagesDataSource,
-        blockData,
-        newContracts,
-      ),
+      getPublicExecutor(this.merkleTree, publicContractsDB, this.l1Tol2MessagesDataSource, blockData),
       new WasmPublicKernelCircuitSimulator(),
       new EmptyPublicProver(),
       globalVariables,
       blockData,
+      publicContractsDB,
     );
   }
 }
@@ -112,6 +100,7 @@ export class PublicProcessor {
     protected publicProver: PublicProver,
     protected globalVariables: GlobalVariables,
     protected blockData: HistoricBlockData,
+    protected publicContractsDB: ContractsDataSourcePublicDB,
 
     private log = createDebugLogger('aztec:sequencer:public-processor'),
   ) {}
@@ -130,6 +119,7 @@ export class PublicProcessor {
     for (const tx of txs) {
       this.log(`Processing tx ${await tx.getTxHash()}`);
       try {
+        await this.publicContractsDB.addNewContracts(tx);
         result.push(await this.processTx(tx));
       } catch (err) {
         this.log.warn(`Error processing tx ${await tx.getTxHash()}: ${err}`);
@@ -137,6 +127,8 @@ export class PublicProcessor {
           tx,
           error: err instanceof Error ? err : new Error('Unknown error'),
         });
+      } finally {
+        await this.publicContractsDB.clearTxContracts();
       }
     }
 
diff --git a/yarn-project/sequencer-client/src/sequencer/sequencer.test.ts b/yarn-project/sequencer-client/src/sequencer/sequencer.test.ts
index 02c7cfd0380e..6ccd8bf46ea9 100644
--- a/yarn-project/sequencer-client/src/sequencer/sequencer.test.ts
+++ b/yarn-project/sequencer-client/src/sequencer/sequencer.test.ts
@@ -70,7 +70,7 @@ describe('sequencer', () => {
     });
 
     publicProcessorFactory = mock<PublicProcessorFactory>({
-      create: (_a, _b_, _c) => Promise.resolve(publicProcessor),
+      create: (_a, _b_) => Promise.resolve(publicProcessor),
     });
 
     l2BlockSource = mock<L2BlockSource>({
diff --git a/yarn-project/sequencer-client/src/sequencer/sequencer.ts b/yarn-project/sequencer-client/src/sequencer/sequencer.ts
index 81d74462eb53..e6ae74f3b276 100644
--- a/yarn-project/sequencer-client/src/sequencer/sequencer.ts
+++ b/yarn-project/sequencer-client/src/sequencer/sequencer.ts
@@ -137,11 +137,7 @@ export class Sequencer {
 
       // Process txs and drop the ones that fail processing
       // We create a fresh processor each time to reset any cached state (eg storage writes)
-      const processor = await this.publicProcessorFactory.create(
-        prevGlobalVariables,
-        newGlobalVariables,
-        validTxs.flatMap(tx => tx.newContracts),
-      );
+      const processor = await this.publicProcessorFactory.create(prevGlobalVariables, newGlobalVariables);
       const [processedTxs, failedTxs] = await processor.process(validTxs);
       if (failedTxs.length > 0) {
         const failedTxData = failedTxs.map(fail => fail.tx);
diff --git a/yarn-project/sequencer-client/src/simulator/public_executor.ts b/yarn-project/sequencer-client/src/simulator/public_executor.ts
index c75cd3628598..cc981ef700b0 100644
--- a/yarn-project/sequencer-client/src/simulator/public_executor.ts
+++ b/yarn-project/sequencer-client/src/simulator/public_executor.ts
@@ -6,7 +6,7 @@ import {
   PublicStateDB,
 } from '@aztec/acir-simulator';
 import { AztecAddress, CircuitsWasm, EthAddress, Fr, FunctionSelector, HistoricBlockData } from '@aztec/circuits.js';
-import { ContractDataSource, ExtendedContractData, L1ToL2MessageSource, MerkleTreeId } from '@aztec/types';
+import { ContractDataSource, ExtendedContractData, L1ToL2MessageSource, MerkleTreeId, Tx } from '@aztec/types';
 import { MerkleTreeOperations, computePublicDataTreeLeafIndex } from '@aztec/world-state';
 
 /**
@@ -17,14 +17,13 @@ import { MerkleTreeOperations, computePublicDataTreeLeafIndex } from '@aztec/wor
  */
 export function getPublicExecutor(
   merkleTree: MerkleTreeOperations,
-  contractDataSource: ContractDataSource,
+  publicContractsDB: PublicContractsDB,
   l1toL2MessageSource: L1ToL2MessageSource,
   blockData: HistoricBlockData,
-  newContracts: ExtendedContractData[] = [],
 ) {
   return new PublicExecutor(
     new WorldStatePublicDB(merkleTree),
-    new ContractsDataSourcePublicDB(contractDataSource, newContracts),
+    publicContractsDB,
     new WorldStateDB(merkleTree, l1toL2MessageSource),
     blockData,
   );
@@ -33,24 +32,46 @@ export function getPublicExecutor(
 /**
  * Implements the PublicContractsDB using a ContractDataSource and a set of new contracts.
  */
-class ContractsDataSourcePublicDB implements PublicContractsDB {
-  constructor(private db: ContractDataSource, private newContracts: ExtendedContractData[] = []) {}
+export class ContractsDataSourcePublicDB implements PublicContractsDB {
+  cache = new Map<string, ExtendedContractData>();
+
+  constructor(private db: ContractDataSource) {}
+
+  /**
+   * Add new contracts from a transaction
+   * @param tx - The transaction to add contracts from.
+   */
+  public addNewContracts(tx: Tx): Promise<void> {
+    for (const contract of tx.newContracts) {
+      this.cache.set(contract.contractData.contractAddress.toString(), contract);
+    }
+
+    return Promise.resolve();
+  }
+
+  /**
+   * Removes new contracts added from transactions
+   */
+  public clearTxContracts(): Promise<void> {
+    this.cache.clear();
+    return Promise.resolve();
+  }
+
   async getBytecode(address: AztecAddress, selector: FunctionSelector): Promise<Buffer | undefined> {
-    const contractData = await this.#getContractData(address);
-    return contractData?.getPublicFunction(selector)?.bytecode;
+    const contract = await this.#getContract(address);
+    return contract?.getPublicFunction(selector)?.bytecode;
   }
   async getIsInternal(address: AztecAddress, selector: FunctionSelector): Promise<boolean | undefined> {
-    const contractData = await this.#getContractData(address);
-    return contractData?.getPublicFunction(selector)?.isInternal;
+    const contract = await this.#getContract(address);
+    return contract?.getPublicFunction(selector)?.isInternal;
   }
   async getPortalContractAddress(address: AztecAddress): Promise<EthAddress | undefined> {
-    const contractData = await this.#getContractData(address);
-    return contractData?.contractData.portalContractAddress;
+    const contract = await this.#getContract(address);
+    return contract?.contractData.portalContractAddress;
   }
 
-  #getContractData(contractAddress: AztecAddress): Promise<ExtendedContractData | undefined> {
-    const contract = this.newContracts.find(c => c.contractData.contractAddress.equals(contractAddress));
-    return Promise.resolve(contract ?? this.db.getExtendedContractData(contractAddress));
+  async #getContract(address: AztecAddress): Promise<ExtendedContractData | undefined> {
+    return this.cache.get(address.toString()) ?? (await this.db.getExtendedContractData(address));
   }
 }