diff --git a/packages/common/src/trees/InMemoryLinkedMerkleTreeStorage.ts b/packages/common/src/trees/InMemoryLinkedMerkleTreeStorage.ts index 54ec04df..7cb9b344 100644 --- a/packages/common/src/trees/InMemoryLinkedMerkleTreeStorage.ts +++ b/packages/common/src/trees/InMemoryLinkedMerkleTreeStorage.ts @@ -9,7 +9,7 @@ export class InMemoryLinkedMerkleTreeStorage implements LinkedMerkleTreeStore { protected leaves: { [key: string]: LinkedLeaf; - } = {}; + } = { "0": { value: 0n, path: 0, nextPath: 0 } }; public getNode(index: bigint, level: number): bigint | undefined { return this.nodes[level]?.[index.toString()]; @@ -23,19 +23,6 @@ export class InMemoryLinkedMerkleTreeStorage implements LinkedMerkleTreeStore { return this.leaves[index.toString()]; } - // This gets the leaf with the closest path. - public getClosestPath(path: number): LinkedLeaf { - // eslint-disable-next-line @typescript-eslint/consistent-type-assertions - let largestLeaf = this.getLeaf(0n) as LinkedLeaf; - while (largestLeaf.nextPath < path) { - // eslint-disable-next-line @typescript-eslint/consistent-type-assertions - const nextIndex = this.getLeafIndex(largestLeaf.nextPath) as bigint; - // eslint-disable-next-line @typescript-eslint/consistent-type-assertions - largestLeaf = this.getLeaf(nextIndex) as LinkedLeaf; - } - return largestLeaf; - } - public setLeaf(index: bigint, value: LinkedLeaf): void { this.leaves[index.toString()] = value; } @@ -49,4 +36,21 @@ export class InMemoryLinkedMerkleTreeStorage implements LinkedMerkleTreeStore { } return BigInt(leafIndex); } + + public getMaximumIndex(): bigint { + return BigInt(Object.keys(this.leaves).length) - 1n; + } + + // This gets the leaf with the closest path. + public getClosestPath(path: number): LinkedLeaf { + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + let largestLeaf = this.getLeaf(0n) as LinkedLeaf; + while (largestLeaf.nextPath < path) { + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + const nextIndex = this.getLeafIndex(largestLeaf.nextPath) as bigint; + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + largestLeaf = this.getLeaf(nextIndex) as LinkedLeaf; + } + return largestLeaf; + } } diff --git a/packages/common/src/trees/LinkedMerkleTree.ts b/packages/common/src/trees/LinkedMerkleTree.ts index 857fe85a..bf27b663 100644 --- a/packages/common/src/trees/LinkedMerkleTree.ts +++ b/packages/common/src/trees/LinkedMerkleTree.ts @@ -19,7 +19,6 @@ class LinkedLeaf extends Struct({ export interface AbstractLinkedMerkleTree { store: LinkedMerkleTreeStore; - readonly leafCount: bigint; /** * Returns a node which lives at a given index and level. * @param level Level of the node. @@ -36,20 +35,20 @@ export interface AbstractLinkedMerkleTree { /** * Sets the value of a leaf node at a given index to a given value. - * @param index Position of the leaf node. - * @param leaf New value. + * @param path of the leaf node. + * @param value New value. */ - setLeaf(index: bigint, leaf: LinkedLeaf): void; + setLeaf(path: number, value: bigint): void; /** * Returns a leaf which lives at a given path. * @param path Index of the node. * @returns The data of the leaf. */ - getLeaf(path: number): LinkedLeaf | undefined; + getLeaf(path: number): LinkedLeaf; /** - * Returns a leaf which lives at a given path. + * Returns a leaf which is closest to a given path. * @param path Index of the node. * @returns The data of the leaf. */ @@ -74,8 +73,6 @@ export interface AbstractLinkedMerkleTreeClass { HEIGHT: number; EMPTY_ROOT: bigint; - - get leafCount(): bigint; } export function createLinkedMerkleTree( @@ -160,6 +157,13 @@ export function createLinkedMerkleTree( ].toString() ); } + + public static dummy() { + return new LinkedMerkleWitness({ + isLeft: Array(height - 1).fill(Bool(false)), + path: Array(height - 1).fill(Field(0)), + }); + } } return class AbstractLinkedRollupMerkleTree @@ -173,10 +177,6 @@ export function createLinkedMerkleTree( .getRoot() .toBigInt(); - public static get leafCount(): bigint { - return 2n ** BigInt(AbstractLinkedRollupMerkleTree.HEIGHT - 1); - } - public static WITNESS = LinkedMerkleWitness; // private in interface @@ -199,14 +199,7 @@ export function createLinkedMerkleTree( } } - public assertIndexRange(index: bigint) { - if (index > this.leafCount) { - throw new Error("Index greater than maximum leaf number"); - } - } - public getNode(level: number, index: bigint): Field { - this.assertIndexRange(index); return Field(this.store.getNode(index, level) ?? this.zeroes[level]); } @@ -215,14 +208,14 @@ export function createLinkedMerkleTree( * @param path path of the node. * @returns The data of the node. */ - public getLeaf(path: number): LinkedLeaf | undefined { + public getLeaf(path: number): LinkedLeaf { const index = this.store.getLeafIndex(path); if (index === undefined) { - return index; + throw new Error("Path does not exist in tree."); } const leaf = this.store.getLeaf(BigInt(index)); if (leaf === undefined) { - return undefined; + throw new Error("Index does not exist in tree."); } return { value: Field(leaf.value), @@ -258,24 +251,68 @@ export function createLinkedMerkleTree( } /** - * Sets the value of a leaf node at a given index to a given value. - * @param index Position of the leaf node. - * @param leaf New value. + * Sets the value of a leaf node at a given path to a given value. + * @param path Position of the leaf node. + * @param value New value. */ - public setLeaf(path: bigint, leaf: LinkedLeaf) { - this.setNode(0, index, leaf); - let currentIndex = index; + public setLeaf(path: number, value: bigint) { + const prevLeaf = this.store.getClosestPath(path); + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + let prevLeafIndex = this.store.getLeafIndex(path) as bigint; + const newPrevLeaf = { + value: prevLeaf.value, + path: prevLeaf.path, + nextPath: path, + }; + this.store.setLeaf(prevLeafIndex, newPrevLeaf); + const prevLeafFields = this.getLeaf(prevLeaf.path); + this.setNode( + 0, + prevLeafIndex, + Poseidon.hash([ + prevLeafFields.value, + prevLeafFields.path, + prevLeafFields.nextPath, + ]) + ); + + const newLeaf = { + value: value, + path: path, + nextPath: prevLeaf.nextPath, + }; + let newLeafIndex = this.store.getMaximumIndex() + 1n; + this.store.setLeaf(newLeafIndex, newLeaf); + const newLeafFields = this.getLeaf(path); + this.setNode( + 0, + newLeafIndex, + Poseidon.hash([ + newLeafFields.value, + newLeafFields.path, + newLeafFields.nextPath, + ]) + ); + for ( let level = 1; level < AbstractLinkedRollupMerkleTree.HEIGHT; level += 1 ) { - currentIndex /= 2n; + prevLeafIndex /= 2n; + newLeafIndex /= 2n; - const left = this.getNode(level - 1, currentIndex * 2n); - const right = this.getNode(level - 1, currentIndex * 2n + 1n); + const leftPrev = this.getNode(level - 1, prevLeafIndex * 2n); + const rightPrev = this.getNode(level - 1, prevLeafIndex * 2n + 1n); + const leftNew = this.getNode(level - 1, newLeafIndex * 2n); + const rightNew = this.getNode(level - 1, newLeafIndex * 2n + 1n); - this.setNode(level, currentIndex, Poseidon.hash([left, right])); + this.setNode( + level, + prevLeafIndex, + Poseidon.hash([leftPrev, rightPrev]) + ); + this.setNode(level, prevLeafIndex, Poseidon.hash([leftNew, rightNew])); } } @@ -283,7 +320,7 @@ export function createLinkedMerkleTree( * Returns the witness (also known as * [Merkle Proof or Merkle Witness](https://computersciencewiki.org/index.php/Merkle_proof)) * for the leaf at the given index. - * @param index Position of the leaf node. + * @param path of the leaf node. * @returns The witness that belongs to the leaf. */ public getWitness(path: number): LinkedMerkleWitness { @@ -314,14 +351,6 @@ export function createLinkedMerkleTree( path: pathArray, }); } - - /** - * Returns the amount of leaf nodes. - * @returns Amount of leaf nodes. - */ - public get leafCount(): bigint { - return AbstractLinkedRollupMerkleTree.leafCount; - } }; } diff --git a/packages/common/src/trees/LinkedMerkleTreeStore.ts b/packages/common/src/trees/LinkedMerkleTreeStore.ts index 00ae5353..fec94874 100644 --- a/packages/common/src/trees/LinkedMerkleTreeStore.ts +++ b/packages/common/src/trees/LinkedMerkleTreeStore.ts @@ -10,6 +10,8 @@ export interface LinkedMerkleTreeStore { getLeafIndex: (path: number) => bigint | undefined; getClosestPath: (path: number) => LinkedLeaf; + + getMaximumIndex: () => bigint; } export type LinkedLeaf = { value: bigint; path: number; nextPath: number };