From 35cd1fa751126f589760da52f76de4b407be9f04 Mon Sep 17 00:00:00 2001 From: ejMina226 <118474890+ejMina226@users.noreply.github.com> Date: Wed, 20 Nov 2024 13:09:14 +0000 Subject: [PATCH] Got syncedCachedLinkedMerkleTreeStoreWorking --- packages/common/src/trees/LinkedMerkleTree.ts | 2 +- .../merkle/SyncCachedLinkedMerkleTreeStore.ts | 16 ++++- .../merkle/CachedLinkedMerkleStore.test.ts | 58 +++++++++++++++---- 3 files changed, 62 insertions(+), 14 deletions(-) diff --git a/packages/common/src/trees/LinkedMerkleTree.ts b/packages/common/src/trees/LinkedMerkleTree.ts index 4fa11643..a51d456a 100644 --- a/packages/common/src/trees/LinkedMerkleTree.ts +++ b/packages/common/src/trees/LinkedMerkleTree.ts @@ -220,7 +220,7 @@ export function createLinkedMerkleTree( Poseidon.hash([previousLevel, previousLevel]).toBigInt() ); } - // We only do the leaf initialisation the store + // We only do the leaf initialisation when the store // has no values. Otherwise, we leave the store // as is to not overwrite any data. if (this.store.getMaximumIndex() <= 0n) { diff --git a/packages/sequencer/src/state/merkle/SyncCachedLinkedMerkleTreeStore.ts b/packages/sequencer/src/state/merkle/SyncCachedLinkedMerkleTreeStore.ts index fb01e542..4102ff01 100644 --- a/packages/sequencer/src/state/merkle/SyncCachedLinkedMerkleTreeStore.ts +++ b/packages/sequencer/src/state/merkle/SyncCachedLinkedMerkleTreeStore.ts @@ -5,6 +5,8 @@ import { LinkedMerkleTreeStore, } from "@proto-kit/common"; +// This is mainly used for supporting the rollbacks we need to do in case a runtimemethod fails +// In this case everything should be preloaded in the parent async service export class SyncCachedLinkedMerkleTreeStore extends InMemoryLinkedMerkleTreeStorage { public constructor(private readonly parent: LinkedMerkleTreeStore) { super(); @@ -26,6 +28,18 @@ export class SyncCachedLinkedMerkleTreeStore extends InMemoryLinkedMerkleTreeSto super.setLeaf(index, value); } + // Need to make sure we call the parent as the super will usually be empty + // The Tree calls this method. + public getLeafIndex(path: bigint): bigint | undefined { + return this.parent.getLeafIndex(path); + } + + // Need to make sure we call the parent as the super will usually be empty + // The tree calls this method. + public getMaximumIndex(): bigint { + return this.parent.getMaximumIndex(); + } + public mergeIntoParent() { if (Object.keys(this.leaves).length === 0) { return; @@ -33,7 +47,7 @@ export class SyncCachedLinkedMerkleTreeStore extends InMemoryLinkedMerkleTreeSto const { nodes, leaves } = this; Object.entries(leaves).forEach(([key, leaf]) => - this.setLeaf(BigInt(key), leaf) + this.parent.setLeaf(BigInt(key), leaf) ); Array.from({ length: LinkedMerkleTree.HEIGHT }).forEach((ignored, level) => Object.entries(nodes[level]).forEach((entry) => { diff --git a/packages/sequencer/test/merkle/CachedLinkedMerkleStore.test.ts b/packages/sequencer/test/merkle/CachedLinkedMerkleStore.test.ts index 07c17b3e..10c1e7ff 100644 --- a/packages/sequencer/test/merkle/CachedLinkedMerkleStore.test.ts +++ b/packages/sequencer/test/merkle/CachedLinkedMerkleStore.test.ts @@ -16,7 +16,7 @@ describe("cached linked merkle store", () => { const cachedStore = await CachedLinkedMerkleTreeStore.new(mainStore); const tmpTree = new LinkedMerkleTree(cachedStore); - await tmpTree.setLeaf(5n, 10n); + tmpTree.setLeaf(5n, 10n); await cachedStore.mergeIntoParent(); cache1 = await CachedLinkedMerkleTreeStore.new(mainStore); @@ -183,51 +183,85 @@ describe("cached linked merkle store", () => { }); it("should cache correctly", async () => { - expect.assertions(9); + expect.assertions(12); const cache2 = new SyncCachedLinkedMerkleTreeStore(cache1); const tree2 = new LinkedMerkleTree(cache2); const leaf1 = tree2.getLeaf(5n); + const leaf1Index = cache2.getLeafIndex(5n); + expectDefined(leaf1Index); await expect( - mainStore.getNodesAsync([{ key: 5n, level: 0 }]) + mainStore.getNodesAsync([{ key: leaf1Index, level: 0 }]) ).resolves.toStrictEqual([ Poseidon.hash([leaf1.value, leaf1.path, leaf1.nextPath]).toBigInt(), ]); - await cache1.preloadKey(5n); - - await tree1.setLeaf(10n, 20n); + tree1.setLeaf(10n, 20n); const leaf2 = tree2.getLeaf(10n); - expect(tree2.getNode(0, 10n).toBigInt()).toBe( + const leaf2Index = cache2.getLeafIndex(10n); + expectDefined(leaf2Index); + expect(tree2.getNode(0, leaf2Index).toBigInt()).toBe( Poseidon.hash([leaf2.value, leaf2.path, leaf2.nextPath]).toBigInt() ); const witness = tree2.getWitness(5n); + // We check tree1 and tree2 have same hash roots. + // The witness is from tree2, which comes from cache2, + // but which because of the sync is really just cache1. expect( - witness.leafCurrent.merkleWitness.calculateRoot(Field(10)).toString() + witness.leafCurrent.merkleWitness + .calculateRoot( + Poseidon.hash([ + witness.leafCurrent.leaf.value, + witness.leafCurrent.leaf.path, + witness.leafCurrent.leaf.nextPath, + ]) + ) + .toString() ).toBe(tree1.getRoot().toString()); + expect( - witness.leafCurrent.merkleWitness.calculateRoot(Field(11)).toString() + witness.leafCurrent.merkleWitness + .calculateRoot(Poseidon.hash([Field(11), Field(5n), Field(10n)])) + .toString() ).not.toBe(tree1.getRoot().toString()); const witness2 = tree1.getWitness(10n); expect( - witness2.leafCurrent.merkleWitness.calculateRoot(Field(20)).toString() + witness2.leafCurrent.merkleWitness + .calculateRoot( + Poseidon.hash([ + Field(20), + Field(10n), + witness2.leafCurrent.leaf.nextPath, // This is the maximum as the the leaf 10n should be the last + ]) + ) + .toString() ).toBe(tree2.getRoot().toString()); - await tree2.setLeaf(15n, 30n); + tree2.setLeaf(15n, 30n); + // Won't be the same as the tree2 works on cache2 and these changes don't + // carry up to cache1. Have to merge into parent for this. expect(tree1.getRoot().toString()).not.toBe(tree2.getRoot().toString()); + // After this the changes should be merged into the parents, i.e. cache1, + // which tree1 has access to. cache2.mergeIntoParent(); + const index15 = cache2.getLeafIndex(15n); + const leaf15 = tree2.getLeaf(15n); + expectDefined(index15); expect(tree1.getRoot().toString()).toBe(tree2.getRoot().toString()); - expect(tree1.getNode(0, 15n).toString()).toBe("30"); + expect(tree1.getNode(0, index15).toString()).toBe( + Poseidon.hash([leaf15.value, leaf15.path, leaf15.nextPath]).toString() + ); + // Now the mainstore has the new 15n root. await cache1.mergeIntoParent(); const cachedStore = await CachedLinkedMerkleTreeStore.new(mainStore);