Skip to content

Commit

Permalink
Get first few tests passing.
Browse files Browse the repository at this point in the history
  • Loading branch information
@ejMina226
ejMina226 committed Nov 18, 2024
1 parent 55dda23 commit d0beafd
Show file tree
Hide file tree
Showing 3 changed files with 46 additions and 74 deletions.
3 changes: 3 additions & 0 deletions packages/sequencer/src/state/async/AsyncLinkedMerkleTreeStore.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,4 +22,7 @@ export interface AsyncLinkedMerkleTreeStore {
getPathLessOrEqual: (path: bigint) => LinkedLeaf;

getMaximumIndex: () => bigint;

// For the preloadedKeys functionality
getLeafByIndex: (index: bigint) => LinkedLeaf | undefined;
}
89 changes: 26 additions & 63 deletions packages/sequencer/src/state/merkle/CachedLinkedMerkleTreeStore.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,18 @@ export class CachedLinkedMerkleTreeStore
super();
}

public static async new(
parent: AsyncLinkedMerkleTreeStore
): Promise<CachedLinkedMerkleTreeStore> {
// do your async stuff here
// now instantiate and return a class
const cachedInstance = new CachedLinkedMerkleTreeStore(parent);
if (parent.getMaximumIndex() > 0) {
await cachedInstance.preloadKey(0n);
}
return cachedInstance;
}

// This gets the nodes from the in memory store (which looks also to be the cache).
public getNode(key: bigint, level: number): bigint | undefined {
return super.getNode(key, level);
Expand Down Expand Up @@ -141,59 +153,9 @@ export class CachedLinkedMerkleTreeStore
super.setLeaf(BigInt(key), leaf);
}

// // This sets the leaves in the cache and in the in-memory tree.
// // It also updates the relevant node at the base level.
// // Note that setNode doesn't carry the change up the tree (i.e. for siblings etc)
// public setLeaf(index: bigint, leaf: LinkedLeaf) {
// super.setLeaf(index, leaf);
// this.writeCache.leaves[index.toString()] = leaf;
//
// this.setNode(
// index,
// 0,
// Poseidon.hash([
// Field(leaf.value),
// Field(leaf.path),
// Field(leaf.nextPath),
// ]).toBigInt()
// );
// }
// This is setLeaf (cache and in-memory) for a list of leaves.
// It checks if it's an insert or update and then updates the relevant
// leaf.
// public writeLeaves(leaves: { path: bigint; value: bigint }[]) {
// leaves.forEach(({ value, path }) => {
// const index = super.getLeafIndex(path);
// // The following checks if we have an insert or update.
// if (index !== undefined) {
// // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
// const linkedLeaf = super.getLeaf(index) as LinkedLeaf;
// this.setLeaf(index, {
// value: value,
// path: path,
// nextPath: linkedLeaf.nextPath,
// });
// } else {
// // This is an insert. Need to change two leaves.
// const nearestLinkedLeaf = super.getPathLessOrEqual(path);
// // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
// const nearestLinkedleafIndex = super.getLeafIndex(
// nearestLinkedLeaf.path
// ) as bigint;
// const lowestUnoccupiedIndex = super.getMaximumIndex() + 1n;
// this.setLeaf(nearestLinkedleafIndex, {
// value: nearestLinkedLeaf.value,
// path: nearestLinkedLeaf.path,
// nextPath: path,
// });
// this.setLeaf(lowestUnoccupiedIndex, {
// value: value,
// path: path,
// nextPath: nearestLinkedLeaf.path,
// });
// }
// });
// }
public getLeafByIndex(index: bigint) {
return super.getLeaf(index);
}

// This gets the nodes from the cache.
// Only used in mergeIntoParent
Expand Down Expand Up @@ -268,21 +230,22 @@ export class CachedLinkedMerkleTreeStore
return this.collectNodesToFetch(pathIndex);
});

for (const path of paths) {
const pathIndex = this.parent.getLeafIndex(path) ?? 0n;
const resultLeaf = // eslint-disable-next-line no-await-in-loop,@typescript-eslint/consistent-type-assertions
(await this.parent.getLeavesAsync([path]))[0] as LinkedLeaf;
super.setLeaf(pathIndex, resultLeaf);
this.writeCache.leaves[pathIndex.toString()] = resultLeaf;
}

const resultsNode = await this.parent.getNodesAsync(nodesToRetrieve);
nodesToRetrieve.forEach(({ key, level }, index) => {
let index = 0;
for (const retrievedNode of nodesToRetrieve) {
const { key, level } = retrievedNode;
const value = resultsNode[index];
if (value !== undefined) {
this.setNode(key, level, value);
if (level === 0) {
// eslint-disable-next-line @typescript-eslint/consistent-type-assertions
const resultLeaf = this.parent.getLeafByIndex(key) as LinkedLeaf;
super.setLeaf(key, resultLeaf);
this.writeCache.leaves[key.toString()] = resultLeaf;
}
}
});
index += 1;
}
}

// This is preloadKeys with just one index/key.
Expand Down
28 changes: 17 additions & 11 deletions packages/sequencer/test/merkle/CachedLinkedMerkleStore.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,25 +13,24 @@ describe("cached linked merkle store", () => {
let tree1: LinkedMerkleTree;

beforeEach(async () => {
const cachedStore = new CachedLinkedMerkleTreeStore(mainStore);
const cachedStore = await CachedLinkedMerkleTreeStore.new(mainStore);

const tmpTree = new LinkedMerkleTree(cachedStore);
tmpTree.setLeaf(5n, 10n);
await cachedStore.mergeIntoParent();

cache1 = new CachedLinkedMerkleTreeStore(mainStore);
cache1 = await CachedLinkedMerkleTreeStore.new(mainStore);
tree1 = new LinkedMerkleTree(cache1);
});

it("should cache multiple keys correctly", async () => {
expect.assertions(7);

const cache2 = new CachedLinkedMerkleTreeStore(cache1);
const tree2 = new LinkedMerkleTree(cache2);

tree1.setLeaf(16n, 16n);
tree1.setLeaf(46n, 46n);

const cache2 = await CachedLinkedMerkleTreeStore.new(cache1);
const tree2 = new LinkedMerkleTree(cache2);
// Need to preload 0n, as well since the nextPath of the leaf would have changed
// when other leaves were added.
await cache2.preloadKeys([0n, 16n, 46n]);
Expand All @@ -45,9 +44,12 @@ describe("cached linked merkle store", () => {
expectDefined(leaf1Index);
expectDefined(leaf2Index);

// Note that 5n hasn't been loaded so indices are off by 1.
expect(leaf1Index).toStrictEqual(1n);
expect(leaf2Index).toStrictEqual(2n);
// The new leaves are at index 2 and 3, as the index 5 is auto-preloaded
// as it is next to 0, and 0 is always preloaded as well as any relevant
// nodes.

expect(leaf1Index).toStrictEqual(2n);
expect(leaf2Index).toStrictEqual(3n);

expect(tree2.getNode(0, leaf1Index).toBigInt()).toBe(
Poseidon.hash([leaf1.value, leaf1.path, leaf1.nextPath]).toBigInt()
Expand All @@ -62,12 +64,16 @@ describe("cached linked merkle store", () => {
});

it("should preload through multiple levels", async () => {
const cache2 = new CachedLinkedMerkleTreeStore(cache1);
const cache2 = await CachedLinkedMerkleTreeStore.new(cache1);

await cache2.preloadKeys([0n, 5n]);

const leaf = tree1.getLeaf(5n);
expect(cache2.getNode(5n, 0)).toStrictEqual(

const leafIndex = cache2.getLeafIndex(5n);
expectDefined(leafIndex);
expect(leafIndex).toStrictEqual(1n);
expect(cache2.getNode(leafIndex, 0)).toStrictEqual(
Poseidon.hash([leaf.value, leaf.path, leaf.nextPath]).toBigInt()
);
});
Expand Down Expand Up @@ -120,7 +126,7 @@ describe("cached linked merkle store", () => {

await cache1.mergeIntoParent();

const cachedStore = new CachedLinkedMerkleTreeStore(mainStore);
const cachedStore = await CachedLinkedMerkleTreeStore.new(mainStore);
await cachedStore.preloadKey(15n);

expect(new LinkedMerkleTree(cachedStore).getRoot().toString()).toBe(
Expand Down

0 comments on commit d0beafd

Please sign in to comment.