From 1e7d42018e667e84bb702b75d7183a71fc576fb5 Mon Sep 17 00:00:00 2001
From: ejMina226 <118474890+ejMina226@users.noreply.github.com>
Date: Wed, 4 Dec 2024 13:56:21 +0000
Subject: [PATCH] Settlement changes

---
 packages/common/src/trees/LinkedMerkleTree.ts      |  4 +++-
 .../contracts/SettlementSmartContract.ts           |  2 +-
 .../settlement/messages/OutgoingMessageArgument.ts | 12 +++++++++---
 .../sequencer/src/settlement/SettlementModule.ts   | 14 ++++++++------
 packages/sequencer/test/settlement/Settlement.ts   |  4 ++--
 5 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/packages/common/src/trees/LinkedMerkleTree.ts b/packages/common/src/trees/LinkedMerkleTree.ts
index 34fe5e841..66d67ea9f 100644
--- a/packages/common/src/trees/LinkedMerkleTree.ts
+++ b/packages/common/src/trees/LinkedMerkleTree.ts
@@ -78,6 +78,8 @@ export interface AbstractLinkedMerkleTree {
   getWitness(path: bigint): LinkedLeafAndMerkleWitness;
 
   dummyWitness(): LinkedMerkleTreeWitness;
+
+  dummy(): LinkedLeafAndMerkleWitness;
 }
 
 export interface AbstractLinkedMerkleTreeClass {
@@ -447,7 +449,7 @@ export function createLinkedMerkleTree(
       });
     }
 
-    private dummy(): LinkedLeafAndMerkleWitness {
+    public dummy(): LinkedLeafAndMerkleWitness {
       return new LinkedLeafAndMerkleWitness({
         merkleWitness: new RollupMerkleTreeWitness({
           // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
diff --git a/packages/protocol/src/settlement/contracts/SettlementSmartContract.ts b/packages/protocol/src/settlement/contracts/SettlementSmartContract.ts
index a0966ce37..54e9e3bac 100644
--- a/packages/protocol/src/settlement/contracts/SettlementSmartContract.ts
+++ b/packages/protocol/src/settlement/contracts/SettlementSmartContract.ts
@@ -320,7 +320,7 @@ export class SettlementSmartContract
       // Check witness
       const path = Path.fromKey(mapPath, Field, counter);
 
-      args.witness
+      args.witness.merkleWitness
         .checkMembership(
           stateRoot,
           path,
diff --git a/packages/protocol/src/settlement/messages/OutgoingMessageArgument.ts b/packages/protocol/src/settlement/messages/OutgoingMessageArgument.ts
index f7cf9ab2b..112383e59 100644
--- a/packages/protocol/src/settlement/messages/OutgoingMessageArgument.ts
+++ b/packages/protocol/src/settlement/messages/OutgoingMessageArgument.ts
@@ -1,17 +1,23 @@
 import { Bool, Provable, Struct } from "o1js";
-import { RollupMerkleTreeWitness } from "@proto-kit/common";
+import {
+  InMemoryLinkedMerkleLeafStore,
+  LinkedLeafAndMerkleWitness,
+  LinkedMerkleTree,
+} from "@proto-kit/common";
 
 import { Withdrawal } from "./Withdrawal";
 
 export const OUTGOING_MESSAGE_BATCH_SIZE = 1;
 
 export class OutgoingMessageArgument extends Struct({
-  witness: RollupMerkleTreeWitness,
+  witness: LinkedLeafAndMerkleWitness,
   value: Withdrawal,
 }) {
   public static dummy(): OutgoingMessageArgument {
     return new OutgoingMessageArgument({
-      witness: RollupMerkleTreeWitness.dummy(),
+      witness: new LinkedMerkleTree(
+        new InMemoryLinkedMerkleLeafStore()
+      ).dummy(),
       value: Withdrawal.dummy(),
     });
   }
diff --git a/packages/sequencer/src/settlement/SettlementModule.ts b/packages/sequencer/src/settlement/SettlementModule.ts
index 85887757b..3063da3ef 100644
--- a/packages/sequencer/src/settlement/SettlementModule.ts
+++ b/packages/sequencer/src/settlement/SettlementModule.ts
@@ -29,8 +29,8 @@ import {
   EventEmittingComponent,
   log,
   noop,
-  RollupMerkleTree,
   AreProofsEnabled,
+  LinkedMerkleTree,
 } from "@proto-kit/common";
 import { Runtime, RuntimeModulesRecord } from "@proto-kit/module";
 
@@ -43,11 +43,11 @@ import { SettlementStorage } from "../storage/repositories/SettlementStorage";
 import { MessageStorage } from "../storage/repositories/MessageStorage";
 import type { MinaBaseLayer } from "../protocol/baselayer/MinaBaseLayer";
 import { Batch, SettleableBatch } from "../storage/model/Batch";
-import { AsyncMerkleTreeStore } from "../state/async/AsyncMerkleTreeStore";
-import { CachedMerkleTreeStore } from "../state/merkle/CachedMerkleTreeStore";
 import { BlockProofSerializer } from "../protocol/production/helpers/BlockProofSerializer";
 import { Settlement } from "../storage/model/Settlement";
 import { FeeStrategy } from "../protocol/baselayer/fees/FeeStrategy";
+import { AsyncLinkedMerkleTreeStore } from "../state/async/AsyncLinkedMerkleTreeStore";
+import { CachedLinkedMerkleTreeStore } from "../state/merkle/CachedLinkedMerkleTreeStore";
 
 import { IncomingMessageAdapter } from "./messages/IncomingMessageAdapter";
 import type { OutgoingMessageQueue } from "./messages/WithdrawalQueue";
@@ -105,7 +105,7 @@ export class SettlementModule
     @inject("OutgoingMessageQueue")
     private readonly outgoingMessageQueue: OutgoingMessageQueue,
     @inject("AsyncMerkleStore")
-    private readonly merkleTreeStore: AsyncMerkleTreeStore,
+    private readonly merkleTreeStore: AsyncLinkedMerkleTreeStore,
     private readonly blockProofSerializer: BlockProofSerializer,
     @inject("TransactionSender")
     private readonly transactionSender: MinaTransactionSender,
@@ -206,8 +206,10 @@ export class SettlementModule
 
     const { settlement } = this.getContracts();
 
-    const cachedStore = new CachedMerkleTreeStore(this.merkleTreeStore);
-    const tree = new RollupMerkleTree(cachedStore);
+    const cachedStore = await CachedLinkedMerkleTreeStore.new(
+      this.merkleTreeStore
+    );
+    const tree = new LinkedMerkleTree(cachedStore);
 
     const [withdrawalModule, withdrawalStateName] =
       this.getSettlementModuleConfig().withdrawalStatePath.split(".");
diff --git a/packages/sequencer/test/settlement/Settlement.ts b/packages/sequencer/test/settlement/Settlement.ts
index 7735552fb..b5f50b987 100644
--- a/packages/sequencer/test/settlement/Settlement.ts
+++ b/packages/sequencer/test/settlement/Settlement.ts
@@ -1,5 +1,5 @@
 /* eslint-disable no-inner-declarations */
-import { log, mapSequential, RollupMerkleTree } from "@proto-kit/common";
+import { LinkedMerkleTree, log, mapSequential } from "@proto-kit/common";
 import { VanillaProtocolModules } from "@proto-kit/library";
 import { Runtime } from "@proto-kit/module";
 import {
@@ -274,7 +274,7 @@ export const settlementTestFn = (
           batch!.proof.publicInput.map((x) => Field(x))
         );
         expect(input.stateRoot.toBigInt()).toStrictEqual(
-          RollupMerkleTree.EMPTY_ROOT
+          LinkedMerkleTree.EMPTY_ROOT
         );
 
         const lastBlock = await blockQueue.getLatestBlock();