diff --git a/docs/sequencer/spec.md b/docs/sequencer/spec.md new file mode 100644 index 000000000..c74e9d2a8 --- /dev/null +++ b/docs/sequencer/spec.md @@ -0,0 +1,76 @@ +## Merkle Tree Stores + +Object we need to store: +(Nodes, Leaves, MaximumIndex) + +Level 1: +Async stores: (InMemory*, Redis*) + +Schema: +Record + +write +getAsync +getMaximumIndexAsync +getLeafLessOrEqualAsync(path) (gives us either our current leaf or previous leaf in case of insert) + +openTransaction() +commit() + +( getLeafByIndex ) + +Level 2: +CachedStore: implements Sync, parent: Async + +Sync: +set +getNode +getLeaf(path) => { leaf: LinkedLeaf, index: bigint } +getMaximumIndex +getLeafLessOrEqual(path) => { leaf: LinkedLeaf, index: bigint } + +Cached: +preloadMerkleWitness(index) +preloadKeys(paths: string[]) +mergeIntoParent() + +Level 3: +SyncCachedStore: implements Sync, parent: Sync +mergeIntoParent() + +preLoading: +input: path +``` +const leaf = parent.getLeaf(path) +if(leaf !== undefined) { + super.cache(leaf); + // Update + preloadMerkleWitness(leaf.index); +} else { + // Insert + const previousLeaf = parent.getLeafLessOrEqual(path); + super.cache(previousLeaf); + preloadMerkleWitness(previousLeaf.index); + const maximumIndex = this.preloadAndGetMaximumINndex(); // super.getMaximumINdex() ?? await parent.getMaximumIndexASync() + preloadMerkleWitness(maximumIndex); +} + +``` + +Sync interface: +Union of LinkedMerkleTreeStore (rename to LinkedLeafStore) + MerkleTreeStore + +Async +Level 1 methods + +InMemoryLeafStore - subset that does leafs + maximumindex +InMemoryMerkleStore - subset that does only merkle nodes + +-> future Redis + +InMemoryAsyncLinkedMerkleTreeStore - implements Async +uses inmemory implementations + + + + diff --git a/packages/api/src/graphql/VanillaGraphqlModules.ts b/packages/api/src/graphql/VanillaGraphqlModules.ts index e7859ec9d..f94d29128 100644 --- a/packages/api/src/graphql/VanillaGraphqlModules.ts +++ b/packages/api/src/graphql/VanillaGraphqlModules.ts @@ -6,7 +6,7 @@ import { QueryGraphqlModule } from "./modules/QueryGraphqlModule"; import { BatchStorageResolver } from "./modules/BatchStorageResolver"; import { NodeStatusResolver } from "./modules/NodeStatusResolver"; import { BlockResolver } from "./modules/BlockResolver"; -import { MerkleWitnessResolver } from "./modules/MerkleWitnessResolver"; +import { LinkedMerkleWitnessResolver as MerkleWitnessResolver } from "./modules/LinkedMerkleWitnessResolver"; export type VanillaGraphqlModulesRecord = { MempoolResolver: typeof MempoolResolver; diff --git a/packages/api/src/graphql/modules/LeafResolver.ts b/packages/api/src/graphql/modules/LeafResolver.ts new file mode 100644 index 000000000..568060233 --- /dev/null +++ b/packages/api/src/graphql/modules/LeafResolver.ts @@ -0,0 +1,28 @@ +import { Field, ObjectType } from "type-graphql"; +import { LinkedLeafStruct } from "@proto-kit/common"; + +@ObjectType() +export class LeafDTO { + public static fromServiceLayerModel(leaf: LinkedLeafStruct) { + return new LeafDTO( + leaf.value.toString(), + leaf.path.toString(), + leaf.nextPath.toString() + ); + } + + @Field() + value: string; + + @Field() + path: string; + + @Field() + nextPath: string; + + private constructor(value: string, path: string, nextPath: string) { + this.value = value; + this.path = path; + this.nextPath = nextPath; + } +} diff --git a/packages/api/src/graphql/modules/LinkedMerkleWitnessResolver.ts b/packages/api/src/graphql/modules/LinkedMerkleWitnessResolver.ts new file mode 100644 index 000000000..1538506f8 --- /dev/null +++ b/packages/api/src/graphql/modules/LinkedMerkleWitnessResolver.ts @@ -0,0 +1,62 @@ +import { Arg, Field, ObjectType, Query } from "type-graphql"; +import { inject } from "tsyringe"; +import { + LinkedLeafAndMerkleWitness, + LinkedMerkleTree, +} from "@proto-kit/common"; +import { CachedLinkedMerkleTreeStore } from "@proto-kit/sequencer/dist/state/merkle/CachedLinkedMerkleTreeStore"; +import { AsyncLinkedMerkleTreeStore } from "@proto-kit/sequencer/dist/state/async/AsyncLinkedMerkleTreeStore"; + +import { GraphqlModule, graphqlModule } from "../GraphqlModule"; + +import { MerkleWitnessDTO } from "./MerkleWitnessResolver"; +import { LeafDTO } from "./LeafResolver"; + +@ObjectType() +export class LinkedMerkleWitnessDTO { + public static fromServiceLayerObject(witness: LinkedLeafAndMerkleWitness) { + const { leaf, merkleWitness } = witness; + const leafDTO = LeafDTO.fromServiceLayerModel(leaf); + const witnessDTO = MerkleWitnessDTO.fromServiceLayerObject(merkleWitness); + return new LinkedMerkleWitnessDTO(leafDTO, witnessDTO); + } + + public constructor(leaf: LeafDTO, witness: MerkleWitnessDTO) { + this.leaf = leaf; + this.merkleWitness = new MerkleWitnessDTO( + witness.siblings, + witness.isLefts + ); + } + + @Field(() => LeafDTO) + public leaf: LeafDTO; + + @Field(() => MerkleWitnessDTO) + public merkleWitness: MerkleWitnessDTO; +} + +@graphqlModule() +export class LinkedMerkleWitnessResolver extends GraphqlModule { + public constructor( + @inject("AsyncMerkleStore") + private readonly treeStore: AsyncLinkedMerkleTreeStore + ) { + super(); + } + + @Query(() => LinkedMerkleWitnessDTO, { + description: + "Allows retrieval of merkle witnesses corresponding to a specific path in the appchain's state tree. These proves are generally retrieved from the current 'proven' state", + }) + public async witness(@Arg("path") path: string) { + const syncStore = await CachedLinkedMerkleTreeStore.new(this.treeStore); + + const tree = new LinkedMerkleTree(syncStore); + await syncStore.preloadKey(BigInt(path)); + + const witness = tree.getWitness(BigInt(path)); + + return LinkedMerkleWitnessDTO.fromServiceLayerObject(witness); + } +} diff --git a/packages/api/src/index.ts b/packages/api/src/index.ts index 9c7d5dbef..7f734c208 100644 --- a/packages/api/src/index.ts +++ b/packages/api/src/index.ts @@ -9,4 +9,5 @@ export * from "./graphql/modules/NodeStatusResolver"; export * from "./graphql/modules/AdvancedNodeStatusResolver"; export * from "./graphql/services/NodeStatusService"; export * from "./graphql/modules/MerkleWitnessResolver"; +export * from "./graphql/modules/LinkedMerkleWitnessResolver"; export * from "./graphql/VanillaGraphqlModules"; diff --git a/packages/common/package.json b/packages/common/package.json index 6fa29f054..ea6f866c4 100644 --- a/packages/common/package.json +++ b/packages/common/package.json @@ -20,7 +20,8 @@ "lodash": "^4.17.21", "loglevel": "^1.8.1", "reflect-metadata": "^0.1.13", - "typescript-memoize": "^1.1.1" + "typescript-memoize": "^1.1.1", + "ts-mixer": "^6.0.3" }, "peerDependencies": { "o1js": "^1.1.0", diff --git a/packages/common/src/index.ts b/packages/common/src/index.ts index 2309ba770..4b8dfc759 100644 --- a/packages/common/src/index.ts +++ b/packages/common/src/index.ts @@ -13,7 +13,11 @@ export * from "./log"; export * from "./events/EventEmittingComponent"; export * from "./events/EventEmitter"; export * from "./trees/MerkleTreeStore"; +export * from "./trees/LinkedMerkleTreeStore"; export * from "./trees/InMemoryMerkleTreeStorage"; export * from "./trees/RollupMerkleTree"; +export * from "./trees/LinkedMerkleTree"; +export * from "./trees/InMemoryLinkedLeafStore"; export * from "./events/EventEmitterProxy"; export * from "./trees/MockAsyncMerkleStore"; +export * from "./trees/InMemoryLinkedMerkleLeafStore"; diff --git a/packages/common/src/trees/InMemoryLinkedLeafStore.ts b/packages/common/src/trees/InMemoryLinkedLeafStore.ts new file mode 100644 index 000000000..0af12a83c --- /dev/null +++ b/packages/common/src/trees/InMemoryLinkedLeafStore.ts @@ -0,0 +1,40 @@ +import { LinkedLeafStore, LinkedLeaf } from "./LinkedMerkleTreeStore"; + +export class InMemoryLinkedLeafStore implements LinkedLeafStore { + public leaves: { + [key: string]: { leaf: LinkedLeaf; index: bigint }; + } = {}; + + public maximumIndex?: bigint; + + public getLeaf( + path: bigint + ): { leaf: LinkedLeaf; index: bigint } | undefined { + return this.leaves[path.toString()]; + } + + public setLeaf(index: bigint, value: LinkedLeaf): void { + const leaf = this.getLeaf(value.path); + if (leaf !== undefined && leaf?.index !== index) { + throw new Error("Cannot change index of existing leaf"); + } + this.leaves[value.path.toString()] = { leaf: value, index: index }; + if (this.maximumIndex === undefined || index > this.maximumIndex) { + this.maximumIndex = index; + } + } + + public getMaximumIndex(): bigint | undefined { + return this.maximumIndex; + } + + // This gets the leaf with the closest path. + public getLeafLessOrEqual( + path: bigint + ): { leaf: LinkedLeaf; index: bigint } | undefined { + return Object.values(this.leaves).find( + (storedLeaf) => + storedLeaf.leaf.nextPath > path && storedLeaf.leaf.path <= path + ); + } +} diff --git a/packages/common/src/trees/InMemoryLinkedMerkleLeafStore.ts b/packages/common/src/trees/InMemoryLinkedMerkleLeafStore.ts new file mode 100644 index 000000000..6d673c3f4 --- /dev/null +++ b/packages/common/src/trees/InMemoryLinkedMerkleLeafStore.ts @@ -0,0 +1,9 @@ +import { Mixin } from "ts-mixer"; + +import { InMemoryLinkedLeafStore } from "./InMemoryLinkedLeafStore"; +import { InMemoryMerkleTreeStorage } from "./InMemoryMerkleTreeStorage"; + +export class InMemoryLinkedMerkleLeafStore extends Mixin( + InMemoryLinkedLeafStore, + InMemoryMerkleTreeStorage +) {} diff --git a/packages/common/src/trees/InMemoryMerkleTreeStorage.ts b/packages/common/src/trees/InMemoryMerkleTreeStorage.ts index 390b87fc8..c042c0d3a 100644 --- a/packages/common/src/trees/InMemoryMerkleTreeStorage.ts +++ b/packages/common/src/trees/InMemoryMerkleTreeStorage.ts @@ -1,7 +1,7 @@ import { MerkleTreeStore } from "./MerkleTreeStore"; export class InMemoryMerkleTreeStorage implements MerkleTreeStore { - protected nodes: { + public nodes: { [key: number]: { [key: string]: bigint; }; diff --git a/packages/common/src/trees/LinkedMerkleTree.ts b/packages/common/src/trees/LinkedMerkleTree.ts new file mode 100644 index 000000000..66d67ea9f --- /dev/null +++ b/packages/common/src/trees/LinkedMerkleTree.ts @@ -0,0 +1,478 @@ +// eslint-disable-next-line max-classes-per-file +import { Bool, Field, Poseidon, Provable, Struct } from "o1js"; + +import { TypedClass } from "../types"; +import { range } from "../utils"; + +import { LinkedMerkleTreeStore } from "./LinkedMerkleTreeStore"; +import { + AbstractMerkleWitness, + createMerkleTree, + maybeSwap, +} from "./RollupMerkleTree"; +import { InMemoryLinkedMerkleLeafStore } from "./InMemoryLinkedMerkleLeafStore"; + +const RollupMerkleTreeWitness = createMerkleTree(40).WITNESS; +export class LinkedLeafStruct extends Struct({ + value: Field, + path: Field, + nextPath: Field, +}) { + public hash(): Field { + return Poseidon.hash(LinkedLeafStruct.toFields(this)); + } +} + +// We use the RollupMerkleTreeWitness here, although we will actually implement +// the RollupMerkleTreeWitnessV2 defined below when instantiating the class. +export class LinkedLeafAndMerkleWitness extends Struct({ + leaf: LinkedLeafStruct, + merkleWitness: RollupMerkleTreeWitness, +}) {} + +class LinkedStructTemplate extends Struct({ + leafPrevious: LinkedLeafAndMerkleWitness, + leafCurrent: LinkedLeafAndMerkleWitness, +}) {} + +export interface AbstractLinkedMerkleWitness extends LinkedStructTemplate {} + +export interface AbstractLinkedMerkleTree { + store: LinkedMerkleTreeStore; + /** + * Returns a node which lives at a given index and level. + * @param level Level of the node. + * @param index Index of the node. + * @returns The data of the node. + */ + getNode(level: number, index: bigint): Field; + + /** + * Returns the root of the [Merkle Tree](https://en.wikipedia.org/wiki/Merkle_tree). + * @returns The root of the Merkle Tree. + */ + getRoot(): Field; + + /** + * Sets the value of a leaf node at a given index to a given value. + * @param path of the leaf node. + * @param value New value. + */ + setLeaf(path: bigint, value?: bigint): LinkedMerkleTreeWitness; + + /** + * Returns a leaf which lives at a given path. + * Errors otherwise. + * @param path Index of the node. + * @returns The data of the leaf. + */ + getLeaf(path: bigint): LinkedLeafStruct | undefined; + + /** + * Returns the witness (also known as + * [Merkle Proof or Merkle Witness](https://computersciencewiki.org/index.php/Merkle_proof)) + * for the leaf at the given path. + * @param path Position of the leaf node. + * @returns The witness that belongs to the leaf. + */ + getWitness(path: bigint): LinkedLeafAndMerkleWitness; + + dummyWitness(): LinkedMerkleTreeWitness; + + dummy(): LinkedLeafAndMerkleWitness; +} + +export interface AbstractLinkedMerkleTreeClass { + new (store: LinkedMerkleTreeStore): AbstractLinkedMerkleTree; + + WITNESS: TypedClass & + typeof LinkedStructTemplate; + + HEIGHT: number; + + EMPTY_ROOT: bigint; +} + +export function createLinkedMerkleTree( + height: number +): AbstractLinkedMerkleTreeClass { + class LinkedMerkleWitness + extends LinkedStructTemplate + implements AbstractLinkedMerkleWitness {} + /** + * The {@link RollupMerkleWitness} class defines a circuit-compatible base class + * for [Merkle Witness'](https://computersciencewiki.org/index.php/Merkle_proof). + */ + // We define the RollupMerkleWitness again here as we want it to have the same height + // as the tree. If we re-used the Witness from the RollupMerkleTree.ts we wouldn't have + // control, whilst having the overhead of creating the RollupTree, since the witness is + // defined from the tree (for the height reason already described). + class RollupMerkleWitnessV2 + extends Struct({ + path: Provable.Array(Field, height - 1), + isLeft: Provable.Array(Bool, height - 1), + }) + implements AbstractMerkleWitness + { + public static height = height; + + public height(): number { + return RollupMerkleWitnessV2.height; + } + + /** + * Calculates a root depending on the leaf value. + * @param leaf Value of the leaf node that belongs to this Witness. + * @returns The calculated root. + */ + public calculateRoot(leaf: Field): Field { + let hash = leaf; + const n = this.height(); + + for (let index = 1; index < n; ++index) { + const isLeft = this.isLeft[index - 1]; + + const [left, right] = maybeSwap(isLeft, hash, this.path[index - 1]); + hash = Poseidon.hash([left, right]); + } + + return hash; + } + + /** + * Calculates the index of the leaf node that belongs to this Witness. + * @returns Index of the leaf. + */ + public calculateIndex(): Field { + let powerOfTwo = Field(1); + let index = Field(0); + const n = this.height(); + + for (let i = 1; i < n; ++i) { + index = Provable.if(this.isLeft[i - 1], index, index.add(powerOfTwo)); + powerOfTwo = powerOfTwo.mul(2); + } + + return index; + } + + public checkMembership(root: Field, key: Field, value: Field): Bool { + const calculatedRoot = this.calculateRoot(value); + const calculatedKey = this.calculateIndex(); + // We don't have to range-check the key, because if it would be greater + // than leafCount, it would not match the computedKey + key.assertEquals(calculatedKey, "Keys of MerkleWitness does not match"); + return root.equals(calculatedRoot); + } + + public checkMembershipSimple(root: Field, value: Field): Bool { + const calculatedRoot = this.calculateRoot(value); + return root.equals(calculatedRoot); + } + + public checkMembershipGetRoots( + root: Field, + key: Field, + value: Field + ): [Bool, Field, Field] { + const calculatedRoot = this.calculateRoot(value); + const calculatedKey = this.calculateIndex(); + key.assertEquals(calculatedKey, "Keys of MerkleWitness does not match"); + return [root.equals(calculatedRoot), root, calculatedRoot]; + } + + public toShortenedEntries() { + return range(0, 5) + .concat(range(this.height() - 4, this.height())) + .map((index) => + [ + this.path[index].toString(), + this.isLeft[index].toString(), + ].toString() + ); + } + + public static dummy() { + return new RollupMerkleWitnessV2({ + isLeft: Array(this.height - 1).fill(Bool(false)), + path: Array(this.height - 1).fill(Field(0)), + }); + } + } + return class AbstractLinkedRollupMerkleTree + implements AbstractLinkedMerkleTree + { + public static HEIGHT = height; + + public static EMPTY_ROOT = new AbstractLinkedRollupMerkleTree( + new InMemoryLinkedMerkleLeafStore() + ) + .getRoot() + .toBigInt(); + + public static WITNESS = LinkedMerkleWitness; + + readonly zeroes: bigint[]; + + readonly store: LinkedMerkleTreeStore; + + public constructor(store: LinkedMerkleTreeStore) { + this.store = store; + this.zeroes = [0n]; + for ( + let index = 1; + index < AbstractLinkedRollupMerkleTree.HEIGHT; + index += 1 + ) { + const previousLevel = Field(this.zeroes[index - 1]); + this.zeroes.push( + Poseidon.hash([previousLevel, previousLevel]).toBigInt() + ); + } + // 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() === undefined) { + this.setLeafInitialisation(); + } + } + + public getNode(level: number, index: bigint): Field { + const node = this.store.getNode(index, level); + return Field(node ?? this.zeroes[level]); + } + + /** + * Returns leaf which lives at a given path. + * Errors if the path is not defined. + * @param path path of the node. + * @returns The data of the node. + */ + public getLeaf(path: bigint): LinkedLeafStruct | undefined { + const storedLeaf = this.store.getLeaf(path); + if (storedLeaf === undefined) { + return undefined; + } + return new LinkedLeafStruct({ + value: Field(storedLeaf.leaf.value), + path: Field(storedLeaf.leaf.path), + nextPath: Field(storedLeaf.leaf.nextPath), + }); + } + + /** + * Returns the root of the [Merkle Tree](https://en.wikipedia.org/wiki/Merkle_tree). + * @returns The root of the Merkle Tree. + */ + public getRoot(): Field { + return this.getNode( + AbstractLinkedRollupMerkleTree.HEIGHT - 1, + 0n + ).toConstant(); + } + + private setNode(level: number, index: bigint, value: Field) { + this.store.setNode(index, level, value.toBigInt()); + } + + /** + * Sets the value of a leaf node at a given index to a given value + * and carry the change through to the tree. + * @param index Position of the leaf node. + * @param leaf New value. + */ + private setMerkleLeaf(index: bigint, leaf: LinkedLeafStruct) { + this.setNode(0, index, leaf.hash()); + let tempIndex = index; + for ( + let level = 1; + level < AbstractLinkedRollupMerkleTree.HEIGHT; + level += 1 + ) { + tempIndex /= 2n; + const leftPrev = this.getNode(level - 1, tempIndex * 2n); + const rightPrev = this.getNode(level - 1, tempIndex * 2n + 1n); + this.setNode(level, tempIndex, Poseidon.hash([leftPrev, rightPrev])); + } + } + + /** + * Sets the value of a node at a given index to a given value. + * @param path Position of the leaf node. + * @param value New value. + */ + public setLeaf(path: bigint, value?: bigint): LinkedMerkleWitness { + if (value === undefined) { + return new LinkedMerkleWitness({ + leafPrevious: this.dummy(), + leafCurrent: this.getWitness(path), + }); + } + const storedLeaf = this.store.getLeaf(path); + const prevLeaf = this.store.getLeafLessOrEqual(path); + if (prevLeaf === undefined) { + throw Error("Prev leaf shouldn't be undefined"); + } + let witnessPrevious; + let index: bigint; + if (storedLeaf === undefined) { + // The above means the path doesn't already exist, and we are inserting, not updating. + // This requires us to update the node with the previous path, as well. + const tempIndex = this.store.getMaximumIndex(); + if (tempIndex === undefined) { + throw Error("Store Max Index not defined"); + } + if (tempIndex + 1n >= 2 ** height) { + throw new Error("Index greater than maximum leaf number"); + } + witnessPrevious = this.getWitness(prevLeaf.leaf.path); + const newPrevLeaf = { + value: prevLeaf.leaf.value, + path: prevLeaf.leaf.path, + nextPath: path, + }; + this.store.setLeaf(prevLeaf.index, newPrevLeaf); + this.setMerkleLeaf( + prevLeaf.index, + new LinkedLeafStruct({ + value: Field(newPrevLeaf.value), + path: Field(newPrevLeaf.path), + nextPath: Field(newPrevLeaf.nextPath), + }) + ); + + index = tempIndex + 1n; + } else { + witnessPrevious = this.dummy(); + index = storedLeaf.index; + } + const newLeaf = { + value: value, + path: path, + nextPath: prevLeaf.leaf.nextPath, + }; + const witnessNext = this.getWitness(newLeaf.path); + this.store.setLeaf(index, newLeaf); + this.setMerkleLeaf( + index, + new LinkedLeafStruct({ + value: Field(newLeaf.value), + path: Field(newLeaf.path), + nextPath: Field(newLeaf.nextPath), + }) + ); + return new LinkedMerkleWitness({ + leafPrevious: witnessPrevious, + leafCurrent: witnessNext, + }); + } + + /** + * Sets the value of a leaf node at initialisation, + * i.e. {vale: 0, path: 0, nextPath: Field.Max} + */ + private setLeafInitialisation() { + // This is the maximum value of the hash + const MAX_FIELD_VALUE: bigint = Field.ORDER - 1n; + this.store.setLeaf(0n, { + value: 0n, + path: 0n, + nextPath: MAX_FIELD_VALUE, + }); + // We now set the leafs in the merkle tree to cascade the values up + // the tree. + this.setMerkleLeaf( + 0n, + new LinkedLeafStruct({ + value: Field(0n), + path: Field(0n), + nextPath: Field(MAX_FIELD_VALUE), + }) + ); + } + + /** + * Returns the witness (also known as + * [Merkle Proof or Merkle Witness](https://computersciencewiki.org/index.php/Merkle_proof)) + * for the leaf at the given path, otherwise returns a witness for the first unused index. + * @param path of the leaf node. + * @returns The witness that belongs to the leaf. + */ + public getWitness(path: bigint): LinkedLeafAndMerkleWitness { + const storedLeaf = this.store.getLeaf(path); + let leaf; + let currentIndex: bigint; + + if (storedLeaf === undefined) { + const storeIndex = this.store.getMaximumIndex(); + if (storeIndex === undefined) { + throw new Error("Store Undefined"); + } + currentIndex = storeIndex + 1n; + leaf = new LinkedLeafStruct({ + value: Field(0), + path: Field(0), + nextPath: Field(0), + }); + } else { + leaf = new LinkedLeafStruct({ + value: Field(storedLeaf.leaf.value), + path: Field(storedLeaf.leaf.path), + nextPath: Field(storedLeaf.leaf.nextPath), + }); + currentIndex = storedLeaf.index; + } + + const pathArray = []; + const isLefts = []; + + for ( + let level = 0; + level < AbstractLinkedRollupMerkleTree.HEIGHT - 1; + level += 1 + ) { + const isLeft = currentIndex % 2n === 0n; + const sibling = this.getNode( + level, + isLeft ? currentIndex + 1n : currentIndex - 1n + ); + isLefts.push(Bool(isLeft)); + pathArray.push(sibling); + currentIndex /= 2n; + } + return new LinkedLeafAndMerkleWitness({ + merkleWitness: new RollupMerkleWitnessV2({ + path: pathArray, + isLeft: isLefts, + }), + leaf: leaf, + }); + } + + public dummy(): LinkedLeafAndMerkleWitness { + return new LinkedLeafAndMerkleWitness({ + merkleWitness: new RollupMerkleTreeWitness({ + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + path: Array(40).fill(Field(0)) as Field[], + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + isLeft: Array(40).fill(new Bool(true)) as Bool[], + }), + leaf: new LinkedLeafStruct({ + value: Field(0), + path: Field(0), + nextPath: Field(0), + }), + }); + } + + public dummyWitness() { + return new LinkedMerkleWitness({ + leafPrevious: this.dummy(), + leafCurrent: this.dummy(), + }); + } + }; +} + +export class LinkedMerkleTree extends createLinkedMerkleTree(40) {} +export class LinkedMerkleTreeWitness extends LinkedMerkleTree.WITNESS {} diff --git a/packages/common/src/trees/LinkedMerkleTreeStore.ts b/packages/common/src/trees/LinkedMerkleTreeStore.ts new file mode 100644 index 000000000..f8efefd65 --- /dev/null +++ b/packages/common/src/trees/LinkedMerkleTreeStore.ts @@ -0,0 +1,22 @@ +import { MerkleTreeStore } from "./MerkleTreeStore"; + +export interface LinkedLeafStore { + setLeaf: (index: bigint, value: LinkedLeaf) => void; + + getLeaf: (path: bigint) => { leaf: LinkedLeaf; index: bigint } | undefined; + + getLeafLessOrEqual: ( + path: bigint + ) => { leaf: LinkedLeaf; index: bigint } | undefined; + + getMaximumIndex: () => bigint | undefined; +} + +export type LinkedLeaf = { value: bigint; path: bigint; nextPath: bigint }; +export interface LinkedMerkleTreeStore + extends LinkedLeafStore, + MerkleTreeStore {} + +export interface PreloadingLinkedMerkleTreeStore extends LinkedMerkleTreeStore { + preloadKeys(path: bigint[]): Promise; +} diff --git a/packages/common/src/trees/RollupMerkleTree.ts b/packages/common/src/trees/RollupMerkleTree.ts index 44049e805..b8b245d16 100644 --- a/packages/common/src/trees/RollupMerkleTree.ts +++ b/packages/common/src/trees/RollupMerkleTree.ts @@ -6,7 +6,7 @@ import { TypedClass } from "../types"; import { MerkleTreeStore } from "./MerkleTreeStore"; import { InMemoryMerkleTreeStorage } from "./InMemoryMerkleTreeStorage"; -class StructTemplate extends Struct({ +export class StructTemplate extends Struct({ path: Provable.Array(Field, 0), isLeft: Provable.Array(Bool, 0), }) {} @@ -16,7 +16,7 @@ export interface AbstractMerkleWitness extends StructTemplate { /** * Calculates a root depending on the leaf value. - * @param leaf Value of the leaf node that belongs to this Witness. + * @param hash Value of the leaf node that belongs to this Witness. * @returns The calculated root. */ calculateRoot(hash: Field): Field; @@ -29,6 +29,8 @@ export interface AbstractMerkleWitness extends StructTemplate { checkMembership(root: Field, key: Field, value: Field): Bool; + checkMembershipSimple(root: Field, value: Field): Bool; + checkMembershipGetRoots( root: Field, key: Field, @@ -115,7 +117,7 @@ export interface AbstractMerkleTreeClass { */ export function createMerkleTree(height: number): AbstractMerkleTreeClass { /** - * The {@link BaseMerkleWitness} class defines a circuit-compatible base class + * The {@link RollupMerkleWitness} class defines a circuit-compatible base class * for [Merkle Witness'](https://computersciencewiki.org/index.php/Merkle_proof). */ class RollupMerkleWitness @@ -176,6 +178,11 @@ export function createMerkleTree(height: number): AbstractMerkleTreeClass { return root.equals(calculatedRoot); } + public checkMembershipSimple(root: Field, value: Field): Bool { + const calculatedRoot = this.calculateRoot(value); + return root.equals(calculatedRoot); + } + public checkMembershipGetRoots( root: Field, key: Field, @@ -200,12 +207,11 @@ export function createMerkleTree(height: number): AbstractMerkleTreeClass { public static dummy() { return new RollupMerkleWitness({ - isLeft: Array(height - 1).fill(Bool(false)), - path: Array(height - 1).fill(Field(0)), + isLeft: Array(this.height - 1).fill(Bool(false)), + path: Array(this.height - 1).fill(Field(0)), }); } } - return class AbstractRollupMerkleTree implements AbstractMerkleTree { public static HEIGHT = height; @@ -348,7 +354,7 @@ export class RollupMerkleTreeWitness extends RollupMerkleTree.WITNESS {} * More efficient version of `maybeSwapBad` which * reuses an intermediate variable */ -function maybeSwap(b: Bool, x: Field, y: Field): [Field, Field] { +export function maybeSwap(b: Bool, x: Field, y: Field): [Field, Field] { const m = b.toField().mul(x.sub(y)); // b*(x - y) const x1 = y.add(m); // y + b*(x - y) const y2 = x.sub(m); // x - b*(x - y) = x + b*(y - x) diff --git a/packages/common/test/trees/LinkedMerkleTree.test.ts b/packages/common/test/trees/LinkedMerkleTree.test.ts new file mode 100644 index 000000000..755bc4473 --- /dev/null +++ b/packages/common/test/trees/LinkedMerkleTree.test.ts @@ -0,0 +1,123 @@ +import { beforeEach } from "@jest/globals"; +import { Field, Poseidon } from "o1js"; + +import { + createLinkedMerkleTree, + InMemoryLinkedMerkleLeafStore, + log, +} from "../../src"; +import { expectDefined } from "../../dist/utils"; + +describe.each([4, 16, 254])("cachedMerkleTree - %s", (height) => { + class LinkedMerkleTree extends createLinkedMerkleTree(height) {} + + let store: InMemoryLinkedMerkleLeafStore; + let tree: LinkedMerkleTree; + + beforeEach(() => { + log.setLevel("INFO"); + + store = new InMemoryLinkedMerkleLeafStore(); + tree = new LinkedMerkleTree(store); + }); + + it("should have the same root when empty", () => { + expect.assertions(1); + + expect(tree.getRoot().toBigInt()).toStrictEqual( + LinkedMerkleTree.EMPTY_ROOT + ); + }); + + it("should have a different root when not empty", () => { + expect.assertions(1); + + tree.setLeaf(1n, 1n); + + expect(tree.getRoot().toBigInt()).not.toStrictEqual( + LinkedMerkleTree.EMPTY_ROOT + ); + }); + + it("should provide correct witnesses", () => { + expect.assertions(1); + + tree.setLeaf(1n, 1n); + tree.setLeaf(5n, 5n); + + const witness = tree.getWitness(5n); + + expect( + witness.merkleWitness + .calculateRoot( + Poseidon.hash([ + witness.leaf.value, + witness.leaf.path, + witness.leaf.nextPath, + ]) + ) + .toBigInt() + ).toStrictEqual(tree.getRoot().toBigInt()); + }); + + it("should have invalid witnesses with wrong values", () => { + expect.assertions(1); + + tree.setLeaf(1n, 1n); + tree.setLeaf(5n, 5n); + + const witness = tree.getWitness(5n); + + expect( + witness.merkleWitness.calculateRoot(Field(6)).toBigInt() + ).not.toStrictEqual(tree.getRoot().toBigInt()); + }); + + it("should have valid witnesses with changed value on the same leafs", () => { + expect.assertions(1); + + tree.setLeaf(1n, 1n); + tree.setLeaf(5n, 5n); + + const witness = tree.getWitness(5n); + + tree.setLeaf(5n, 10n); + + expect( + witness.merkleWitness + .calculateRoot( + Poseidon.hash([Field(10), witness.leaf.path, witness.leaf.nextPath]) + ) + .toBigInt() + ).toStrictEqual(tree.getRoot().toBigInt()); + }); + + it("should return zeroNode", () => { + expect.assertions(4); + const MAX_FIELD_VALUE: bigint = Field.ORDER - 1n; + const zeroLeaf = tree.getLeaf(0n); + expectDefined(zeroLeaf); + expect(zeroLeaf.value.toBigInt()).toStrictEqual(0n); + expect(zeroLeaf.path.toBigInt()).toStrictEqual(0n); + expect(zeroLeaf.nextPath.toBigInt()).toStrictEqual(MAX_FIELD_VALUE); + }); +}); + +// Separate describe here since we only want small trees for this test. +describe("Error check", () => { + class LinkedMerkleTree extends createLinkedMerkleTree(4) {} + let store: InMemoryLinkedMerkleLeafStore; + let tree: LinkedMerkleTree; + + it("throw for invalid index", () => { + log.setLevel("INFO"); + + store = new InMemoryLinkedMerkleLeafStore(); + tree = new LinkedMerkleTree(store); + expect(() => { + for (let i = 0; i < 2n ** BigInt(4) + 1n; i++) { + tree.setLeaf(BigInt(i), 2n); + } + }).toThrow("Index greater than maximum leaf number"); + }); +}); diff --git a/packages/persistance/src/RedisConnection.ts b/packages/persistance/src/RedisConnection.ts index e9bd6a9f8..80e10eeb0 100644 --- a/packages/persistance/src/RedisConnection.ts +++ b/packages/persistance/src/RedisConnection.ts @@ -5,7 +5,7 @@ import { } from "@proto-kit/sequencer"; import { DependencyFactory } from "@proto-kit/common"; -import { RedisMerkleTreeStore } from "./services/redis/RedisMerkleTreeStore"; +import { RedisLinkedMerkleTreeStore } from "./services/redis/RedisLinkedMerkleTreeStore"; export interface RedisConnectionConfig { host: string; @@ -39,13 +39,13 @@ export class RedisConnectionModule > { return { asyncMerkleStore: { - useFactory: () => new RedisMerkleTreeStore(this), + useFactory: () => new RedisLinkedMerkleTreeStore(this), }, unprovenMerkleStore: { - useFactory: () => new RedisMerkleTreeStore(this, "unproven"), + useFactory: () => new RedisLinkedMerkleTreeStore(this, "unproven"), }, blockTreeStore: { - useFactory: () => new RedisMerkleTreeStore(this, "blockHash"), + useFactory: () => new RedisLinkedMerkleTreeStore(this, "blockHash"), }, }; } diff --git a/packages/persistance/src/index.ts b/packages/persistance/src/index.ts index 208211217..51ef2c6f9 100644 --- a/packages/persistance/src/index.ts +++ b/packages/persistance/src/index.ts @@ -15,3 +15,4 @@ export * from "./services/prisma/mappers/StateTransitionMapper"; export * from "./services/prisma/mappers/TransactionMapper"; export * from "./services/prisma/mappers/BlockResultMapper"; export * from "./services/redis/RedisMerkleTreeStore"; +export * from "./services/redis/RedisLinkedMerkleTreeStore"; diff --git a/packages/persistance/src/services/redis/RedisLinkedMerkleTreeStore.ts b/packages/persistance/src/services/redis/RedisLinkedMerkleTreeStore.ts new file mode 100644 index 000000000..7616a546a --- /dev/null +++ b/packages/persistance/src/services/redis/RedisLinkedMerkleTreeStore.ts @@ -0,0 +1,83 @@ +import { MerkleTreeNode, MerkleTreeNodeQuery } from "@proto-kit/sequencer"; +import { LinkedLeaf, log, noop } from "@proto-kit/common"; +import { AsyncLinkedMerkleTreeStore } from "@proto-kit/sequencer/dist/state/async/AsyncLinkedMerkleTreeStore"; + +import type { RedisConnection } from "../../RedisConnection"; + +export class RedisLinkedMerkleTreeStore implements AsyncLinkedMerkleTreeStore { + private cache: MerkleTreeNode[] = []; + + public constructor( + private readonly connection: RedisConnection, + private readonly mask: string = "base" + ) {} + + private getKey(node: MerkleTreeNodeQuery): string { + return `${this.mask}:${node.level}:${node.key.toString()}`; + } + + public async openTransaction(): Promise { + noop(); + } + + public async commit(): Promise { + const start = Date.now(); + const array: [string, string][] = this.cache.map( + ({ key, level, value }) => [this.getKey({ key, level }), value.toString()] + ); + + if (array.length === 0) { + return; + } + + try { + await this.connection.redisClient.mSet(array.flat(1)); + } catch (error) { + log.error(error); + } + log.trace( + `Committing ${array.length} kv-pairs took ${Date.now() - start} ms` + ); + + this.cache = []; + } + + public async getNodesAsync( + nodes: MerkleTreeNodeQuery[] + ): Promise<(bigint | undefined)[]> { + if (nodes.length === 0) { + return []; + } + + const keys = nodes.map((node) => this.getKey(node)); + + const result = await this.connection.redisClient.mGet(keys); + + return result.map((x) => (x !== null ? BigInt(x) : undefined)); + } + + public writeNodes(nodes: MerkleTreeNode[]): void { + this.cache = this.cache.concat(nodes); + } + + public writeLeaves(leaves: { leaf: LinkedLeaf; index: bigint }[]) {} + + public getLeavesAsync(paths: bigint[]) { + return Promise.resolve([undefined]); + } + + public getMaximumIndexAsync() { + return Promise.resolve(0n); + } + + public getLeafLessOrEqualAsync(path: bigint) { + return Promise.resolve({ + leaf: { + value: 0n, + path: 0n, + nextPath: 0n, + }, + index: 0n, + }); + } +} diff --git a/packages/protocol/src/model/StateTransitionProvableBatch.ts b/packages/protocol/src/model/StateTransitionProvableBatch.ts index cbd7283be..1620610f2 100644 --- a/packages/protocol/src/model/StateTransitionProvableBatch.ts +++ b/packages/protocol/src/model/StateTransitionProvableBatch.ts @@ -1,10 +1,9 @@ import { Bool, Provable, Struct } from "o1js"; +import { InMemoryLinkedMerkleLeafStore, range } from "@proto-kit/common"; import { - InMemoryMerkleTreeStorage, - range, - RollupMerkleTree, - RollupMerkleTreeWitness, -} from "@proto-kit/common"; + LinkedMerkleTree, + LinkedMerkleTreeWitness, +} from "@proto-kit/common/dist/trees/LinkedMerkleTree"; import { constants } from "../Constants"; @@ -67,7 +66,7 @@ export class StateTransitionProvableBatch extends Struct({ ), merkleWitnesses: Provable.Array( - RollupMerkleTreeWitness, + LinkedMerkleTreeWitness, constants.stateTransitionProverBatchSize ), }) { @@ -76,7 +75,7 @@ export class StateTransitionProvableBatch extends Struct({ transition: ProvableStateTransition; type: ProvableStateTransitionType; }[], - merkleWitnesses: RollupMerkleTreeWitness[] + merkleWitnesses: LinkedMerkleTreeWitness[] ): StateTransitionProvableBatch { const batch = transitions.map((entry) => entry.transition); const transitionTypes = transitions.map((entry) => entry.type); @@ -96,9 +95,7 @@ export class StateTransitionProvableBatch extends Struct({ batch.push(ProvableStateTransition.dummy()); transitionTypes.push(ProvableStateTransitionType.normal); witnesses.push( - new RollupMerkleTree(new InMemoryMerkleTreeStorage()).getWitness( - BigInt(0) - ) + new LinkedMerkleTree(new InMemoryLinkedMerkleLeafStore()).dummyWitness() ); } return new StateTransitionProvableBatch({ @@ -111,7 +108,7 @@ export class StateTransitionProvableBatch extends Struct({ public static fromTransitions( transitions: ProvableStateTransition[], protocolTransitions: ProvableStateTransition[], - merkleWitnesses: RollupMerkleTreeWitness[] + merkleWitnesses: LinkedMerkleTreeWitness[] ): StateTransitionProvableBatch { const array = transitions.slice().concat(protocolTransitions); @@ -138,7 +135,7 @@ export class StateTransitionProvableBatch extends Struct({ private constructor(object: { batch: ProvableStateTransition[]; transitionTypes: ProvableStateTransitionType[]; - merkleWitnesses: RollupMerkleTreeWitness[]; + merkleWitnesses: LinkedMerkleTreeWitness[]; }) { super(object); } diff --git a/packages/protocol/src/prover/statetransition/StateTransitionProver.ts b/packages/protocol/src/prover/statetransition/StateTransitionProver.ts index 51301625a..11698a524 100644 --- a/packages/protocol/src/prover/statetransition/StateTransitionProver.ts +++ b/packages/protocol/src/prover/statetransition/StateTransitionProver.ts @@ -1,12 +1,13 @@ import { AreProofsEnabled, + LinkedLeafStruct, PlainZkProgram, provableMethod, - RollupMerkleTreeWitness, ZkProgrammable, } from "@proto-kit/common"; -import { Field, Provable, SelfProof, ZkProgram } from "o1js"; +import { Bool, Field, Provable, SelfProof, ZkProgram } from "o1js"; import { injectable } from "tsyringe"; +import { LinkedMerkleTreeWitness } from "@proto-kit/common/dist/trees/LinkedMerkleTree"; import { constants } from "../../Constants"; import { ProvableStateTransition } from "../../model/StateTransition"; @@ -157,7 +158,7 @@ export class StateTransitionProverProgrammable extends ZkProgrammable< const transitions = transitionBatch.batch; const types = transitionBatch.transitionTypes; - const merkleWitness = transitionBatch.merkleWitnesses; + const { merkleWitnesses } = transitionBatch; for ( let index = 0; index < constants.stateTransitionProverBatchSize; @@ -167,7 +168,7 @@ export class StateTransitionProverProgrammable extends ZkProgrammable< state, transitions[index], types[index], - merkleWitness[index], + merkleWitnesses[index], index ); } @@ -183,26 +184,128 @@ export class StateTransitionProverProgrammable extends ZkProgrammable< state: StateTransitionProverExecutionState, transition: ProvableStateTransition, type: ProvableStateTransitionType, - merkleWitness: RollupMerkleTreeWitness, + merkleWitness: LinkedMerkleTreeWitness, index = 0 ) { - const membershipValid = merkleWitness.checkMembership( + const isUpdate = merkleWitness.leafPrevious.leaf.nextPath.equals(Field(0)); + + const isDummy = transition.path.equals(0); + const isNotDummy = isDummy.not(); + + // The following checks if this is an update or insert + // If it's an update then the leafCurrent will be the current leaf, + // rather than the zero/dummy leaf if it's an insert. + // If it's an insert then we need to check the leafPrevious is + // a valid leaf, i.e. path is less than transition.path and nextPath + // greater than transition.path. + // Even if we're just reading (rather than writing) then we expect + // the path for the current leaf to be populated. + const pathValid = Provable.if( + isUpdate, // nextPath equal to 0 only if it's a dummy., which is when we update + merkleWitness.leafCurrent.leaf.path.equals(transition.path), // update + merkleWitness.leafPrevious.leaf.path + .lessThan(transition.path) + .and( + merkleWitness.leafPrevious.leaf.nextPath.greaterThan(transition.path) + ) // insert + ); + // This is for dummy STs + Provable.if(isNotDummy, pathValid, new Bool(true)).assertTrue(); + + // Only if we're doing an insert is this valid. + const previousWitnessValid = + merkleWitness.leafPrevious.merkleWitness.checkMembershipSimple( + state.stateRoot, + merkleWitness.leafPrevious.leaf.hash() + ); + + // Combine previousWitnessValid and if it's an update + // it should just be true, as the prev leaf is just a dummy leaf + // so should always be true. + const prevWitnessOrCurrentWitness = Provable.if( + isUpdate, + Bool(true), + previousWitnessValid + ); + + // We need to check the sequencer had fetched the correct previousLeaf, + // specifically that the previousLeaf is what is verified. + // We check the stateRoot matches. + // For an insert the prev leaf is not a dummy, + // and for an update the prev leaf is a dummy. + + // We assert that the previous witness is valid in case of this one being an update + Provable.if( + isNotDummy, + prevWitnessOrCurrentWitness, + Bool(true) + ).assertTrue(); + + // Need to calculate the new state root after the previous leaf is changed. + // This is only relevant if it's an insert. If an update, we will just use + // the existing state root. + const rootWithLeafChanged = + merkleWitness.leafPrevious.merkleWitness.calculateRoot( + new LinkedLeafStruct({ + value: merkleWitness.leafPrevious.leaf.value, + path: merkleWitness.leafPrevious.leaf.path, + nextPath: transition.path, + }).hash() + ); + + const rootAfterFirstStep = Provable.if( + isUpdate, state.stateRoot, - transition.path, - transition.from.value + rootWithLeafChanged ); - membershipValid - .or(transition.from.isSome.not()) - .assertTrue( - errors.merkleWitnessNotCorrect( - index, - type.isNormal().toBoolean() ? "normal" : "protocol" - ) + // Need to check the second leaf is correct, i.e. leafCurrent. + // is what the sequencer claims it is. + // Again, we check whether we have an update or insert as the value + // depends on this. If insert then we have the current path would be 0. + // We use the existing state root if it's only an update as the prev leaf + // wouldn't have changed and therefore the state root should be the same. + const currentWitnessLeaf = Provable.if( + isUpdate, + new LinkedLeafStruct({ + value: transition.from.value, + path: transition.path, + nextPath: merkleWitness.leafCurrent.leaf.nextPath, + }).hash(), + Field(0) + ); + const currentWitnessValid = + merkleWitness.leafCurrent.merkleWitness.checkMembershipSimple( + rootAfterFirstStep, + currentWitnessLeaf ); - const newRoot = merkleWitness.calculateRoot(transition.to.value); + Provable.if(isNotDummy, currentWitnessValid, Bool(true)).assertTrue(); + // Compute the new final root. + // For an insert we have to hash the new leaf and use the leafPrev's nextPath + // For an update we just use the new value, but keep the leafCurrents + // next path the same. + const newCurrentNextPath = Provable.if( + isUpdate, + merkleWitness.leafCurrent.leaf.nextPath, + merkleWitness.leafPrevious.leaf.nextPath + ); + + const newCurrentLeaf = new LinkedLeafStruct({ + value: transition.to.value, + path: transition.path, + nextPath: newCurrentNextPath, + }); + + const newRoot = merkleWitness.leafCurrent.merkleWitness.calculateRoot( + newCurrentLeaf.hash() + ); + + // TODO Make sure that path == 0 -> both isSomes == false + + // This is checking if we have a read or write. + // If a read the state root should stay the same. state.stateRoot = Provable.if( transition.to.isSome, newRoot, @@ -217,8 +320,6 @@ export class StateTransitionProverProgrammable extends ZkProgrammable< state.protocolStateRoot ); - const isNotDummy = transition.path.equals(Field(0)).not(); - state.stateTransitionList.pushIf( transition, isNotDummy.and(type.isNormal()) 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/sdk/src/graphql/GraphqlQueryTransportModule.ts b/packages/sdk/src/graphql/GraphqlQueryTransportModule.ts index 7a28965e8..c69122f74 100644 --- a/packages/sdk/src/graphql/GraphqlQueryTransportModule.ts +++ b/packages/sdk/src/graphql/GraphqlQueryTransportModule.ts @@ -2,7 +2,7 @@ import { QueryTransportModule } from "@proto-kit/sequencer"; import { Field } from "o1js"; import { inject, injectable } from "tsyringe"; import { gql } from "@urql/core"; -import { RollupMerkleTreeWitness } from "@proto-kit/common"; +import { LinkedLeafAndMerkleWitness } from "@proto-kit/common"; import { AppChainModule } from "../appChain/AppChainModule"; @@ -64,12 +64,19 @@ export class GraphqlQueryTransportModule public async merkleWitness( key: Field - ): Promise { + ): Promise { const query = gql` query Witness($path: String!) { witness(path: $path) { - siblings - isLefts + leaf { + value + path + nextPath + } + merkleWitness { + siblings + isLefts + } } } `; @@ -87,21 +94,26 @@ export class GraphqlQueryTransportModule } if ( - witnessJson.siblings === undefined || - witnessJson.isLefts === undefined + witnessJson.leaf === undefined || + witnessJson.merkleWitness.siblings === undefined || + witnessJson.merkleWitness.isLefts === undefined ) { throw new Error("Witness json object malformed"); } - assertStringArray(witnessJson.siblings); - assertBooleanArray(witnessJson.isLefts); + assertStringArray(witnessJson.merkleWitness.siblings); + assertBooleanArray(witnessJson.merkleWitness.isLefts); - return new RollupMerkleTreeWitness( - RollupMerkleTreeWitness.fromJSON({ + return new LinkedLeafAndMerkleWitness( + LinkedLeafAndMerkleWitness.fromJSON({ // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment - path: witnessJson.siblings, - // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment - isLeft: witnessJson.isLefts, + leaf: witnessJson.leaf, + merkleWitness: { + // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment + path: witnessJson.merkleWitness.siblings, + // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment + isLeft: witnessJson.merkleWitness.isLefts, + }, }) ); } diff --git a/packages/sdk/src/query/StateServiceQueryModule.ts b/packages/sdk/src/query/StateServiceQueryModule.ts index ec78f44aa..deb4194e8 100644 --- a/packages/sdk/src/query/StateServiceQueryModule.ts +++ b/packages/sdk/src/query/StateServiceQueryModule.ts @@ -1,14 +1,17 @@ import { AsyncStateService, - CachedMerkleTreeStore, QueryTransportModule, Sequencer, SequencerModulesRecord, - AsyncMerkleTreeStore, } from "@proto-kit/sequencer"; import { Field } from "o1js"; import { inject, injectable } from "tsyringe"; -import { RollupMerkleTree, RollupMerkleTreeWitness } from "@proto-kit/common"; +import { + LinkedLeafAndMerkleWitness, + LinkedMerkleTree, +} from "@proto-kit/common"; +import { CachedLinkedMerkleTreeStore } from "@proto-kit/sequencer/dist/state/merkle/CachedLinkedMerkleTreeStore"; +import { AsyncLinkedMerkleTreeStore } from "@proto-kit/sequencer/dist/state/async/AsyncLinkedMerkleTreeStore"; import { AppChainModule } from "../appChain/AppChainModule"; @@ -29,8 +32,8 @@ export class StateServiceQueryModule ); } - public get treeStore(): AsyncMerkleTreeStore { - return this.sequencer.dependencyContainer.resolve("AsyncMerkleStore"); + public get treeStore(): AsyncLinkedMerkleTreeStore { + return this.sequencer.dependencyContainer.resolve("AsyncLinkedMerkleStore"); } public get(key: Field) { @@ -39,11 +42,11 @@ export class StateServiceQueryModule public async merkleWitness( path: Field - ): Promise { - const syncStore = new CachedMerkleTreeStore(this.treeStore); + ): Promise { + const syncStore = await CachedLinkedMerkleTreeStore.new(this.treeStore); await syncStore.preloadKey(path.toBigInt()); - const tree = new RollupMerkleTree(syncStore); + const tree = new LinkedMerkleTree(syncStore); return tree.getWitness(path.toBigInt()); } diff --git a/packages/sequencer/src/helpers/query/QueryBuilderFactory.ts b/packages/sequencer/src/helpers/query/QueryBuilderFactory.ts index eb69ca104..1029f5d57 100644 --- a/packages/sequencer/src/helpers/query/QueryBuilderFactory.ts +++ b/packages/sequencer/src/helpers/query/QueryBuilderFactory.ts @@ -1,4 +1,4 @@ -import { TypedClass, RollupMerkleTreeWitness } from "@proto-kit/common"; +import { TypedClass, LinkedLeafAndMerkleWitness } from "@proto-kit/common"; import { Runtime, RuntimeModule, @@ -24,13 +24,13 @@ export type PickByType = { export interface QueryGetterState { get: () => Promise; path: () => string; - merkleWitness: () => Promise; + merkleWitness: () => Promise; } export interface QueryGetterStateMap { get: (key: Key) => Promise; path: (key: Key) => string; - merkleWitness: (key: Key) => Promise; + merkleWitness: (key: Key) => Promise; } export type PickStateProperties = PickByType>; diff --git a/packages/sequencer/src/helpers/query/QueryTransportModule.ts b/packages/sequencer/src/helpers/query/QueryTransportModule.ts index eae3290c3..0998d1194 100644 --- a/packages/sequencer/src/helpers/query/QueryTransportModule.ts +++ b/packages/sequencer/src/helpers/query/QueryTransportModule.ts @@ -1,7 +1,9 @@ import { Field } from "o1js"; -import { RollupMerkleTreeWitness } from "@proto-kit/common"; +import { LinkedLeafAndMerkleWitness } from "@proto-kit/common"; export interface QueryTransportModule { get: (key: Field) => Promise; - merkleWitness: (key: Field) => Promise; + merkleWitness: ( + key: Field + ) => Promise; } diff --git a/packages/sequencer/src/protocol/production/BatchProducerModule.ts b/packages/sequencer/src/protocol/production/BatchProducerModule.ts index 967d34ac4..95561a0b6 100644 --- a/packages/sequencer/src/protocol/production/BatchProducerModule.ts +++ b/packages/sequencer/src/protocol/production/BatchProducerModule.ts @@ -8,7 +8,7 @@ import { NetworkState, } from "@proto-kit/protocol"; import { Field, Proof } from "o1js"; -import { log, noop, RollupMerkleTree } from "@proto-kit/common"; +import { LinkedMerkleTree, log, noop } from "@proto-kit/common"; import { sequencerModule, @@ -17,11 +17,12 @@ import { import { BatchStorage } from "../../storage/repositories/BatchStorage"; import { SettleableBatch } from "../../storage/model/Batch"; import { CachedStateService } from "../../state/state/CachedStateService"; -import { CachedMerkleTreeStore } from "../../state/merkle/CachedMerkleTreeStore"; import { AsyncStateService } from "../../state/async/AsyncStateService"; import { AsyncMerkleTreeStore } from "../../state/async/AsyncMerkleTreeStore"; import { BlockResult, BlockWithResult } from "../../storage/model/Block"; import { VerificationKeyService } from "../runtime/RuntimeVerificationKeyService"; +import { CachedLinkedMerkleTreeStore } from "../../state/merkle/CachedLinkedMerkleTreeStore"; +import { AsyncLinkedMerkleTreeStore } from "../../state/async/AsyncLinkedMerkleTreeStore"; import { BlockProverParameters } from "./tasks/BlockProvingTask"; import { StateTransitionProofParameters } from "./tasks/StateTransitionTaskParameters"; @@ -53,7 +54,7 @@ export interface BlockWithPreviousResult { interface BatchMetadata { batch: SettleableBatch; stateService: CachedStateService; - merkleStore: CachedMerkleTreeStore; + merkleStore: CachedLinkedMerkleTreeStore; } const errors = { @@ -77,7 +78,7 @@ export class BatchProducerModule extends SequencerModule { @inject("AsyncStateService") private readonly asyncStateService: AsyncStateService, @inject("AsyncMerkleStore") - private readonly merkleStore: AsyncMerkleTreeStore, + private readonly merkleStore: AsyncLinkedMerkleTreeStore, @inject("BatchStorage") private readonly batchStorage: BatchStorage, @inject("BlockTreeStore") private readonly blockTreeStore: AsyncMerkleTreeStore, @@ -212,7 +213,7 @@ export class BatchProducerModule extends SequencerModule { ): Promise<{ proof: Proof; stateService: CachedStateService; - merkleStore: CachedMerkleTreeStore; + merkleStore: CachedLinkedMerkleTreeStore; fromNetworkState: NetworkState; toNetworkState: NetworkState; }> { @@ -222,7 +223,7 @@ export class BatchProducerModule extends SequencerModule { const stateServices = { stateService: new CachedStateService(this.asyncStateService), - merkleStore: new CachedMerkleTreeStore(this.merkleStore), + merkleStore: await CachedLinkedMerkleTreeStore.new(this.merkleStore), }; const blockTraces: BlockTrace[] = []; @@ -267,7 +268,7 @@ export class BatchProducerModule extends SequencerModule { this.blockTreeStore, Field( blockWithPreviousResult.lastBlockResult?.stateRoot ?? - RollupMerkleTree.EMPTY_ROOT + LinkedMerkleTree.EMPTY_ROOT ), blockWithPreviousResult.block ); diff --git a/packages/sequencer/src/protocol/production/TransactionTraceService.ts b/packages/sequencer/src/protocol/production/TransactionTraceService.ts index 0ccb5c136..de6b1203d 100644 --- a/packages/sequencer/src/protocol/production/TransactionTraceService.ts +++ b/packages/sequencer/src/protocol/production/TransactionTraceService.ts @@ -10,20 +10,21 @@ import { StateTransitionProverPublicInput, StateTransitionType, } from "@proto-kit/protocol"; -import { MAX_FIELD, RollupMerkleTree } from "@proto-kit/common"; +import { MAX_FIELD } from "@proto-kit/common"; import { Bool, Field } from "o1js"; import chunk from "lodash/chunk"; +import { LinkedMerkleTree } from "@proto-kit/common/dist/trees/LinkedMerkleTree"; import { distinctByString } from "../../helpers/utils"; -import { CachedMerkleTreeStore } from "../../state/merkle/CachedMerkleTreeStore"; import { CachedStateService } from "../../state/state/CachedStateService"; -import { SyncCachedMerkleTreeStore } from "../../state/merkle/SyncCachedMerkleTreeStore"; import type { TransactionExecutionResult, BlockWithResult, } from "../../storage/model/Block"; import { AsyncMerkleTreeStore } from "../../state/async/AsyncMerkleTreeStore"; import { VerificationKeyService } from "../runtime/RuntimeVerificationKeyService"; +import { CachedLinkedMerkleTreeStore } from "../../state/merkle/CachedLinkedMerkleTreeStore"; +import { SyncCachedLinkedMerkleTreeStore } from "../../state/merkle/SyncCachedLinkedMerkleTreeStore"; import type { TransactionTrace, BlockTrace } from "./BatchProducerModule"; import { StateTransitionProofParameters } from "./tasks/StateTransitionTaskParameters"; @@ -79,7 +80,7 @@ export class TransactionTraceService { traces: TransactionTrace[], stateServices: { stateService: CachedStateService; - merkleStore: CachedMerkleTreeStore; + merkleStore: CachedLinkedMerkleTreeStore; }, blockHashTreeStore: AsyncMerkleTreeStore, beforeBlockStateRoot: Field, @@ -105,8 +106,8 @@ export class TransactionTraceService { await stateServices.merkleStore.preloadKey(0n); fromStateRoot = Field( - stateServices.merkleStore.getNode(0n, RollupMerkleTree.HEIGHT - 1) ?? - RollupMerkleTree.EMPTY_ROOT + stateServices.merkleStore.getNode(0n, LinkedMerkleTree.HEIGHT - 1) ?? + LinkedMerkleTree.EMPTY_ROOT ); stParameters = [ @@ -173,7 +174,7 @@ export class TransactionTraceService { executionResult: TransactionExecutionResult, stateServices: { stateService: CachedStateService; - merkleStore: CachedMerkleTreeStore; + merkleStore: CachedLinkedMerkleTreeStore; }, verificationKeyService: VerificationKeyService, networkState: NetworkState, @@ -258,7 +259,7 @@ export class TransactionTraceService { } private async createMerkleTrace( - merkleStore: CachedMerkleTreeStore, + merkleStore: CachedLinkedMerkleTreeStore, stateTransitions: UntypedStateTransition[], protocolTransitions: UntypedStateTransition[], runtimeSuccess: boolean @@ -268,14 +269,15 @@ export class TransactionTraceService { }> { const keys = this.allKeys(protocolTransitions.concat(stateTransitions)); - const runtimeSimulationMerkleStore = new SyncCachedMerkleTreeStore( + const tree = new LinkedMerkleTree(merkleStore); + // TODO Consolidate + const runtimeSimulationMerkleStore = new SyncCachedLinkedMerkleTreeStore( merkleStore ); await merkleStore.preloadKeys(keys.map((key) => key.toBigInt())); - const tree = new RollupMerkleTree(merkleStore); - const runtimeTree = new RollupMerkleTree(runtimeSimulationMerkleStore); + const runtimeTree = new LinkedMerkleTree(runtimeSimulationMerkleStore); // const runtimeTree = new RollupMerkleTree(merkleStore); const initialRoot = tree.getRoot(); @@ -322,18 +324,23 @@ export class TransactionTraceService { const provableTransition = transition.toProvable(); - const witness = usedTree.getWitness(provableTransition.path.toBigInt()); + let witness; if (provableTransition.to.isSome.toBoolean()) { - usedTree.setLeaf( + witness = usedTree.setLeaf( provableTransition.path.toBigInt(), - provableTransition.to.value + provableTransition.to.value.toBigInt() ); stateRoot = usedTree.getRoot(); if (StateTransitionType.isProtocol(type)) { protocolStateRoot = stateRoot; } + } else { + witness = usedTree.setLeaf( + provableTransition.path.toBigInt(), + undefined + ); } // Push transition to respective hashlist diff --git a/packages/sequencer/src/protocol/production/sequencing/BlockProducerModule.ts b/packages/sequencer/src/protocol/production/sequencing/BlockProducerModule.ts index 106833f48..c6a145279 100644 --- a/packages/sequencer/src/protocol/production/sequencing/BlockProducerModule.ts +++ b/packages/sequencer/src/protocol/production/sequencing/BlockProducerModule.ts @@ -16,11 +16,12 @@ import { } from "../../../sequencer/builder/SequencerModule"; import { BlockQueue } from "../../../storage/repositories/BlockStorage"; import { PendingTransaction } from "../../../mempool/PendingTransaction"; -import { AsyncMerkleTreeStore } from "../../../state/async/AsyncMerkleTreeStore"; import { AsyncStateService } from "../../../state/async/AsyncStateService"; import { Block, BlockWithResult } from "../../../storage/model/Block"; import { CachedStateService } from "../../../state/state/CachedStateService"; import { MessageStorage } from "../../../storage/repositories/MessageStorage"; +import { AsyncLinkedMerkleTreeStore } from "../../../state/async/AsyncLinkedMerkleTreeStore"; +import { AsyncMerkleTreeStore } from "../../../state/async/AsyncMerkleTreeStore"; import { TransactionExecutionService } from "./TransactionExecutionService"; @@ -39,7 +40,7 @@ export class BlockProducerModule extends SequencerModule { @inject("UnprovenStateService") private readonly unprovenStateService: AsyncStateService, @inject("UnprovenMerkleStore") - private readonly unprovenMerkleStore: AsyncMerkleTreeStore, + private readonly unprovenMerkleStore: AsyncLinkedMerkleTreeStore, @inject("BlockQueue") private readonly blockQueue: BlockQueue, @inject("BlockTreeStore") diff --git a/packages/sequencer/src/protocol/production/sequencing/TransactionExecutionService.ts b/packages/sequencer/src/protocol/production/sequencing/TransactionExecutionService.ts index dae2a7b7a..f957c1e35 100644 --- a/packages/sequencer/src/protocol/production/sequencing/TransactionExecutionService.ts +++ b/packages/sequencer/src/protocol/production/sequencing/TransactionExecutionService.ts @@ -25,8 +25,8 @@ import { Bool, Field, Poseidon } from "o1js"; import { AreProofsEnabled, log, - RollupMerkleTree, mapSequential, + LinkedMerkleTree, } from "@proto-kit/common"; import { MethodParameterEncoder, @@ -38,8 +38,6 @@ import { import { PendingTransaction } from "../../../mempool/PendingTransaction"; import { CachedStateService } from "../../../state/state/CachedStateService"; import { distinctByString } from "../../../helpers/utils"; -import { CachedMerkleTreeStore } from "../../../state/merkle/CachedMerkleTreeStore"; -import { AsyncMerkleTreeStore } from "../../../state/async/AsyncMerkleTreeStore"; import { TransactionExecutionResult, Block, @@ -48,6 +46,10 @@ import { } from "../../../storage/model/Block"; import { UntypedStateTransition } from "../helpers/UntypedStateTransition"; import type { StateRecord } from "../BatchProducerModule"; +import { CachedLinkedMerkleTreeStore } from "../../../state/merkle/CachedLinkedMerkleTreeStore"; +import { AsyncLinkedMerkleTreeStore } from "../../../state/async/AsyncLinkedMerkleTreeStore"; +import { CachedMerkleTreeStore } from "../../../state/merkle/CachedMerkleTreeStore"; +import { AsyncMerkleTreeStore } from "../../../state/async/AsyncMerkleTreeStore"; const errors = { methodIdNotFound: (methodId: string) => @@ -322,7 +324,7 @@ export class TransactionExecutionService { public async generateMetadataForNextBlock( block: Block, - merkleTreeStore: AsyncMerkleTreeStore, + merkleTreeStore: AsyncLinkedMerkleTreeStore, blockHashTreeStore: AsyncMerkleTreeStore, modifyTreeStore = true ): Promise { @@ -339,8 +341,9 @@ export class TransactionExecutionService { return Object.assign(accumulator, diff); }, {}); - const inMemoryStore = new CachedMerkleTreeStore(merkleTreeStore); - const tree = new RollupMerkleTree(inMemoryStore); + const inMemoryStore = + await CachedLinkedMerkleTreeStore.new(merkleTreeStore); + const tree = new LinkedMerkleTree(inMemoryStore); const blockHashInMemoryStore = new CachedMerkleTreeStore( blockHashTreeStore ); @@ -359,7 +362,7 @@ export class TransactionExecutionService { Object.entries(combinedDiff).forEach(([key, state]) => { const treeValue = state !== undefined ? Poseidon.hash(state) : Field(0); - tree.setLeaf(BigInt(key), treeValue); + tree.setLeaf(BigInt(key), treeValue.toBigInt()); }); const stateRoot = tree.getRoot(); diff --git a/packages/sequencer/src/protocol/production/tasks/StateTransitionTaskParameters.ts b/packages/sequencer/src/protocol/production/tasks/StateTransitionTaskParameters.ts index 3408f7b5a..5c17bda3a 100644 --- a/packages/sequencer/src/protocol/production/tasks/StateTransitionTaskParameters.ts +++ b/packages/sequencer/src/protocol/production/tasks/StateTransitionTaskParameters.ts @@ -3,8 +3,8 @@ import { ProvableStateTransitionType, StateTransitionProverPublicInput, } from "@proto-kit/protocol"; -import { RollupMerkleTreeWitness } from "@proto-kit/common"; import { Bool } from "o1js"; +import { LinkedMerkleTreeWitness } from "@proto-kit/common/dist/trees/LinkedMerkleTree"; import { TaskSerializer } from "../../../worker/flow/Task"; @@ -14,7 +14,7 @@ export interface StateTransitionProofParameters { transition: ProvableStateTransition; type: ProvableStateTransitionType; }[]; - merkleWitnesses: RollupMerkleTreeWitness[]; + merkleWitnesses: LinkedMerkleTreeWitness[]; } interface StateTransitionParametersJSON { @@ -23,7 +23,7 @@ interface StateTransitionParametersJSON { transition: ReturnType; type: boolean; }[]; - merkleWitnesses: ReturnType[]; + merkleWitnesses: ReturnType[]; } export class StateTransitionParametersSerializer @@ -43,7 +43,7 @@ export class StateTransitionParametersSerializer }), merkleWitnesses: parameters.merkleWitnesses.map((witness) => - RollupMerkleTreeWitness.toJSON(witness) + LinkedMerkleTreeWitness.toJSON(witness) ), } satisfies StateTransitionParametersJSON); } @@ -69,7 +69,7 @@ export class StateTransitionParametersSerializer merkleWitnesses: parsed.merkleWitnesses.map( (witness) => - new RollupMerkleTreeWitness(RollupMerkleTreeWitness.fromJSON(witness)) + new LinkedMerkleTreeWitness(LinkedMerkleTreeWitness.fromJSON(witness)) ), }; } 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/src/state/async/AsyncLinkedMerkleTreeStore.ts b/packages/sequencer/src/state/async/AsyncLinkedMerkleTreeStore.ts new file mode 100644 index 000000000..e2e6eb6f4 --- /dev/null +++ b/packages/sequencer/src/state/async/AsyncLinkedMerkleTreeStore.ts @@ -0,0 +1,25 @@ +import { LinkedLeaf } from "@proto-kit/common"; + +import { MerkleTreeNode, MerkleTreeNodeQuery } from "./AsyncMerkleTreeStore"; + +export type StoredLeaf = { leaf: LinkedLeaf; index: bigint }; + +export interface AsyncLinkedMerkleTreeStore { + openTransaction: () => Promise; + + commit: () => Promise; + + writeNodes: (nodes: MerkleTreeNode[]) => void; + + writeLeaves: (leaves: StoredLeaf[]) => void; + + getNodesAsync: ( + nodes: MerkleTreeNodeQuery[] + ) => Promise<(bigint | undefined)[]>; + + getLeavesAsync: (paths: bigint[]) => Promise<(StoredLeaf | undefined)[]>; + + getMaximumIndexAsync: () => Promise; + + getLeafLessOrEqualAsync: (path: bigint) => Promise; +} diff --git a/packages/sequencer/src/state/merkle/CachedLinkedMerkleTreeStore.ts b/packages/sequencer/src/state/merkle/CachedLinkedMerkleTreeStore.ts new file mode 100644 index 000000000..5ddae4fd8 --- /dev/null +++ b/packages/sequencer/src/state/merkle/CachedLinkedMerkleTreeStore.ts @@ -0,0 +1,300 @@ +import { + log, + InMemoryLinkedLeafStore, + LinkedLeaf, + InMemoryMerkleTreeStorage, + PreloadingLinkedMerkleTreeStore, + mapSequential, +} from "@proto-kit/common"; + +import { + MerkleTreeNode, + MerkleTreeNodeQuery, +} from "../async/AsyncMerkleTreeStore"; +import { AsyncLinkedMerkleTreeStore } from "../async/AsyncLinkedMerkleTreeStore"; + +export class CachedLinkedMerkleTreeStore + implements PreloadingLinkedMerkleTreeStore +{ + private writeCache: { + nodes: { + [key: number]: { + [key: string]: bigint; + }; + }; + leaves: { + [key: string]: { leaf: LinkedLeaf; index: bigint }; + }; + } = { nodes: {}, leaves: {} }; + + private readonly leafStore = new InMemoryLinkedLeafStore(); + + private readonly nodeStore = new InMemoryMerkleTreeStorage(); + + private constructor(private readonly parent: AsyncLinkedMerkleTreeStore) {} + + public static async new( + parent: AsyncLinkedMerkleTreeStore + ): Promise { + const cachedInstance = new CachedLinkedMerkleTreeStore(parent); + await cachedInstance.preloadMaximumIndex(); + 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 this.nodeStore.getNode(key, level); + } + + // This gets the nodes from the in memory store. + // If the node is not in the in-memory store it goes to the parent (i.e. + // what's put in the constructor). + public async getNodesAsync( + nodes: MerkleTreeNodeQuery[] + ): Promise<(bigint | undefined)[]> { + const results = Array(nodes.length).fill(undefined); + + const toFetch: MerkleTreeNodeQuery[] = []; + + nodes.forEach((node, index) => { + const localResult = this.getNode(node.key, node.level); + if (localResult !== undefined) { + results[index] = localResult; + } else { + toFetch.push(node); + } + }); + + // Reverse here, so that we can use pop() later + const fetchResult = (await this.parent.getNodesAsync(toFetch)).reverse(); + + results.forEach((result, index) => { + if (result === undefined) { + results[index] = fetchResult.pop(); + } + }); + + return results; + } + + // This sets the nodes in the cache and in the in-memory tree. + public setNode(key: bigint, level: number, value: bigint) { + this.nodeStore.setNode(key, level, value); + (this.writeCache.nodes[level] ??= {})[key.toString()] = value; + } + + // This is basically setNode (cache and in-memory) for a list of nodes. + // Looks only to be used in the mergeIntoParent + public writeNodes(nodes: MerkleTreeNode[]) { + nodes.forEach(({ key, level, value }) => { + this.setNode(key, level, value); + }); + } + + // This gets the leaves and the nodes from the in memory store. + // If the leaf is not in the in-memory store it goes to the parent (i.e. + // what's put in the constructor). + public async getLeavesAsync(paths: bigint[]) { + const results = Array<{ leaf: LinkedLeaf; index: bigint } | undefined>( + paths.length + ).fill(undefined); + + const toFetch: bigint[] = []; + + paths.forEach((path, index) => { + const localResult = this.getLeaf(path); + if (localResult !== undefined) { + results[index] = localResult; + } else { + toFetch.push(path); + } + }); + + // Reverse here, so that we can use pop() later + const fetchResult = (await this.parent.getLeavesAsync(toFetch)).reverse(); + + results.forEach((result, index) => { + if (result === undefined) { + results[index] = fetchResult.pop(); + } + }); + + return results; + } + + // This is just used in the mergeIntoParent. + // It doesn't need any fancy logic and just updates the leaves. + // I don't think we need to coordinate this with the nodes + // or do any calculations. Just a straight copy and paste. + public writeLeaves(leaves: { leaf: LinkedLeaf; index: bigint }[]) { + leaves.forEach(({ leaf, index }) => { + this.setLeaf(index, leaf); + }); + } + + public setLeaf(index: bigint, leaf: LinkedLeaf) { + this.writeCache.leaves[leaf.path.toString()] = { leaf: leaf, index: index }; + this.leafStore.setLeaf(index, leaf); + } + + // This gets the nodes from the cache. + // Only used in mergeIntoParent + public getWrittenNodes(): { + [key: number]: { + [key: string]: bigint; + }; + } { + return this.writeCache.nodes; + } + + // This gets the leaves from the cache. + // Only used in mergeIntoParent + public getWrittenLeaves(): { leaf: LinkedLeaf; index: bigint }[] { + return Object.values(this.writeCache.leaves); + } + + // This resets the cache (not the in memory tree). + public resetWrittenTree() { + this.writeCache = { nodes: {}, leaves: {} }; + } + + // Used only in the preloadKeys + // Basically, gets all of the relevant nodes (and siblings) in the Merkle tree + // at the various levels required to produce a witness for the given index (at level 0). + // But only gets those that aren't already in the cache. + private collectNodesToFetch(index: bigint) { + // This is hardcoded, but should be changed. + const HEIGHT = 40n; + const leafCount = 2n ** (HEIGHT - 1n); + + let currentIndex = index >= leafCount ? index % leafCount : index; + + const nodesToRetrieve: MerkleTreeNodeQuery[] = []; + + for (let level = 0; level < HEIGHT; level++) { + const key = currentIndex; + + const isLeft = key % 2n === 0n; + const siblingKey = isLeft ? key + 1n : key - 1n; + + // Only preload node if it is not already preloaded. + // We also don't want to overwrite because changes will get lost (tracing) + if (this.getNode(key, level) === undefined) { + nodesToRetrieve.push({ + key, + level, + }); + if (level === 0) { + log.trace(`Queued preloading of ${key} @ ${level}`); + } + } + + if (this.getNode(siblingKey, level) === undefined) { + nodesToRetrieve.push({ + key: siblingKey, + level, + }); + } + currentIndex /= 2n; + } + return nodesToRetrieve; + } + + protected async preloadMaximumIndex() { + if (this.leafStore.getMaximumIndex() === undefined) { + this.leafStore.maximumIndex = await this.parent.getMaximumIndexAsync(); + } + } + + public async preloadNodes(indexes: bigint[]) { + const nodesToRetrieve = indexes.flatMap((key) => + this.collectNodesToFetch(key) + ); + + const results = await this.parent.getNodesAsync(nodesToRetrieve); + nodesToRetrieve.forEach(({ key, level }, index) => { + const value = results[index]; + if (value !== undefined) { + this.setNode(key, level, value); + } + }); + } + + public getLeaf(path: bigint) { + return this.leafStore.getLeaf(path); + } + + // Takes a list of paths and for each key collects the relevant nodes from the + // parent tree and sets the leaf and node in the cached tree (and in-memory tree). + public async preloadKey(path: bigint) { + const leaf = + this.leafStore.getLeaf(path) ?? + (await this.parent.getLeavesAsync([path]))[0]; + if (leaf !== undefined) { + this.leafStore.setLeaf(leaf.index, leaf.leaf); + // Update + await this.preloadNodes([leaf.index]); + } else { + // Insert + const previousLeaf = + this.leafStore.getLeafLessOrEqual(path) ?? + (await this.parent.getLeafLessOrEqualAsync(path)); + if (previousLeaf === undefined) { + throw Error("Previous Leaf should never be empty"); + } + this.leafStore.setLeaf(previousLeaf.index, previousLeaf.leaf); + await this.preloadNodes([previousLeaf.index]); + const maximumIndex = + this.leafStore.getMaximumIndex() ?? + (await this.parent.getMaximumIndexAsync()); + if (maximumIndex === undefined) { + throw Error("Maximum index should be defined in parent."); + } + await this.preloadNodes([maximumIndex + 1n]); + } + } + + public async preloadKeys(paths: bigint[]): Promise { + await mapSequential(paths, (x) => this.preloadKey(x)); + } + + // This merges the cache into the parent tree and resets the cache, but not the + // in-memory merkle tree. + public async mergeIntoParent(): Promise { + // In case no state got set we can skip this step + if (Object.keys(this.writeCache.leaves).length === 0) { + return; + } + + await this.parent.openTransaction(); + const nodes = this.getWrittenNodes(); + const leaves = this.getWrittenLeaves(); + + this.parent.writeLeaves(Object.values(leaves)); + const writes = Object.keys(nodes).flatMap((levelString) => { + const level = Number(levelString); + return Object.entries(nodes[level]).map( + ([key, value]) => { + return { + key: BigInt(key), + level, + value, + }; + } + ); + }); + + this.parent.writeNodes(writes); + + await this.parent.commit(); + this.resetWrittenTree(); + } + + public getLeafLessOrEqual(path: bigint) { + return this.leafStore.getLeafLessOrEqual(path); + } + + public getMaximumIndex() { + return this.leafStore.getMaximumIndex(); + } +} diff --git a/packages/sequencer/src/state/merkle/SyncCachedLinkedMerkleTreeStore.ts b/packages/sequencer/src/state/merkle/SyncCachedLinkedMerkleTreeStore.ts new file mode 100644 index 000000000..f38d9990b --- /dev/null +++ b/packages/sequencer/src/state/merkle/SyncCachedLinkedMerkleTreeStore.ts @@ -0,0 +1,79 @@ +import { + LinkedLeaf, + LinkedMerkleTree, + InMemoryLinkedLeafStore, + InMemoryMerkleTreeStorage, + PreloadingLinkedMerkleTreeStore, +} from "@proto-kit/common"; + +import { StoredLeaf } from "../async/AsyncLinkedMerkleTreeStore"; + +// 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 + implements PreloadingLinkedMerkleTreeStore +{ + private readonly leafStore = new InMemoryLinkedLeafStore(); + + private readonly nodeStore = new InMemoryMerkleTreeStorage(); + + public constructor( + private readonly parent: PreloadingLinkedMerkleTreeStore + ) {} + + public getNode(key: bigint, level: number): bigint | undefined { + return ( + this.nodeStore.getNode(key, level) ?? this.parent.getNode(key, level) + ); + } + + public setNode(key: bigint, level: number, value: bigint) { + this.nodeStore.setNode(key, level, value); + } + + public getLeaf(path: bigint): StoredLeaf | undefined { + return this.leafStore.getLeaf(path) ?? this.parent.getLeaf(path); + } + + public setLeaf(index: bigint, value: LinkedLeaf) { + this.leafStore.setLeaf(index, value); + } + + // Need to make sure we call the parent as the super will usually be empty + // The tree calls this method. + public getMaximumIndex(): bigint | undefined { + return (this.leafStore.getMaximumIndex() ?? -1) > + (this.parent.getMaximumIndex() ?? -1) + ? this.leafStore.getMaximumIndex() + : this.parent.getMaximumIndex(); + } + + public getLeafLessOrEqual(path: bigint): StoredLeaf | undefined { + return ( + this.leafStore.getLeafLessOrEqual(path) ?? + this.parent.getLeafLessOrEqual(path) + ); + } + + public async preloadKeys(path: bigint[]) { + await this.parent.preloadKeys(path); + } + + public mergeIntoParent() { + if (Object.keys(this.leafStore.leaves).length === 0) { + return; + } + + Object.values(this.leafStore.leaves).forEach(({ leaf, index }) => + this.parent.setLeaf(index, leaf) + ); + Array.from({ length: LinkedMerkleTree.HEIGHT }).forEach((ignored, level) => + Object.entries(this.nodeStore.nodes[level]).forEach((entry) => { + this.parent.setNode(BigInt(entry[0]), level, entry[1]); + }) + ); + + this.leafStore.leaves = {}; + this.nodeStore.nodes = {}; + } +} diff --git a/packages/sequencer/src/storage/StorageDependencyFactory.ts b/packages/sequencer/src/storage/StorageDependencyFactory.ts index ed660d8fa..27c1a29b2 100644 --- a/packages/sequencer/src/storage/StorageDependencyFactory.ts +++ b/packages/sequencer/src/storage/StorageDependencyFactory.ts @@ -5,6 +5,7 @@ import { } from "@proto-kit/common"; import { AsyncStateService } from "../state/async/AsyncStateService"; +import { AsyncLinkedMerkleTreeStore } from "../state/async/AsyncLinkedMerkleTreeStore"; import { AsyncMerkleTreeStore } from "../state/async/AsyncMerkleTreeStore"; import { BatchStorage } from "./repositories/BatchStorage"; @@ -15,12 +16,12 @@ import { TransactionStorage } from "./repositories/TransactionStorage"; export interface StorageDependencyMinimumDependencies extends DependencyRecord { asyncStateService: DependencyDeclaration; - asyncMerkleStore: DependencyDeclaration; + asyncMerkleStore: DependencyDeclaration; batchStorage: DependencyDeclaration; blockQueue: DependencyDeclaration; blockStorage: DependencyDeclaration; unprovenStateService: DependencyDeclaration; - unprovenMerkleStore: DependencyDeclaration; + unprovenMerkleStore: DependencyDeclaration; blockTreeStore: DependencyDeclaration; messageStorage: DependencyDeclaration; settlementStorage: DependencyDeclaration; diff --git a/packages/sequencer/src/storage/inmemory/InMemoryAsyncLinkedMerkleTreeStore.ts b/packages/sequencer/src/storage/inmemory/InMemoryAsyncLinkedMerkleTreeStore.ts new file mode 100644 index 000000000..877efedbd --- /dev/null +++ b/packages/sequencer/src/storage/inmemory/InMemoryAsyncLinkedMerkleTreeStore.ts @@ -0,0 +1,90 @@ +import { + InMemoryLinkedLeafStore, + InMemoryMerkleTreeStorage, + LinkedLeaf, + LinkedMerkleTreeStore, + noop, +} from "@proto-kit/common"; + +import { AsyncLinkedMerkleTreeStore } from "../../state/async/AsyncLinkedMerkleTreeStore"; +import { + MerkleTreeNode, + MerkleTreeNodeQuery, +} from "../../state/async/AsyncMerkleTreeStore"; + +export class InMemoryAsyncLinkedMerkleTreeStore + implements AsyncLinkedMerkleTreeStore, LinkedMerkleTreeStore +{ + private readonly leafStore = new InMemoryLinkedLeafStore(); + + private readonly nodeStore = new InMemoryMerkleTreeStorage(); + + public async openTransaction(): Promise { + noop(); + } + + public async commit(): Promise { + noop(); + } + + public writeNodes(nodes: MerkleTreeNode[]): void { + nodes.forEach(({ key, level, value }) => + this.nodeStore.setNode(key, level, value) + ); + } + + // This is using the index/key + public writeLeaves(leaves: { leaf: LinkedLeaf; index: bigint }[]) { + leaves.forEach(({ leaf, index }) => { + this.leafStore.setLeaf(index, leaf); + }); + } + + public async getNodesAsync( + nodes: MerkleTreeNodeQuery[] + ): Promise<(bigint | undefined)[]> { + return nodes.map(({ key, level }) => this.nodeStore.getNode(key, level)); + } + + public async getLeavesAsync(paths: bigint[]) { + return paths.map((path) => { + const leaf = this.leafStore.getLeaf(path); + if (leaf !== undefined) { + return leaf; + } + return undefined; + }); + } + + public async getMaximumIndexAsync() { + return this.leafStore.getMaximumIndex(); + } + + public async getLeafLessOrEqualAsync(path: bigint) { + return this.leafStore.getLeafLessOrEqual(path); + } + + public setLeaf(index: bigint, value: LinkedLeaf) { + this.leafStore.setLeaf(index, value); + } + + public getLeaf(path: bigint) { + return this.leafStore.getLeaf(path); + } + + public getLeafLessOrEqual(path: bigint) { + return this.leafStore.getLeafLessOrEqual(path); + } + + public getMaximumIndex() { + return this.leafStore.getMaximumIndex(); + } + + public setNode(key: bigint, level: number, value: bigint) { + this.nodeStore.setNode(key, level, value); + } + + public getNode(key: bigint, level: number) { + return this.nodeStore.getNode(key, level); + } +} diff --git a/packages/sequencer/src/storage/inmemory/InMemoryDatabase.ts b/packages/sequencer/src/storage/inmemory/InMemoryDatabase.ts index 963909589..be06030ec 100644 --- a/packages/sequencer/src/storage/inmemory/InMemoryDatabase.ts +++ b/packages/sequencer/src/storage/inmemory/InMemoryDatabase.ts @@ -9,7 +9,7 @@ import { StorageDependencyMinimumDependencies } from "../StorageDependencyFactor import { Database } from "../Database"; import { InMemoryBlockStorage } from "./InMemoryBlockStorage"; -import { InMemoryAsyncMerkleTreeStore } from "./InMemoryAsyncMerkleTreeStore"; +import { InMemoryAsyncLinkedMerkleTreeStore } from "./InMemoryAsyncLinkedMerkleTreeStore"; import { InMemoryBatchStorage } from "./InMemoryBatchStorage"; import { InMemoryMessageStorage } from "./InMemoryMessageStorage"; import { InMemorySettlementStorage } from "./InMemorySettlementStorage"; @@ -20,7 +20,7 @@ export class InMemoryDatabase extends SequencerModule implements Database { public dependencies(): StorageDependencyMinimumDependencies { return { asyncMerkleStore: { - useClass: InMemoryAsyncMerkleTreeStore, + useClass: InMemoryAsyncLinkedMerkleTreeStore, }, asyncStateService: { useFactory: () => new CachedStateService(undefined), @@ -38,10 +38,10 @@ export class InMemoryDatabase extends SequencerModule implements Database { useFactory: () => new CachedStateService(undefined), }, unprovenMerkleStore: { - useClass: InMemoryAsyncMerkleTreeStore, + useClass: InMemoryAsyncLinkedMerkleTreeStore, }, blockTreeStore: { - useClass: InMemoryAsyncMerkleTreeStore, + useClass: InMemoryAsyncLinkedMerkleTreeStore, }, messageStorage: { useClass: InMemoryMessageStorage, diff --git a/packages/sequencer/test/merkle/CachedLinkedMerkleStore.test.ts b/packages/sequencer/test/merkle/CachedLinkedMerkleStore.test.ts new file mode 100644 index 000000000..4345ef761 --- /dev/null +++ b/packages/sequencer/test/merkle/CachedLinkedMerkleStore.test.ts @@ -0,0 +1,542 @@ +import { + expectDefined, + LinkedLeafStruct, + LinkedMerkleTree, +} from "@proto-kit/common"; +import { beforeEach, expect } from "@jest/globals"; +import { Field, Poseidon } from "o1js"; + +import { CachedLinkedMerkleTreeStore } from "../../src/state/merkle/CachedLinkedMerkleTreeStore"; +import { InMemoryAsyncLinkedMerkleTreeStore } from "../../src/storage/inmemory/InMemoryAsyncLinkedMerkleTreeStore"; +import { SyncCachedLinkedMerkleTreeStore } from "../../src/state/merkle/SyncCachedLinkedMerkleTreeStore"; + +describe("cached linked merkle store", () => { + let mainStore: InMemoryAsyncLinkedMerkleTreeStore; + + let cache1: CachedLinkedMerkleTreeStore; + let tree1: LinkedMerkleTree; + + beforeEach(async () => { + mainStore = new InMemoryAsyncLinkedMerkleTreeStore(); + + const cachedStore = await CachedLinkedMerkleTreeStore.new(mainStore); + + const tmpTree = new LinkedMerkleTree(cachedStore); + tmpTree.setLeaf(5n, 10n); + await cachedStore.mergeIntoParent(); + + cache1 = await CachedLinkedMerkleTreeStore.new(mainStore); + tree1 = new LinkedMerkleTree(cache1); + }); + + it("should cache multiple keys correctly", async () => { + expect.assertions(11); + await cache1.preloadKeys([16n, 46n]); + tree1.setLeaf(16n, 16n); + tree1.setLeaf(46n, 46n); + + const cache2 = new SyncCachedLinkedMerkleTreeStore(cache1); + const tree2 = new LinkedMerkleTree(cache2); + + const leaf1 = tree1.getLeaf(16n); + const leaf2 = tree1.getLeaf(46n); + + expectDefined(leaf1); + expectDefined(leaf2); + + const storedLeaf1 = cache2.getLeaf(16n); + const storedLeaf2 = cache2.getLeaf(46n); + + expectDefined(storedLeaf1); + expectDefined(storedLeaf2); + + expect(storedLeaf1.index).toStrictEqual(2n); + expect(storedLeaf2.index).toStrictEqual(3n); + + expect(tree2.getNode(0, storedLeaf1.index).toBigInt()).toBe( + leaf1.hash().toBigInt() + ); + expect(tree2.getNode(0, storedLeaf2.index).toBigInt()).toBe( + leaf2.hash().toBigInt() + ); + + expect(tree2.getLeaf(16n)).toEqual(leaf1); + expect(tree2.getLeaf(46n)).toEqual(leaf2); + + expect(tree2.getRoot().toString()).toStrictEqual( + tree1.getRoot().toString() + ); + }); + + it("simple test - check hash of updated node is updated", async () => { + // main store already has 0n and 5n paths defined. + // preloading 10n should load up 5n in the cache1 leaf and node stores. + await cache1.preloadKeys([10n]); + + expectDefined(cache1.getLeaf(5n)); + expectDefined(cache1.getNode(1n, 0)); + + tree1.setLeaf(10n, 10n); + await cache1.mergeIntoParent(); + + const leaf5 = tree1.getLeaf(5n); + const leaf10 = tree1.getLeaf(10n); + expectDefined(leaf5); + expectDefined(leaf10); + + const storedLeaf5 = cache1.getLeaf(5n); + const storedLeaf10 = cache1.getLeaf(10n); + + expectDefined(storedLeaf5); + expectDefined(storedLeaf10); + + expect(storedLeaf5).toStrictEqual({ + leaf: { value: 10n, path: 5n, nextPath: 10n }, + index: 1n, + }); + expect(storedLeaf10.index).toStrictEqual(2n); + + // Check leaves were hashed properly when added to nodes/merkle-tree + expect(cache1.getNode(storedLeaf10.index, 0)).toStrictEqual( + leaf10.hash().toBigInt() + ); + expect(cache1.getNode(storedLeaf5.index, 0)).toStrictEqual( + leaf5.hash().toBigInt() + ); + }); + + it("should preload through multiple levels and insert correctly at right index", async () => { + await cache1.preloadKeys([10n, 11n, 12n, 13n]); + + tree1.setLeaf(10n, 10n); + tree1.setLeaf(11n, 11n); + tree1.setLeaf(12n, 12n); + tree1.setLeaf(13n, 13n); + await cache1.mergeIntoParent(); + + const cache2 = new SyncCachedLinkedMerkleTreeStore(cache1); + await cache2.preloadKeys([14n]); + + const tree2 = new LinkedMerkleTree(cache2); + tree2.setLeaf(14n, 14n); + + const leaf = tree1.getLeaf(5n); + const leaf2 = tree2.getLeaf(14n); + + expectDefined(leaf); + expectDefined(leaf2); + + const storedLeaf5 = cache2.getLeaf(5n); + const storedLeaf10 = cache2.getLeaf(10n); + const storedLeaf11 = cache2.getLeaf(11n); + const storedLeaf12 = cache2.getLeaf(12n); + const storedLeaf13 = cache2.getLeaf(13n); + const storedLeaf14 = cache2.getLeaf(14n); + + expectDefined(storedLeaf5); + expectDefined(storedLeaf10); + expectDefined(storedLeaf11); + expectDefined(storedLeaf12); + expectDefined(storedLeaf13); + expectDefined(storedLeaf14); + + expect(storedLeaf5.index).toStrictEqual(1n); + expect(storedLeaf10.index).toStrictEqual(2n); + expect(storedLeaf11.index).toStrictEqual(3n); + expect(storedLeaf12.index).toStrictEqual(4n); + expect(storedLeaf13.index).toStrictEqual(5n); + expect(storedLeaf14.index).toStrictEqual(6n); + + // Check leaves were hashed properly when added to nodes/merkle-tree + expect(cache1.getNode(storedLeaf5.index, 0)).toStrictEqual( + leaf.hash().toBigInt() + ); + expect(cache2.getNode(storedLeaf14.index, 0)).toStrictEqual( + leaf2.hash().toBigInt() + ); + }); + + it("should preload through multiple levels and insert correctly at right index - harder", async () => { + await cache1.preloadKeys([10n, 100n, 200n, 300n, 400n, 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); + + const cache2 = new SyncCachedLinkedMerkleTreeStore(cache1); + await cache2.preloadKeys([14n]); + const tree2 = new LinkedMerkleTree(cache2); + tree2.setLeaf(14n, 14n); + + const leaf = tree1.getLeaf(5n); + const leaf2 = tree2.getLeaf(14n); + + expectDefined(leaf); + expectDefined(leaf2); + + const storedLeaf5 = cache2.getLeaf(5n); + const storedLeaf10 = cache2.getLeaf(10n); + const storedLeaf100 = cache2.getLeaf(100n); + const storedLeaf200 = cache2.getLeaf(200n); + const storedLeaf300 = cache2.getLeaf(300n); + const storedLeaf400 = cache2.getLeaf(400n); + const storedLeaf500 = cache2.getLeaf(500n); + const storedLeaf14 = cache2.getLeaf(14n); + + expectDefined(storedLeaf5); + expectDefined(storedLeaf10); + expectDefined(storedLeaf100); + expectDefined(storedLeaf200); + expectDefined(storedLeaf300); + expectDefined(storedLeaf400); + expectDefined(storedLeaf500); + expectDefined(storedLeaf14); + + expect(storedLeaf5.index).toStrictEqual(1n); + expect(storedLeaf10.index).toStrictEqual(2n); + expect(storedLeaf100.index).toStrictEqual(3n); + expect(storedLeaf200?.index).toStrictEqual(4n); + expect(storedLeaf300?.index).toStrictEqual(5n); + expect(storedLeaf400.index).toStrictEqual(6n); + expect(storedLeaf500.index).toStrictEqual(7n); + expect(storedLeaf14.index).toStrictEqual(8n); + + expect(cache1.getNode(storedLeaf5.index, 0)).toStrictEqual( + leaf.hash().toBigInt() + ); + expect(cache2.getNode(storedLeaf14.index, 0)).toStrictEqual( + leaf2.hash().toBigInt() + ); + expect(tree1.getRoot()).not.toEqual(tree2.getRoot()); + await cache2.mergeIntoParent(); + expect(tree1.getRoot()).toEqual(tree2.getRoot()); + }); + + it("mimic transaction execution service", async () => { + expect.assertions(18); + + const treeCache1 = new LinkedMerkleTree(cache1); + await cache1.preloadKeys([10n, 20n]); + treeCache1.setLeaf(10n, 10n); + treeCache1.setLeaf(20n, 20n); + await cache1.mergeIntoParent(); + + const cache2 = new SyncCachedLinkedMerkleTreeStore(cache1); + const treeCache2 = new LinkedMerkleTree(cache2); + await cache2.preloadKeys([7n]); + treeCache2.setLeaf(7n, 7n); + cache2.mergeIntoParent(); + + const leaves = await cache1.getLeavesAsync([0n, 5n, 7n, 10n, 20n]); + expectDefined(leaves[0]); + expectDefined(leaves[1]); + expectDefined(leaves[2]); + expectDefined(leaves[3]); + expectDefined(leaves[4]); + + expect(leaves[0]?.leaf).toEqual({ + value: 0n, + path: 0n, + nextPath: 5n, + }); + expect(leaves[1]?.leaf).toEqual({ + value: 10n, + path: 5n, + nextPath: 7n, + }); + expect(leaves[2]?.leaf).toEqual({ + value: 7n, + path: 7n, + nextPath: 10n, + }); + expect(leaves[3]?.leaf).toEqual({ + value: 10n, + path: 10n, + nextPath: 20n, + }); + expect(leaves[4]?.leaf).toEqual({ + value: 20n, + path: 20n, + nextPath: Field.ORDER - 1n, + }); + + const storedLeaf5 = cache1.getLeaf(5n); + const storedLeaf7 = cache1.getLeaf(7n); + const storedLeaf10 = cache1.getLeaf(10n); + const storedLeaf20 = cache1.getLeaf(20n); + + expectDefined(storedLeaf5); + await expect( + cache1.getNodesAsync([{ key: storedLeaf5.index, level: 0 }]) + ).resolves.toStrictEqual([ + Poseidon.hash([Field(10), Field(5), Field(7)]).toBigInt(), + ]); + expectDefined(storedLeaf7); + await expect( + cache1.getNodesAsync([{ key: storedLeaf7.index, level: 0 }]) + ).resolves.toStrictEqual([ + Poseidon.hash([Field(7), Field(7), Field(10)]).toBigInt(), + ]); + expectDefined(storedLeaf10); + await expect( + cache1.getNodesAsync([{ key: storedLeaf10.index, level: 0 }]) + ).resolves.toStrictEqual([ + Poseidon.hash([Field(10), Field(10), Field(20)]).toBigInt(), + ]); + expectDefined(storedLeaf20); + await expect( + cache1.getNodesAsync([{ key: storedLeaf20.index, level: 0 }]) + ).resolves.toStrictEqual([ + Poseidon.hash([Field(20), Field(20), Field(Field.ORDER - 1n)]).toBigInt(), + ]); + }); + + it("should cache correctly", async () => { + expect.assertions(15); + + const cache2 = new SyncCachedLinkedMerkleTreeStore(cache1); + const tree2 = new LinkedMerkleTree(cache2); + + await cache2.preloadKeys([5n]); + const leaf1 = tree2.getLeaf(5n); + const storedLeaf1 = cache2.getLeaf(5n); + expectDefined(leaf1); + expectDefined(storedLeaf1); + await expect( + mainStore.getNodesAsync([{ key: storedLeaf1.index, level: 0 }]) + ).resolves.toStrictEqual([ + Poseidon.hash([leaf1.value, leaf1.path, leaf1.nextPath]).toBigInt(), + ]); + + tree1.setLeaf(10n, 20n); + + const leaf2 = tree2.getLeaf(10n); + const storedLeaf2 = cache2.getLeaf(10n); + expectDefined(leaf2); + expectDefined(storedLeaf2); + expect(tree2.getNode(0, storedLeaf2.index).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.merkleWitness + .calculateRoot( + Poseidon.hash([ + witness.leaf.value, + witness.leaf.path, + witness.leaf.nextPath, + ]) + ) + .toString() + ).toBe(tree1.getRoot().toString()); + + expect( + witness.merkleWitness + .calculateRoot(Poseidon.hash([Field(11), Field(5n), Field(10n)])) + .toString() + ).not.toBe(tree1.getRoot().toString()); + + const witness2 = tree1.getWitness(10n); + + expect( + witness2.merkleWitness + .calculateRoot( + Poseidon.hash([ + Field(20), + Field(10n), + witness2.leaf.nextPath, // This is the maximum as the the leaf 10n should be the last + ]) + ) + .toString() + ).toBe(tree2.getRoot().toString()); + + 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 storedLeaf15 = cache2.getLeaf(15n); + const leaf15 = tree2.getLeaf(15n); + expectDefined(leaf15); + expectDefined(storedLeaf15); + expect(tree1.getRoot().toString()).toBe(tree2.getRoot().toString()); + expect(tree1.getNode(0, storedLeaf15.index).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); + await cachedStore.preloadKey(15n); + + expect(new LinkedMerkleTree(cachedStore).getRoot().toString()).toBe( + tree2.getRoot().toString() + ); + }); + + it("mimic transaction execution service further", async () => { + expect.assertions(16); + + const mStore = new InMemoryAsyncLinkedMerkleTreeStore(); + const mCache = await CachedLinkedMerkleTreeStore.new(mStore); + const mCache2 = new SyncCachedLinkedMerkleTreeStore(mCache); + const treeCache1 = new LinkedMerkleTree(mCache); + const treeCache2 = new LinkedMerkleTree(mCache2); + + await mCache.preloadKeys([5n]); + treeCache1.setLeaf(10n, 10n); + treeCache1.setLeaf(20n, 20n); + + await mCache2.preloadKeys([7n]); + treeCache2.setLeaf(7n, 7n); + mCache2.mergeIntoParent(); + + const leaves = await mCache.getLeavesAsync([0n, 7n, 10n, 20n]); + expectDefined(leaves[0]); + expectDefined(leaves[1]); + expectDefined(leaves[2]); + expectDefined(leaves[3]); + + expect(leaves[0]?.leaf).toEqual({ + value: 0n, + path: 0n, + nextPath: 7n, + }); + expect(leaves[1]?.leaf).toEqual({ + value: 7n, + path: 7n, + nextPath: 10n, + }); + expect(leaves[2]?.leaf).toEqual({ + value: 10n, + path: 10n, + nextPath: 20n, + }); + expect(leaves[3]?.leaf).toEqual({ + value: 20n, + path: 20n, + nextPath: Field.ORDER - 1n, + }); + + const storedLeaf0 = mCache.getLeaf(0n); + const storedLeaf7 = mCache.getLeaf(7n); + const storedLeaf10 = mCache.getLeaf(10n); + const storedLeaf20 = mCache.getLeaf(20n); + + expectDefined(storedLeaf0); + await expect( + mCache.getNodesAsync([{ key: storedLeaf0.index, level: 0 }]) + ).resolves.toStrictEqual([ + Poseidon.hash([Field(0), Field(0), Field(7)]).toBigInt(), + ]); + expectDefined(storedLeaf7); + await expect( + mCache.getNodesAsync([{ key: storedLeaf7.index, level: 0 }]) + ).resolves.toStrictEqual([ + Poseidon.hash([Field(7), Field(7), Field(10)]).toBigInt(), + ]); + expectDefined(storedLeaf10); + await expect( + mCache.getNodesAsync([{ key: storedLeaf10.index, level: 0 }]) + ).resolves.toStrictEqual([ + Poseidon.hash([Field(10), Field(10), Field(20)]).toBigInt(), + ]); + expectDefined(storedLeaf20); + await expect( + mCache.getNodesAsync([{ key: storedLeaf20.index, level: 0 }]) + ).resolves.toStrictEqual([ + Poseidon.hash([Field(20), Field(20), Field(Field.ORDER - 1n)]).toBigInt(), + ]); + }); + + it("mimic block production test && ST Prover", async () => { + // main store already has 0n and 5n paths defined. + // preloading 10n should load up 5n in the cache1 leaf and node stores. + await cache1.preloadKeys([10n, 10n, 10n]); + const state = tree1.getRoot(); + + // This is an insert as 10n is not already in the tree. + const witness1 = tree1.setLeaf(10n, 10n); + // This checks the right previous leaf was found. + expect( + witness1.leafPrevious.merkleWitness + .checkMembershipSimple( + state, + new LinkedLeafStruct({ + value: Field(witness1.leafPrevious.leaf.value), + path: Field(witness1.leafPrevious.leaf.path), + nextPath: Field(witness1.leafPrevious.leaf.nextPath), + }).hash() + ) + .toBoolean() + ).toStrictEqual(true); + + // We now look to the state after the prevLeaf is changed. + // The prev leaf should be the 5n. + const rootAfterFirstChange = + witness1.leafPrevious.merkleWitness.calculateRoot( + new LinkedLeafStruct({ + value: Field(witness1.leafPrevious.leaf.value), + path: Field(witness1.leafPrevious.leaf.path), + nextPath: Field(10n), + }).hash() + ); + expect( + witness1.leafCurrent.merkleWitness.calculateRoot(Field(0)).toBigInt() + ).toStrictEqual(rootAfterFirstChange.toBigInt()); + + // We now check that right hashing was done to get to the current root. + expect( + witness1.leafCurrent.merkleWitness + .checkMembershipSimple( + tree1.getRoot(), + new LinkedLeafStruct({ + value: Field(10n), + path: Field(10n), + nextPath: Field(witness1.leafPrevious.leaf.nextPath), + }).hash() + ) + .toBoolean() + ).toStrictEqual(true); + + // Now we update the node at 10n, + const witness2 = tree1.setLeaf(10n, 8n); + // We now check that right hashing was done to get to the current root. + expect( + witness2.leafCurrent.merkleWitness.calculateRoot( + new LinkedLeafStruct({ + value: Field(8n), + path: Field(10n), + nextPath: Field(witness2.leafCurrent.leaf.nextPath), + }).hash() + ) + ).toStrictEqual(tree1.getRoot()); + + // Now we update the node at 10n, again, + const witness3 = tree1.setLeaf(10n, 4n); + // We now check that right hashing was done to get to the current root. + expect( + witness3.leafCurrent.merkleWitness.calculateRoot( + new LinkedLeafStruct({ + value: Field(4n), + path: Field(10n), + nextPath: Field(witness2.leafCurrent.leaf.nextPath), + }).hash() + ) + ).toStrictEqual(tree1.getRoot()); + }); +}); diff --git a/packages/sequencer/test/merkle/CachedMerkleStore.test.ts b/packages/sequencer/test/merkle/CachedMerkleStore.test.ts index 4d882d7db..1b6cbf378 100644 --- a/packages/sequencer/test/merkle/CachedMerkleStore.test.ts +++ b/packages/sequencer/test/merkle/CachedMerkleStore.test.ts @@ -25,7 +25,7 @@ describe("cached merkle store", () => { tree1 = new RollupMerkleTree(cache1); }); - it("should cache multiple keys corretly", async () => { + it("should cache multiple keys correctly", async () => { expect.assertions(3); const cache2 = new CachedMerkleTreeStore(cache1); 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(); diff --git a/packages/stack/src/scripts/graphql/server.ts b/packages/stack/src/scripts/graphql/server.ts index 6f6cc3d86..b70ae1255 100644 --- a/packages/stack/src/scripts/graphql/server.ts +++ b/packages/stack/src/scripts/graphql/server.ts @@ -40,7 +40,7 @@ import { GraphqlSequencerModule, GraphqlServer, MempoolResolver, - MerkleWitnessResolver, + LinkedMerkleWitnessResolver as MerkleWitnessResolver, NodeStatusResolver, QueryGraphqlModule, BlockResolver, diff --git a/packages/stack/test/graphql/graphql.test.ts b/packages/stack/test/graphql/graphql.test.ts index fa1c18603..e15a8886f 100644 --- a/packages/stack/test/graphql/graphql.test.ts +++ b/packages/stack/test/graphql/graphql.test.ts @@ -170,8 +170,8 @@ describe("graphql client test", () => { expect(witness).toBeDefined(); // Check if this works, i.e. if it correctly parsed - expect(witness!.calculateRoot(Field(0)).toBigInt()).toBeGreaterThanOrEqual( - 0n - ); + expect( + witness!.merkleWitness.calculateRoot(Field(0)).toBigInt() + ).toBeGreaterThanOrEqual(0n); }); });