Move changes to getNearestPath to another method on the cache and cal…

…l separately.

packages/common/src/trees/InMemoryLinkedMerkleTreeStorage.ts

@@ -49,7 +49,7 @@ export class InMemoryLinkedMerkleTreeStorage implements LinkedMerkleTreeStore {
   }
 
   // This gets the leaf with the closest path.
-  public getLeafLessOrEqual(path: bigint): Promise<LinkedLeaf> {
+  public getLeafLessOrEqual(path: bigint): LinkedLeaf {
     // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
     let largestLeaf = this.getLeaf(0n) as LinkedLeaf;
     while (largestLeaf.nextPath <= path) {
@@ -58,6 +58,6 @@ export class InMemoryLinkedMerkleTreeStorage implements LinkedMerkleTreeStore {
       // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
       largestLeaf = this.getLeaf(nextIndex) as LinkedLeaf;
     }
-    return Promise.resolve(largestLeaf);
+    return largestLeaf;
   }
 }

packages/common/src/trees/LinkedMerkleTree.ts

@@ -53,7 +53,7 @@ export interface AbstractLinkedMerkleTree {
    * @param path of the leaf node.
    * @param value New value.
    */
-  setLeaf(path: bigint, value: bigint): Promise<LinkedMerkleTreeWitness>;
+  setLeaf(path: bigint, value: bigint): LinkedMerkleTreeWitness;
 
   /**
    * Returns a leaf which lives at a given path.
@@ -300,9 +300,9 @@ export function createLinkedMerkleTree(
      * @param path Position of the leaf node.
      * @param value New value.
      */
-    public async setLeaf(path: bigint, value: bigint) {
+    public setLeaf(path: bigint, value: bigint) {
       let index = this.store.getLeafIndex(path);
-      const prevLeaf = await this.store.getLeafLessOrEqual(path);
+      const prevLeaf = this.store.getLeafLessOrEqual(path);
       let witnessPrevious;
       if (index === undefined) {
         // The above means the path doesn't already exist, and we are inserting, not updating.

packages/common/src/trees/LinkedMerkleTreeStore.ts

@@ -11,7 +11,7 @@ export interface LinkedMerkleTreeStore extends MerkleTreeStore {
 
   getLeafIndex: (path: bigint) => bigint | undefined;
 
-  getLeafLessOrEqual: (path: bigint) => Promise<LinkedLeaf>;
+  getLeafLessOrEqual: (path: bigint) => LinkedLeaf;
 
   getMaximumIndex: () => bigint;
 }

packages/sequencer/src/state/merkle/CachedLinkedMerkleTreeStore.ts

@@ -182,9 +182,12 @@ export class CachedLinkedMerkleTreeStore
     return this.writeCache.leaves;
   }
 
-  // This gets the leaf with the path equal or just less than
-  // the given path.
-  public async getLeafLessOrEqual(path: bigint): Promise<LinkedLeaf> {
+  // This ensures all the keys needed to be loaded
+  // to find the closest path are loaded.
+  // A bit repetitive as we basically repeat the process
+  // (without the loading) when we find the closest leaf.
+  // TODO: see how we could sue a returned value.
+  public async loadUpKeysForClosestPath(path: bigint): Promise<void> {
     // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
     let largestLeaf = this.getLeaf(0n) as LinkedLeaf;
     while (largestLeaf.nextPath <= path) {
@@ -200,7 +203,6 @@ export class CachedLinkedMerkleTreeStore
       }
       largestLeaf = nextLeaf;
     }
-    return largestLeaf;
   }
 
   // This resets the cache (not the in memory tree).

packages/sequencer/src/storage/inmemory/InMemoryAsyncLinkedMerkleTreeStore.ts

@@ -59,11 +59,6 @@ export class InMemoryAsyncLinkedMerkleTreeStore
     return this.store.getMaximumIndex();
   }
 
-  // This gets the leaf with the closest path.
-  public getPathLessOrEqual(path: bigint) {
-    return this.store.getPathLessOrEqual(path);
-  }
-
   public getLeafByIndex(index: bigint) {
     return this.store.getLeaf(index);
   }

packages/sequencer/test/merkle/CachedLinkedMerkleStore.test.ts

@@ -26,8 +26,8 @@ describe("cached linked merkle store", () => {
   it("should cache multiple keys correctly", async () => {
     expect.assertions(10);
 
-    await tree1.setLeaf(16n, 16n);
-    await tree1.setLeaf(46n, 46n);
+    tree1.setLeaf(16n, 16n);
+    tree1.setLeaf(46n, 46n);
 
     const cache2 = await CachedLinkedMerkleTreeStore.new(cache1);
     const tree2 = new LinkedMerkleTree(cache2);
@@ -68,10 +68,10 @@ describe("cached linked merkle store", () => {
   });
 
   it("should preload through multiple levels and insert correctly at right index", async () => {
-    await tree1.setLeaf(10n, 10n);
-    await tree1.setLeaf(11n, 11n);
-    await tree1.setLeaf(12n, 12n);
-    await tree1.setLeaf(13n, 13n);
+    tree1.setLeaf(10n, 10n);
+    tree1.setLeaf(11n, 11n);
+    tree1.setLeaf(12n, 12n);
+    tree1.setLeaf(13n, 13n);
 
     // Nodes 0 and 5 should be auto-preloaded when cache2 is created
     // as 0 is the first and 5 is its sibling. Similarly, 12 and 13
@@ -85,7 +85,8 @@ describe("cached linked merkle store", () => {
     // When we set this leaf the missing nodes are preloaded
     // as when we do a set we have to go through all the leaves to find
     // the one with the nextPath that is suitable
-    await tree2.setLeaf(14n, 14n);
+    await cache2.loadUpKeysForClosestPath(14n);
+    tree2.setLeaf(14n, 14n);
 
     const leaf = tree1.getLeaf(5n);
     const leaf2 = tree2.getLeaf(14n);
@@ -120,12 +121,12 @@ describe("cached linked merkle store", () => {
   });
 
   it("should preload through multiple levels and insert correctly at right index - harder", async () => {
-    await tree1.setLeaf(10n, 10n);
-    await tree1.setLeaf(100n, 100n);
-    await tree1.setLeaf(200n, 200n);
-    await tree1.setLeaf(300n, 300n);
-    await tree1.setLeaf(400n, 400n);
-    await tree1.setLeaf(500n, 500n);
+    tree1.setLeaf(10n, 10n);
+    tree1.setLeaf(100n, 100n);
+    tree1.setLeaf(200n, 200n);
+    tree1.setLeaf(300n, 300n);
+    tree1.setLeaf(400n, 400n);
+    tree1.setLeaf(500n, 500n);
 
     // Nodes 0 and 5 should be auto-preloaded when cache2 is created
     // as 0 is the first and 5 is its sibling. Similarly, 400 and 500
@@ -141,6 +142,7 @@ describe("cached linked merkle store", () => {
     // the one with the nextPath that is suitable and this preloads that are missing before.
     // This means 10n will be preloaded and since 100n is its sibling this will be preloaded, too.
     // Note that the nodes 200n and 300n are not preloaded.
+    await cache2.loadUpKeysForClosestPath(14n);
     await tree2.setLeaf(14n, 14n);
 
     const leaf = tree1.getLeaf(5n);