diff --git a/yarn-project/merkle-tree/src/index.ts b/yarn-project/merkle-tree/src/index.ts index de19e295bfd..5181cecfc15 100644 --- a/yarn-project/merkle-tree/src/index.ts +++ b/yarn-project/merkle-tree/src/index.ts @@ -4,8 +4,12 @@ export * from './interfaces/merkle_tree.js'; export * from './interfaces/update_only_tree.js'; export * from './pedersen.js'; export * from './sparse_tree/sparse_tree.js'; -export * from './standard_indexed_tree/standard_indexed_tree.js'; +export { LowLeafWitnessData, StandardIndexedTree } from './standard_indexed_tree/standard_indexed_tree.js'; export * from './standard_tree/standard_tree.js'; export { INITIAL_LEAF } from './tree_base.js'; export { newTree } from './new_tree.js'; export { loadTree } from './load_tree.js'; +export * from './snapshots/snapshot_builder.js'; +export * from './snapshots/full_snapshot.js'; +export * from './snapshots/append_only_snapshot.js'; +export * from './snapshots/indexed_tree_snapshot.js'; diff --git a/yarn-project/merkle-tree/src/interfaces/append_only_tree.ts b/yarn-project/merkle-tree/src/interfaces/append_only_tree.ts index f4ecaabe157..77dd7ae9e5d 100644 --- a/yarn-project/merkle-tree/src/interfaces/append_only_tree.ts +++ b/yarn-project/merkle-tree/src/interfaces/append_only_tree.ts @@ -1,9 +1,10 @@ +import { TreeSnapshotBuilder } from '../snapshots/snapshot_builder.js'; import { MerkleTree } from './merkle_tree.js'; /** * A Merkle tree that supports only appending leaves and not updating existing leaves. */ -export interface AppendOnlyTree extends MerkleTree { +export interface AppendOnlyTree extends MerkleTree, TreeSnapshotBuilder { /** * Appends a set of leaf values to the tree. * @param leaves - The set of leaves to be appended. diff --git a/yarn-project/merkle-tree/src/interfaces/update_only_tree.ts b/yarn-project/merkle-tree/src/interfaces/update_only_tree.ts index 59a82d0b118..6bd5c024d0c 100644 --- a/yarn-project/merkle-tree/src/interfaces/update_only_tree.ts +++ b/yarn-project/merkle-tree/src/interfaces/update_only_tree.ts @@ -1,11 +1,12 @@ import { LeafData } from '@aztec/types'; +import { TreeSnapshotBuilder } from '../snapshots/snapshot_builder.js'; import { MerkleTree } from './merkle_tree.js'; /** * A Merkle tree that supports updates at arbitrary indices but not appending. */ -export interface UpdateOnlyTree extends MerkleTree { +export interface UpdateOnlyTree extends MerkleTree, TreeSnapshotBuilder { /** * Updates a leaf at a given index in the tree. * @param leaf - The leaf value to be updated. diff --git a/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.test.ts b/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.test.ts new file mode 100644 index 00000000000..b66eb2af22b --- /dev/null +++ b/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.test.ts @@ -0,0 +1,28 @@ +import levelup, { LevelUp } from 'levelup'; + +import { Pedersen, StandardTree, newTree } from '../index.js'; +import { createMemDown } from '../test/utils/create_mem_down.js'; +import { AppendOnlySnapshotBuilder } from './append_only_snapshot.js'; +import { describeSnapshotBuilderTestSuite } from './snapshot_builder_test_suite.js'; + +describe('AppendOnlySnapshot', () => { + let tree: StandardTree; + let snapshotBuilder: AppendOnlySnapshotBuilder; + let db: LevelUp; + + beforeEach(async () => { + db = levelup(createMemDown()); + const hasher = new Pedersen(); + tree = await newTree(StandardTree, db, hasher, 'test', 4); + snapshotBuilder = new AppendOnlySnapshotBuilder(db, tree, hasher); + }); + + describeSnapshotBuilderTestSuite( + () => tree, + () => snapshotBuilder, + async tree => { + const newLeaves = Array.from({ length: 2 }).map(() => Buffer.from(Math.random().toString())); + await tree.appendLeaves(newLeaves); + }, + ); +}); diff --git a/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts b/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts new file mode 100644 index 00000000000..b530e981b27 --- /dev/null +++ b/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts @@ -0,0 +1,232 @@ +import { Hasher, SiblingPath } from '@aztec/types'; + +import { LevelUp } from 'levelup'; + +import { AppendOnlyTree } from '../interfaces/append_only_tree.js'; +import { TreeBase } from '../tree_base.js'; +import { TreeSnapshot, TreeSnapshotBuilder } from './snapshot_builder.js'; + +// stores the last block that modified this node +const nodeModifiedAtBlockKey = (treeName: string, level: number, index: bigint) => + `snapshot:node:${treeName}:${level}:${index}:block`; + +// stores the value of the node at the above block +const historicalNodeKey = (treeName: string, level: number, index: bigint) => + `snapshot:node:${treeName}:${level}:${index}:value`; + +// metadata for a snapshot +const snapshotRootKey = (treeName: string, block: number) => `snapshot:root:${treeName}:${block}`; +const snapshotNumLeavesKey = (treeName: string, block: number) => `snapshot:numLeaves:${treeName}:${block}`; + +/** + * A more space-efficient way of storing snapshots of AppendOnlyTrees that trades space need for slower + * sibling path reads. + * + * Complexity: + * + * N - count of non-zero nodes in tree + * M - count of snapshots + * H - tree height + * + * Space complexity: O(N + M) (N nodes - stores the last snapshot for each node and M - ints, for each snapshot stores up to which leaf its written to) + * Sibling path access: + * Best case: O(H) database reads + O(1) hashes + * Worst case: O(H) database reads + O(H) hashes + */ +export class AppendOnlySnapshotBuilder implements TreeSnapshotBuilder { + constructor(private db: LevelUp, private tree: TreeBase & AppendOnlyTree, private hasher: Hasher) {} + async getSnapshot(block: number): Promise { + const meta = await this.#getSnapshotMeta(block); + + if (typeof meta === 'undefined') { + throw new Error(`Snapshot for tree ${this.tree.getName()} at block ${block} does not exist`); + } + + return new AppendOnlySnapshot(this.db, block, meta.numLeaves, meta.root, this.tree, this.hasher); + } + + async snapshot(block: number): Promise { + const meta = await this.#getSnapshotMeta(block); + if (typeof meta !== 'undefined') { + // no-op, we already have a snapshot + return new AppendOnlySnapshot(this.db, block, meta.numLeaves, meta.root, this.tree, this.hasher); + } + + const batch = this.db.batch(); + const root = this.tree.getRoot(false); + const depth = this.tree.getDepth(); + const treeName = this.tree.getName(); + const queue: [Buffer, number, bigint][] = [[root, 0, 0n]]; + + // walk the tree in BF and store latest nodes + while (queue.length > 0) { + const [node, level, index] = queue.shift()!; + + const historicalValue = await this.db.get(historicalNodeKey(treeName, level, index)).catch(() => undefined); + if (!historicalValue || !node.equals(historicalValue)) { + // we've never seen this node before or it's different than before + // update the historical tree and tag it with the block that modified it + batch.put(nodeModifiedAtBlockKey(treeName, level, index), String(block)); + batch.put(historicalNodeKey(treeName, level, index), node); + } else { + // if this node hasn't changed, that means, nothing below it has changed either + continue; + } + + if (level + 1 > depth) { + // short circuit if we've reached the leaf level + // otherwise getNode might throw if we ask for the children of a leaf + continue; + } + + // these could be undefined because zero hashes aren't stored in the tree + const [lhs, rhs] = await Promise.all([ + this.tree.getNode(level + 1, 2n * index), + this.tree.getNode(level + 1, 2n * index + 1n), + ]); + + if (lhs) { + queue.push([lhs, level + 1, 2n * index]); + } + + if (rhs) { + queue.push([rhs, level + 1, 2n * index + 1n]); + } + } + + const numLeaves = this.tree.getNumLeaves(false); + batch.put(snapshotNumLeavesKey(treeName, block), String(numLeaves)); + batch.put(snapshotRootKey(treeName, block), root); + await batch.write(); + + return new AppendOnlySnapshot(this.db, block, numLeaves, root, this.tree, this.hasher); + } + + async #getSnapshotMeta(block: number): Promise< + | { + /** The root of the tree snapshot */ + root: Buffer; + /** The number of leaves in the tree snapshot */ + numLeaves: bigint; + } + | undefined + > { + try { + const treeName = this.tree.getName(); + const root = await this.db.get(snapshotRootKey(treeName, block)); + const numLeaves = BigInt(await this.db.get(snapshotNumLeavesKey(treeName, block))); + return { root, numLeaves }; + } catch (err) { + return undefined; + } + } +} + +/** + * a + */ +class AppendOnlySnapshot implements TreeSnapshot { + constructor( + private db: LevelUp, + private block: number, + private leafCount: bigint, + private historicalRoot: Buffer, + private tree: TreeBase & AppendOnlyTree, + private hasher: Hasher, + ) {} + + public async getSiblingPath(index: bigint): Promise> { + const path: Buffer[] = []; + const depth = this.tree.getDepth(); + let level = depth; + + while (level > 0) { + const isRight = index & 0x01n; + const siblingIndex = isRight ? index - 1n : index + 1n; + + const sibling = await this.#getHistoricalNodeValue(level, siblingIndex); + path.push(sibling); + + level -= 1; + index >>= 1n; + } + + return new SiblingPath(depth as N, path); + } + + getDepth(): number { + return this.tree.getDepth(); + } + + getNumLeaves(): bigint { + return this.leafCount; + } + + getRoot(): Buffer { + // we could recompute it, but it's way cheaper to just store the root + return this.historicalRoot; + } + + async getLeafValue(index: bigint): Promise { + const leafLevel = this.getDepth(); + const blockNumber = await this.#getBlockNumberThatModifiedNode(leafLevel, index); + + // leaf hasn't been set yet + if (typeof blockNumber === 'undefined') { + return undefined; + } + + // leaf was set some time in the past + if (blockNumber <= this.block) { + return this.db.get(historicalNodeKey(this.tree.getName(), leafLevel, index)); + } + + // leaf has been set but in a block in the future + return undefined; + } + + async #getHistoricalNodeValue(level: number, index: bigint): Promise { + const blockNumber = await this.#getBlockNumberThatModifiedNode(level, index); + + // node has never been set + if (typeof blockNumber === 'undefined') { + return this.tree.getZeroHash(level); + } + + // node was set some time in the past + if (blockNumber <= this.block) { + return this.db.get(historicalNodeKey(this.tree.getName(), level, index)); + } + + // the node has been modified since this snapshot was taken + // because we're working with an AppendOnly tree, historical leaves never change + // so what we do instead is rebuild this Merkle path up using zero hashes as needed + // worst case this will do O(H) hashes + // + // we first check if this subtree was touched by the block + // compare how many leaves this block added to the leaf interval of this subtree + // if they don't intersect then the whole subtree was a hash of zero + // if they do then we need to rebuild the merkle tree + const depth = this.tree.getDepth(); + const leafStart = index * 2n ** BigInt(depth - level); + if (leafStart >= this.leafCount) { + return this.tree.getZeroHash(level); + } + + const [lhs, rhs] = await Promise.all([ + this.#getHistoricalNodeValue(level + 1, 2n * index), + this.#getHistoricalNodeValue(level + 1, 2n * index + 1n), + ]); + + return this.hasher.hash(lhs, rhs); + } + + async #getBlockNumberThatModifiedNode(level: number, index: bigint): Promise { + try { + const value: Buffer | string = await this.db.get(nodeModifiedAtBlockKey(this.tree.getName(), level, index)); + return parseInt(value.toString(), 10); + } catch (err) { + return undefined; + } + } +} diff --git a/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts b/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts new file mode 100644 index 00000000000..d77204beafa --- /dev/null +++ b/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts @@ -0,0 +1,221 @@ +import { SiblingPath } from '@aztec/types'; + +import { LevelUp, LevelUpChain } from 'levelup'; + +import { TreeBase } from '../tree_base.js'; +import { TreeSnapshot, TreeSnapshotBuilder } from './snapshot_builder.js'; + +// key for a node's children +const snapshotChildKey = (node: Buffer, child: 0 | 1) => + Buffer.concat([Buffer.from('snapshot:node:'), node, Buffer.from(':' + child)]); + +// metadata for a snapshot +const snapshotRootKey = (treeName: string, block: number) => `snapshot:root:${treeName}:${block}`; +const snapshotNumLeavesKey = (treeName: string, block: number) => `snapshot:numLeaves:${treeName}:${block}`; + +/** + * Builds a full snapshot of a tree. This implementation works for any Merkle tree and stores + * it in a database in a similar way to how a tree is stored in memory, using pointers. + * + * Sharing the same database between versions and trees is recommended as the trees would share + * structure. + * + * Implement the protected method `handleLeaf` to store any additional data you need for each leaf. + * + * Complexity: + * N - count of non-zero nodes in tree + * M - count of snapshots + * H - tree height + * Worst case space complexity: O(N * M) + * Sibling path access: O(H) database reads + */ +export abstract class BaseFullTreeSnapshotBuilder + implements TreeSnapshotBuilder +{ + constructor(protected db: LevelUp, protected tree: T) {} + + async snapshot(block: number): Promise { + const snapshotMetadata = await this.#getSnapshotMeta(block); + + if (snapshotMetadata) { + return this.openSnapshot(snapshotMetadata.root, snapshotMetadata.numLeaves); + } + + const batch = this.db.batch(); + const root = this.tree.getRoot(false); + const numLeaves = this.tree.getNumLeaves(false); + const depth = this.tree.getDepth(); + const queue: [Buffer, number, bigint][] = [[root, 0, 0n]]; + + // walk the tree breadth-first and store each of its nodes in the database + // for each node we save two keys + // :0 -> + // :1 -> + while (queue.length > 0) { + const [node, level, i] = queue.shift()!; + // check if the database already has a child for this tree + // if it does, then we know we've seen the whole subtree below it before + // and we don't have to traverse it anymore + // we use the left child here, but it could be anything that shows we've stored the node before + const exists: Buffer | undefined = await this.db.get(snapshotChildKey(node, 0)).catch(() => undefined); + if (exists) { + continue; + } + + if (level + 1 > depth) { + // short circuit if we've reached the leaf level + // otherwise getNode might throw if we ask for the children of a leaf + this.handleLeaf(i, node, batch); + continue; + } + + const [lhs, rhs] = await Promise.all([ + this.tree.getNode(level + 1, 2n * i), + this.tree.getNode(level + 1, 2n * i + 1n), + ]); + + // we want the zero hash at the children's level, not the node's level + const zeroHash = this.tree.getZeroHash(level + 1); + + batch.put(snapshotChildKey(node, 0), lhs ?? zeroHash); + batch.put(snapshotChildKey(node, 1), rhs ?? zeroHash); + + // enqueue the children only if they're not zero hashes + if (lhs) { + queue.push([lhs, level + 1, 2n * i]); + } + + if (rhs) { + queue.push([rhs, level + 1, 2n * i + 1n]); + } + } + + batch.put(snapshotRootKey(this.tree.getName(), block), root); + batch.put(snapshotNumLeavesKey(this.tree.getName(), block), String(numLeaves)); + await batch.write(); + + return this.openSnapshot(root, numLeaves); + } + + protected handleLeaf(_index: bigint, _node: Buffer, _batch: LevelUpChain) { + return; + } + + async getSnapshot(version: number): Promise { + const snapshotMetadata = await this.#getSnapshotMeta(version); + + if (!snapshotMetadata) { + throw new Error(`Version ${version} does not exist for tree ${this.tree.getName()}`); + } + + return this.openSnapshot(snapshotMetadata.root, snapshotMetadata.numLeaves); + } + + protected abstract openSnapshot(root: Buffer, numLeaves: bigint): S; + + async #getSnapshotMeta(block: number): Promise< + | { + /** The root of the tree snapshot */ + root: Buffer; + /** The number of leaves in the tree snapshot */ + numLeaves: bigint; + } + | undefined + > { + try { + const treeName = this.tree.getName(); + const root = await this.db.get(snapshotRootKey(treeName, block)); + const numLeaves = BigInt(await this.db.get(snapshotNumLeavesKey(treeName, block))); + return { root, numLeaves }; + } catch (err) { + return undefined; + } + } +} + +/** + * A source of sibling paths from a snapshot tree + */ +export class BaseFullTreeSnapshot implements TreeSnapshot { + constructor( + protected db: LevelUp, + protected historicRoot: Buffer, + protected numLeaves: bigint, + protected tree: TreeBase, + ) {} + + async getSiblingPath(index: bigint): Promise> { + const siblings: Buffer[] = []; + + for await (const [_node, sibling] of this.pathFromRootToLeaf(index)) { + siblings.push(sibling); + } + + // we got the siblings we were looking for, but they are in root-leaf order + // reverse them here so we have leaf-root (what SiblingPath expects) + siblings.reverse(); + + return new SiblingPath(this.tree.getDepth() as N, siblings); + } + + async getLeafValue(index: bigint): Promise { + let leafNode: Buffer | undefined = undefined; + for await (const [node, _sibling] of this.pathFromRootToLeaf(index)) { + leafNode = node; + } + + return leafNode; + } + + getDepth(): number { + return this.tree.getDepth(); + } + + getRoot(): Buffer { + return this.historicRoot; + } + + getNumLeaves(): bigint { + return this.numLeaves; + } + + protected async *pathFromRootToLeaf(leafIndex: bigint) { + const root = this.historicRoot; + const pathFromRoot = this.#getPathFromRoot(leafIndex); + + let node: Buffer = root; + for (let i = 0; i < pathFromRoot.length; i++) { + // get both children. We'll need both anyway (one to keep track of, the other to walk down to) + const children: [Buffer, Buffer] = await Promise.all([ + this.db.get(snapshotChildKey(node, 0)), + this.db.get(snapshotChildKey(node, 1)), + ]).catch(() => [this.tree.getZeroHash(i + 1), this.tree.getZeroHash(i + 1)]); + const next = children[pathFromRoot[i]]; + const sibling = children[(pathFromRoot[i] + 1) % 2]; + + yield [next, sibling]; + + node = next; + } + } + + /** + * Calculates the path from the root to the target leaf. Returns an array of 0s and 1s, + * each 0 represents walking down a left child and each 1 walking down to the child on the right. + * + * @param leafIndex - The target leaf + * @returns An array of 0s and 1s + */ + #getPathFromRoot(leafIndex: bigint): ReadonlyArray<0 | 1> { + const path: Array<0 | 1> = []; + let level = this.tree.getDepth(); + while (level > 0) { + path.push(leafIndex & 0x01n ? 1 : 0); + leafIndex >>= 1n; + level--; + } + + path.reverse(); + return path; + } +} diff --git a/yarn-project/merkle-tree/src/snapshots/full_snapshot.test.ts b/yarn-project/merkle-tree/src/snapshots/full_snapshot.test.ts new file mode 100644 index 00000000000..3f2cc2af791 --- /dev/null +++ b/yarn-project/merkle-tree/src/snapshots/full_snapshot.test.ts @@ -0,0 +1,27 @@ +import levelup, { LevelUp } from 'levelup'; + +import { Pedersen, StandardTree, newTree } from '../index.js'; +import { createMemDown } from '../test/utils/create_mem_down.js'; +import { FullTreeSnapshotBuilder } from './full_snapshot.js'; +import { describeSnapshotBuilderTestSuite } from './snapshot_builder_test_suite.js'; + +describe('FullSnapshotBuilder', () => { + let tree: StandardTree; + let snapshotBuilder: FullTreeSnapshotBuilder; + let db: LevelUp; + + beforeEach(async () => { + db = levelup(createMemDown()); + tree = await newTree(StandardTree, db, new Pedersen(), 'test', 4); + snapshotBuilder = new FullTreeSnapshotBuilder(db, tree); + }); + + describeSnapshotBuilderTestSuite( + () => tree, + () => snapshotBuilder, + async () => { + const newLeaves = Array.from({ length: 2 }).map(() => Buffer.from(Math.random().toString())); + await tree.appendLeaves(newLeaves); + }, + ); +}); diff --git a/yarn-project/merkle-tree/src/snapshots/full_snapshot.ts b/yarn-project/merkle-tree/src/snapshots/full_snapshot.ts new file mode 100644 index 00000000000..c78d0ebb188 --- /dev/null +++ b/yarn-project/merkle-tree/src/snapshots/full_snapshot.ts @@ -0,0 +1,26 @@ +import { TreeBase } from '../tree_base.js'; +import { BaseFullTreeSnapshot, BaseFullTreeSnapshotBuilder } from './base_full_snapshot.js'; +import { TreeSnapshot, TreeSnapshotBuilder } from './snapshot_builder.js'; + +/** + * Builds a full snapshot of a tree. This implementation works for any Merkle tree and stores + * it in a database in a similar way to how a tree is stored in memory, using pointers. + * + * Sharing the same database between versions and trees is recommended as the trees would share + * structure. + * + * Complexity: + * N - count of non-zero nodes in tree + * M - count of snapshots + * H - tree height + * Worst case space complexity: O(N * M) + * Sibling path access: O(H) database reads + */ +export class FullTreeSnapshotBuilder + extends BaseFullTreeSnapshotBuilder + implements TreeSnapshotBuilder +{ + protected openSnapshot(root: Buffer, numLeaves: bigint): TreeSnapshot { + return new BaseFullTreeSnapshot(this.db, root, numLeaves, this.tree); + } +} diff --git a/yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.test.ts b/yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.test.ts new file mode 100644 index 00000000000..631531fcfc2 --- /dev/null +++ b/yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.test.ts @@ -0,0 +1,96 @@ +import levelup, { LevelUp } from 'levelup'; + +import { Pedersen, newTree } from '../index.js'; +import { StandardIndexedTreeWithAppend } from '../standard_indexed_tree/test/standard_indexed_tree_with_append.js'; +import { createMemDown } from '../test/utils/create_mem_down.js'; +import { IndexedTreeSnapshotBuilder } from './indexed_tree_snapshot.js'; +import { describeSnapshotBuilderTestSuite } from './snapshot_builder_test_suite.js'; + +describe('IndexedTreeSnapshotBuilder', () => { + let db: LevelUp; + let tree: StandardIndexedTreeWithAppend; + let snapshotBuilder: IndexedTreeSnapshotBuilder; + + beforeEach(async () => { + db = levelup(createMemDown()); + tree = await newTree(StandardIndexedTreeWithAppend, db, new Pedersen(), 'test', 4); + snapshotBuilder = new IndexedTreeSnapshotBuilder(db, tree); + }); + + describeSnapshotBuilderTestSuite( + () => tree, + () => snapshotBuilder, + async () => { + const newLeaves = Array.from({ length: 2 }).map(() => Buffer.from(Math.random().toString())); + await tree.appendLeaves(newLeaves); + }, + ); + + describe('getSnapshot', () => { + it('returns historical leaf data', async () => { + await tree.appendLeaves([Buffer.from('a'), Buffer.from('b'), Buffer.from('c')]); + await tree.commit(); + const expectedLeavesAtBlock1 = await Promise.all([ + tree.getLatestLeafDataCopy(0, false), + tree.getLatestLeafDataCopy(1, false), + tree.getLatestLeafDataCopy(2, false), + // id'expect these to be undefined, but leaf 3 isn't? + // must be some indexed-tree quirk I don't quite understand yet + tree.getLatestLeafDataCopy(3, false), + tree.getLatestLeafDataCopy(4, false), + tree.getLatestLeafDataCopy(5, false), + ]); + + await snapshotBuilder.snapshot(1); + + await tree.appendLeaves([Buffer.from('d'), Buffer.from('e'), Buffer.from('f')]); + await tree.commit(); + const expectedLeavesAtBlock2 = await Promise.all([ + tree.getLatestLeafDataCopy(0, false), + tree.getLatestLeafDataCopy(1, false), + tree.getLatestLeafDataCopy(2, false), + tree.getLatestLeafDataCopy(3, false), + tree.getLatestLeafDataCopy(4, false), + tree.getLatestLeafDataCopy(5, false), + ]); + + await snapshotBuilder.snapshot(2); + + const snapshot1 = await snapshotBuilder.getSnapshot(1); + const actualLeavesAtBlock1 = await Promise.all([ + snapshot1.getLatestLeafDataCopy(0n), + snapshot1.getLatestLeafDataCopy(1n), + snapshot1.getLatestLeafDataCopy(2n), + snapshot1.getLatestLeafDataCopy(3n), + snapshot1.getLatestLeafDataCopy(4n), + snapshot1.getLatestLeafDataCopy(5n), + ]); + expect(actualLeavesAtBlock1).toEqual(expectedLeavesAtBlock1); + + const snapshot2 = await snapshotBuilder.getSnapshot(2); + const actualLeavesAtBlock2 = await Promise.all([ + snapshot2.getLatestLeafDataCopy(0n), + snapshot2.getLatestLeafDataCopy(1n), + snapshot2.getLatestLeafDataCopy(2n), + snapshot2.getLatestLeafDataCopy(3n), + snapshot2.getLatestLeafDataCopy(4n), + snapshot2.getLatestLeafDataCopy(5n), + ]); + expect(actualLeavesAtBlock2).toEqual(expectedLeavesAtBlock2); + }); + }); + + describe('findIndexOfPreviousValue', () => { + it('returns the index of the leaf with the closest value to the given value', async () => { + await tree.appendLeaves([Buffer.from('a'), Buffer.from('f'), Buffer.from('d')]); + await tree.commit(); + const snapshot = await snapshotBuilder.snapshot(1); + const historicalPrevValue = tree.findIndexOfPreviousValue(2n, false); + + await tree.appendLeaves([Buffer.from('c'), Buffer.from('b'), Buffer.from('e')]); + await tree.commit(); + + await expect(snapshot.findIndexOfPreviousValue(2n)).resolves.toEqual(historicalPrevValue); + }); + }); +}); diff --git a/yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.ts b/yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.ts new file mode 100644 index 00000000000..6725bd394e5 --- /dev/null +++ b/yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.ts @@ -0,0 +1,92 @@ +import { toBufferBE } from '@aztec/foundation/bigint-buffer'; +import { LeafData } from '@aztec/types'; + +import { LevelUp, LevelUpChain } from 'levelup'; + +import { IndexedTree } from '../interfaces/indexed_tree.js'; +import { decodeTreeValue, encodeTreeValue } from '../standard_indexed_tree/standard_indexed_tree.js'; +import { TreeBase } from '../tree_base.js'; +import { BaseFullTreeSnapshot, BaseFullTreeSnapshotBuilder } from './base_full_snapshot.js'; +import { IndexedTreeSnapshot, TreeSnapshotBuilder } from './snapshot_builder.js'; + +const snapshotLeafValue = (node: Buffer, index: bigint) => + Buffer.concat([Buffer.from('snapshot:leaf:'), node, Buffer.from(':' + index)]); + +/** a */ +export class IndexedTreeSnapshotBuilder + extends BaseFullTreeSnapshotBuilder + implements TreeSnapshotBuilder +{ + constructor(db: LevelUp, tree: IndexedTree & TreeBase) { + super(db, tree); + } + + protected openSnapshot(root: Buffer, numLeaves: bigint): IndexedTreeSnapshot { + return new IndexedTreeSnapshotImpl(this.db, root, numLeaves, this.tree); + } + + protected handleLeaf(index: bigint, node: Buffer, batch: LevelUpChain) { + const leafData = this.tree.getLatestLeafDataCopy(Number(index), false); + if (leafData) { + batch.put(snapshotLeafValue(node, index), encodeTreeValue(leafData)); + } + } +} + +/** A snapshot of an indexed tree at a particular point in time */ +class IndexedTreeSnapshotImpl extends BaseFullTreeSnapshot implements IndexedTreeSnapshot { + async getLeafValue(index: bigint): Promise { + const leafData = await this.getLatestLeafDataCopy(index); + return leafData ? toBufferBE(leafData.value, 32) : undefined; + } + + async getLatestLeafDataCopy(index: bigint): Promise { + const leafNode = await super.getLeafValue(index); + const leafValue = await this.db.get(snapshotLeafValue(leafNode!, index)).catch(() => undefined); + if (leafValue) { + return decodeTreeValue(leafValue); + } else { + return undefined; + } + } + + async findIndexOfPreviousValue(newValue: bigint): Promise<{ + /** + * The index of the found leaf. + */ + index: number; + /** + * A flag indicating if the corresponding leaf's value is equal to `newValue`. + */ + alreadyPresent: boolean; + }> { + const numLeaves = this.getNumLeaves(); + const diff: bigint[] = []; + + for (let i = 0; i < numLeaves; i++) { + // this is very inefficient + const storedLeaf = await this.getLatestLeafDataCopy(BigInt(i))!; + + // The stored leaf can be undefined if it addresses an empty leaf + // If the leaf is empty we do the same as if the leaf was larger + if (storedLeaf === undefined) { + diff.push(newValue); + } else if (storedLeaf.value > newValue) { + diff.push(newValue); + } else if (storedLeaf.value === newValue) { + return { index: i, alreadyPresent: true }; + } else { + diff.push(newValue - storedLeaf.value); + } + } + + let minIndex = 0; + for (let i = 1; i < diff.length; i++) { + if (diff[i] < diff[minIndex]) { + minIndex = i; + } + } + + return { index: minIndex, alreadyPresent: false }; + } +} diff --git a/yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts b/yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts new file mode 100644 index 00000000000..a6722306301 --- /dev/null +++ b/yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts @@ -0,0 +1,75 @@ +import { LeafData, SiblingPath } from '@aztec/types'; + +/** + * An interface for a tree that can record snapshots of its contents. + */ +export interface TreeSnapshotBuilder { + /** + * Creates a snapshot of the tree at the given version. + * @param block - The version to snapshot the tree at. + */ + snapshot(block: number): Promise; + + /** + * Returns a snapshot of the tree at the given version. + * @param block - The version of the snapshot to return. + */ + getSnapshot(block: number): Promise; +} + +/** + * A tree snapshot + */ +export interface TreeSnapshot { + /** + * Returns the current root of the tree. + */ + getRoot(): Buffer; + + /** + * Returns the number of leaves in the tree. + */ + getDepth(): number; + + /** + * Returns the number of leaves in the tree. + */ + getNumLeaves(): bigint; + + /** + * Returns the value of a leaf at the specified index. + * @param index - The index of the leaf value to be returned. + */ + getLeafValue(index: bigint): Promise; + + /** + * Returns the sibling path for a requested leaf index. + * @param index - The index of the leaf for which a sibling path is required. + */ + getSiblingPath(index: bigint): Promise>; +} + +/** A snapshot of an indexed tree */ +export interface IndexedTreeSnapshot extends TreeSnapshot { + /** + * Gets the historical data for a leaf + * @param index - The index of the leaf to get the data for + */ + getLatestLeafDataCopy(index: bigint): Promise; + + /** + * Finds the index of the largest leaf whose value is less than or equal to the provided value. + * @param newValue - The new value to be inserted into the tree. + * @returns The found leaf index and a flag indicating if the corresponding leaf's value is equal to `newValue`. + */ + findIndexOfPreviousValue(newValue: bigint): Promise<{ + /** + * The index of the found leaf. + */ + index: number; + /** + * A flag indicating if the corresponding leaf's value is equal to `newValue`. + */ + alreadyPresent: boolean; + }>; +} diff --git a/yarn-project/merkle-tree/src/snapshots/snapshot_builder_test_suite.ts b/yarn-project/merkle-tree/src/snapshots/snapshot_builder_test_suite.ts new file mode 100644 index 00000000000..3b66c36164c --- /dev/null +++ b/yarn-project/merkle-tree/src/snapshots/snapshot_builder_test_suite.ts @@ -0,0 +1,197 @@ +import { TreeBase } from '../tree_base.js'; +import { TreeSnapshotBuilder } from './snapshot_builder.js'; + +/** Creates a test suit for snapshots */ +export function describeSnapshotBuilderTestSuite( + getTree: () => T, + getSnapshotBuilder: () => S, + modifyTree: (tree: T) => Promise, +) { + describe('SnapshotBuilder', () => { + let tree: T; + let snapshotBuilder: S; + let leaves: bigint[]; + + beforeEach(() => { + tree = getTree(); + snapshotBuilder = getSnapshotBuilder(); + + leaves = Array.from({ length: 4 }).map(() => BigInt(Math.floor(Math.random() * 2 ** tree.getDepth()))); + }); + + describe('snapshot', () => { + it('takes snapshots', async () => { + await modifyTree(tree); + await tree.commit(); + await expect(snapshotBuilder.snapshot(1)).resolves.toBeDefined(); + }); + + it('is idempotent', async () => { + await modifyTree(tree); + await tree.commit(); + + const block = 1; + const snapshot = await snapshotBuilder.snapshot(block); + await expect(snapshotBuilder.snapshot(block)).resolves.toEqual(snapshot); + }); + + it('returns the same path if tree has not diverged', async () => { + await modifyTree(tree); + await tree.commit(); + const snapshot = await snapshotBuilder.snapshot(1); + + const historicPaths = await Promise.all(leaves.map(leaf => snapshot.getSiblingPath(leaf))); + const expectedPaths = await Promise.all(leaves.map(leaf => tree.getSiblingPath(leaf, false))); + + for (const [index, path] of historicPaths.entries()) { + expect(path).toEqual(expectedPaths[index]); + } + }); + + it('returns historic paths if tree has diverged and no new snapshots have been taken', async () => { + await modifyTree(tree); + await tree.commit(); + const snapshot = await snapshotBuilder.snapshot(1); + + const expectedPaths = await Promise.all(leaves.map(leaf => tree.getSiblingPath(leaf, false))); + + await modifyTree(tree); + await tree.commit(); + + const historicPaths = await Promise.all(leaves.map(leaf => snapshot.getSiblingPath(leaf))); + + for (const [index, path] of historicPaths.entries()) { + expect(path).toEqual(expectedPaths[index]); + } + }); + + it('retains old snapshots even if new one are created', async () => { + await modifyTree(tree); + await tree.commit(); + + const expectedPaths = await Promise.all(leaves.map(leaf => tree.getSiblingPath(leaf, false))); + + const snapshot = await snapshotBuilder.snapshot(1); + + await modifyTree(tree); + await tree.commit(); + + await snapshotBuilder.snapshot(2); + + // check that snapshot 2 has not influenced snapshot(1) at all + const historicPaths = await Promise.all(leaves.map(leaf => snapshot.getSiblingPath(leaf))); + + for (const [index, path] of historicPaths.entries()) { + expect(path).toEqual(expectedPaths[index]); + } + }); + + it('retains old snapshots even if new one are created and the tree diverges', async () => { + await modifyTree(tree); + await tree.commit(); + + const expectedPaths = await Promise.all(leaves.map(leaf => tree.getSiblingPath(leaf, false))); + + const snapshot = await snapshotBuilder.snapshot(1); + + await modifyTree(tree); + await tree.commit(); + + await snapshotBuilder.snapshot(2); + + await modifyTree(tree); + await tree.commit(); + + // check that snapshot 2 has not influenced snapshot(1) at all + // and that the diverging tree does not influence the old snapshot + const historicPaths = await Promise.all(leaves.map(leaf => snapshot.getSiblingPath(leaf))); + + for (const [index, path] of historicPaths.entries()) { + expect(path).toEqual(expectedPaths[index]); + } + }); + }); + + describe('getSnapshot', () => { + it('returns old snapshots', async () => { + await modifyTree(tree); + await tree.commit(); + const expectedPaths = await Promise.all(leaves.map(leaf => tree.getSiblingPath(leaf, false))); + await snapshotBuilder.snapshot(1); + + for (let i = 2; i < 5; i++) { + await modifyTree(tree); + await tree.commit(); + await snapshotBuilder.snapshot(i); + } + + const firstSnapshot = await snapshotBuilder.getSnapshot(1); + const historicPaths = await Promise.all(leaves.map(leaf => firstSnapshot.getSiblingPath(leaf))); + + for (const [index, path] of historicPaths.entries()) { + expect(path).toEqual(expectedPaths[index]); + } + }); + + it('throws if an unknown snapshot is requested', async () => { + await modifyTree(tree); + await tree.commit(); + await snapshotBuilder.snapshot(1); + + await expect(snapshotBuilder.getSnapshot(2)).rejects.toThrow(); + }); + }); + + describe('getRoot', () => { + it('returns the historical root of the tree when the snapshot was taken', async () => { + await modifyTree(tree); + await tree.commit(); + const snapshot = await snapshotBuilder.snapshot(1); + const historicalRoot = tree.getRoot(false); + + await modifyTree(tree); + await tree.commit(); + + expect(snapshot.getRoot()).toEqual(historicalRoot); + expect(snapshot.getRoot()).not.toEqual(tree.getRoot(false)); + }); + }); + + describe('getDepth', () => { + it('returns the same depth as the tree', async () => { + await modifyTree(tree); + await tree.commit(); + const snapshot = await snapshotBuilder.snapshot(1); + expect(snapshot.getDepth()).toEqual(tree.getDepth()); + }); + }); + + describe('getNumLeaves', () => { + it('returns the historical leaves count when the snapshot was taken', async () => { + await modifyTree(tree); + await tree.commit(); + const snapshot = await snapshotBuilder.snapshot(1); + const historicalNumLeaves = tree.getNumLeaves(false); + + await modifyTree(tree); + await tree.commit(); + + expect(snapshot.getNumLeaves()).toEqual(historicalNumLeaves); + }); + }); + + describe('getLeafValue', () => { + it('returns the historical leaf value when the snapshot was taken', async () => { + await modifyTree(tree); + await tree.commit(); + const snapshot = await snapshotBuilder.snapshot(1); + const historicalLeafValue = await tree.getLeafValue(0n, false); + + await modifyTree(tree); + await tree.commit(); + + await expect(snapshot.getLeafValue(0n)).resolves.toEqual(historicalLeafValue); + }); + }); + }); +} diff --git a/yarn-project/merkle-tree/src/sparse_tree/sparse_tree.ts b/yarn-project/merkle-tree/src/sparse_tree/sparse_tree.ts index 463c6431fa9..92cdc4152fc 100644 --- a/yarn-project/merkle-tree/src/sparse_tree/sparse_tree.ts +++ b/yarn-project/merkle-tree/src/sparse_tree/sparse_tree.ts @@ -1,10 +1,14 @@ import { UpdateOnlyTree } from '../interfaces/update_only_tree.js'; +import { FullTreeSnapshotBuilder } from '../snapshots/full_snapshot.js'; +import { TreeSnapshot } from '../snapshots/snapshot_builder.js'; import { INITIAL_LEAF, TreeBase } from '../tree_base.js'; /** * A Merkle tree implementation that uses a LevelDB database to store the tree. */ export class SparseTree extends TreeBase implements UpdateOnlyTree { + #snapshotBuilder = new FullTreeSnapshotBuilder(this.db, this); + /** * Updates a leaf in the tree. * @param leaf - New contents of the leaf. @@ -29,4 +33,12 @@ export class SparseTree extends TreeBase implements UpdateOnlyTree { this.cachedSize = (this.cachedSize ?? this.size) + 1n; } } + + public snapshot(block: number): Promise { + return this.#snapshotBuilder.snapshot(block); + } + + public getSnapshot(block: number): Promise { + return this.#snapshotBuilder.getSnapshot(block); + } } diff --git a/yarn-project/merkle-tree/src/standard_indexed_tree/standard_indexed_tree.ts b/yarn-project/merkle-tree/src/standard_indexed_tree/standard_indexed_tree.ts index f9d44353fa9..ebbf3a3d0ee 100644 --- a/yarn-project/merkle-tree/src/standard_indexed_tree/standard_indexed_tree.ts +++ b/yarn-project/merkle-tree/src/standard_indexed_tree/standard_indexed_tree.ts @@ -4,6 +4,8 @@ import { createDebugLogger } from '@aztec/foundation/log'; import { LeafData, SiblingPath } from '@aztec/types'; import { BatchInsertionResult, IndexedTree } from '../interfaces/indexed_tree.js'; +import { IndexedTreeSnapshotBuilder } from '../snapshots/indexed_tree_snapshot.js'; +import { IndexedTreeSnapshot } from '../snapshots/snapshot_builder.js'; import { TreeBase } from '../tree_base.js'; const log = createDebugLogger('aztec:standard-indexed-tree'); @@ -54,15 +56,14 @@ function getEmptyLowLeafWitness(treeHeight: N): LowLeafWitness }; } -// eslint-disable-next-line @typescript-eslint/no-unused-vars -const encodeTreeValue = (leafData: LeafData) => { +export const encodeTreeValue = (leafData: LeafData) => { const valueAsBuffer = toBufferBE(leafData.value, 32); const indexAsBuffer = toBufferBE(leafData.nextIndex, 32); const nextValueAsBuffer = toBufferBE(leafData.nextValue, 32); return Buffer.concat([valueAsBuffer, indexAsBuffer, nextValueAsBuffer]); }; -const decodeTreeValue = (buf: Buffer) => { +export const decodeTreeValue = (buf: Buffer) => { const value = toBigIntBE(buf.subarray(0, 32)); const nextIndex = toBigIntBE(buf.subarray(32, 64)); const nextValue = toBigIntBE(buf.subarray(64, 96)); @@ -77,6 +78,8 @@ const decodeTreeValue = (buf: Buffer) => { * Indexed merkle tree. */ export class StandardIndexedTree extends TreeBase implements IndexedTree { + #snapshotBuilder = new IndexedTreeSnapshotBuilder(this.db, this); + protected leaves: LeafData[] = []; protected cachedLeaves: { [key: number]: LeafData } = {}; @@ -514,6 +517,14 @@ export class StandardIndexedTree extends TreeBase implements IndexedTree { return fullSiblingPath.getSubtreeSiblingPath(subtreeHeight); } + snapshot(blockNumber: number): Promise { + return this.#snapshotBuilder.snapshot(blockNumber); + } + + getSnapshot(block: number): Promise { + return this.#snapshotBuilder.getSnapshot(block); + } + /** * Encodes leaves and appends them to a tree. * @param leaves - Leaves to encode. diff --git a/yarn-project/merkle-tree/src/standard_tree/standard_tree.ts b/yarn-project/merkle-tree/src/standard_tree/standard_tree.ts index 2c479168def..0b92572a4b8 100644 --- a/yarn-project/merkle-tree/src/standard_tree/standard_tree.ts +++ b/yarn-project/merkle-tree/src/standard_tree/standard_tree.ts @@ -1,3 +1,4 @@ +import { AppendOnlySnapshotBuilder, TreeSnapshot } from '../index.js'; import { AppendOnlyTree } from '../interfaces/append_only_tree.js'; import { TreeBase } from '../tree_base.js'; @@ -5,6 +6,8 @@ import { TreeBase } from '../tree_base.js'; * A Merkle tree implementation that uses a LevelDB database to store the tree. */ export class StandardTree extends TreeBase implements AppendOnlyTree { + #snapshotBuilder = new AppendOnlySnapshotBuilder(this.db, this, this.hasher); + /** * Appends the given leaves to the tree. * @param leaves - The leaves to append. @@ -13,4 +16,12 @@ export class StandardTree extends TreeBase implements AppendOnlyTree { public async appendLeaves(leaves: Buffer[]): Promise { await super.appendLeaves(leaves); } + + public snapshot(block: number): Promise { + return this.#snapshotBuilder.snapshot(block); + } + + public getSnapshot(block: number): Promise { + return this.#snapshotBuilder.getSnapshot(block); + } } diff --git a/yarn-project/merkle-tree/src/tree_base.ts b/yarn-project/merkle-tree/src/tree_base.ts index 6b715280380..c57a0499171 100644 --- a/yarn-project/merkle-tree/src/tree_base.ts +++ b/yarn-project/merkle-tree/src/tree_base.ts @@ -150,6 +150,26 @@ export abstract class TreeBase implements MerkleTree { return this.getLatestValueAtIndex(this.depth, index, includeUncommitted); } + public getNode(level: number, index: bigint): Promise { + if (level < 0 || level > this.depth) { + throw Error('Invalid level: ' + level); + } + + if (index < 0 || index >= 2n ** BigInt(level)) { + throw Error('Invalid index: ' + index); + } + + return this.dbGet(indexToKeyHash(this.name, level, index)); + } + + public getZeroHash(level: number): Buffer { + if (level <= 0 || level > this.depth) { + throw new Error('Invalid level'); + } + + return this.zeroHashes[level - 1]; + } + /** * Clears the cache. */ diff --git a/yarn-project/world-state/src/merkle-tree/merkle_tree_snapshot_operations_facade.ts b/yarn-project/world-state/src/merkle-tree/merkle_tree_snapshot_operations_facade.ts new file mode 100644 index 00000000000..93c94d19163 --- /dev/null +++ b/yarn-project/world-state/src/merkle-tree/merkle_tree_snapshot_operations_facade.ts @@ -0,0 +1,143 @@ +import { Fr } from '@aztec/circuits.js'; +import { BatchInsertionResult, IndexedTreeSnapshot, TreeSnapshot } from '@aztec/merkle-tree'; +import { LeafData, MerkleTreeId, SiblingPath } from '@aztec/types'; + +import { CurrentTreeRoots, HandleL2BlockResult, MerkleTreeDb, MerkleTreeOperations, TreeInfo } from '../index.js'; + +/** + * Merkle tree operations on readonly tree snapshots. + */ +export class MerkleTreeSnapshotOperationsFacade implements MerkleTreeOperations { + #treesDb: MerkleTreeDb; + #blockNumber: number; + #treeSnapshots: ReadonlyArray = []; + + constructor(trees: MerkleTreeDb, blockNumber: number) { + this.#treesDb = trees; + this.#blockNumber = blockNumber; + } + + async #getTreeSnapshot(merkleTreeId: number): Promise { + if (this.#treeSnapshots[merkleTreeId]) { + return this.#treeSnapshots[merkleTreeId]; + } + + this.#treeSnapshots = await this.#treesDb.getSnapshot(this.#blockNumber); + return this.#treeSnapshots[merkleTreeId]!; + } + + async findLeafIndex(treeId: MerkleTreeId, value: Buffer): Promise { + const tree = await this.#getTreeSnapshot(treeId); + const numLeaves = tree.getNumLeaves(); + for (let i = 0n; i < numLeaves; i++) { + const currentValue = await tree.getLeafValue(i); + if (currentValue && currentValue.equals(value)) { + return i; + } + } + return undefined; + } + + getLatestGlobalVariablesHash(): Promise { + return Promise.reject(new Error('not implemented')); + } + + async getLeafData(treeId: MerkleTreeId.NULLIFIER_TREE, index: number): Promise { + const snapshot = (await this.#getTreeSnapshot(treeId)) as IndexedTreeSnapshot; + return snapshot.getLatestLeafDataCopy(BigInt(index)); + } + + async getLeafValue(treeId: MerkleTreeId, index: bigint): Promise { + const snapshot = await this.#getTreeSnapshot(treeId); + return snapshot.getLeafValue(BigInt(index)); + } + + getPreviousValueIndex( + _treeId: MerkleTreeId.NULLIFIER_TREE, + _value: bigint, + ): Promise<{ + /** + * The index of the found leaf. + */ + index: number; + /** + * A flag indicating if the corresponding leaf's value is equal to `newValue`. + */ + alreadyPresent: boolean; + }> { + return Promise.reject(new Error('not implemented')); + } + + async getSiblingPath(treeId: MerkleTreeId, index: bigint): Promise> { + const snapshot = await this.#getTreeSnapshot(treeId); + return snapshot.getSiblingPath(index); + } + + async getTreeInfo(treeId: MerkleTreeId): Promise { + const snapshot = await this.#getTreeSnapshot(treeId); + return { + depth: snapshot.getDepth(), + root: snapshot.getRoot(), + size: snapshot.getNumLeaves(), + treeId, + }; + } + + async getTreeRoots(): Promise { + const snapshots = await Promise.all([ + this.#getTreeSnapshot(MerkleTreeId.CONTRACT_TREE), + this.#getTreeSnapshot(MerkleTreeId.NULLIFIER_TREE), + this.#getTreeSnapshot(MerkleTreeId.NOTE_HASH_TREE), + this.#getTreeSnapshot(MerkleTreeId.PUBLIC_DATA_TREE), + this.#getTreeSnapshot(MerkleTreeId.L1_TO_L2_MESSAGES_TREE), + this.#getTreeSnapshot(MerkleTreeId.BLOCKS_TREE), + ]); + + return { + blocksTreeRoot: snapshots[MerkleTreeId.BLOCKS_TREE].getRoot(), + contractDataTreeRoot: snapshots[MerkleTreeId.CONTRACT_TREE].getRoot(), + l1Tol2MessagesTreeRoot: snapshots[MerkleTreeId.L1_TO_L2_MESSAGES_TREE].getRoot(), + noteHashTreeRoot: snapshots[MerkleTreeId.NOTE_HASH_TREE].getRoot(), + nullifierTreeRoot: snapshots[MerkleTreeId.NULLIFIER_TREE].getRoot(), + publicDataTreeRoot: snapshots[MerkleTreeId.PUBLIC_DATA_TREE].getRoot(), + }; + } + + appendLeaves(): Promise { + return Promise.reject(new Error('Tree snapshot operations are read-only')); + } + + batchInsert(): Promise< + BatchInsertionResult + > { + return Promise.reject(new Error('Tree snapshot operations are read-only')); + } + + updateBlocksTree(): Promise { + return Promise.reject(new Error('Tree snapshot operations are read-only')); + } + + commit(): Promise { + return Promise.reject(new Error('Tree snapshot operations are read-only')); + } + + handleL2Block(): Promise { + return Promise.reject(new Error('Tree snapshot operations are read-only')); + } + + rollback(): Promise { + return Promise.reject(new Error('Tree snapshot operations are read-only')); + } + + updateHistoricBlocksTree(): Promise { + return Promise.reject(new Error('Tree snapshot operations are read-only')); + } + + updateLatestGlobalVariablesHash(): Promise { + return Promise.reject(new Error('Tree snapshot operations are read-only')); + } + + updateLeaf(): Promise { + return Promise.reject(new Error('Tree snapshot operations are read-only')); + } +} diff --git a/yarn-project/world-state/src/synchronizer/server_world_state_synchronizer.ts b/yarn-project/world-state/src/synchronizer/server_world_state_synchronizer.ts index 896b85a55de..06d304b7e35 100644 --- a/yarn-project/world-state/src/synchronizer/server_world_state_synchronizer.ts +++ b/yarn-project/world-state/src/synchronizer/server_world_state_synchronizer.ts @@ -8,6 +8,7 @@ import { LevelUp } from 'levelup'; import { HandleL2BlockResult, MerkleTreeOperations, MerkleTrees } from '../index.js'; import { MerkleTreeOperationsFacade } from '../merkle-tree/merkle_tree_operations_facade.js'; +import { MerkleTreeSnapshotOperationsFacade } from '../merkle-tree/merkle_tree_snapshot_operations_facade.js'; import { WorldStateConfig } from './config.js'; import { WorldStateRunningState, WorldStateStatus, WorldStateSynchronizer } from './world_state_synchronizer.js'; @@ -52,6 +53,10 @@ export class ServerWorldStateSynchronizer implements WorldStateSynchronizer { return new MerkleTreeOperationsFacade(this.merkleTreeDb, false); } + public getSnapshot(blockNumber: number): MerkleTreeOperations { + return new MerkleTreeSnapshotOperationsFacade(this.merkleTreeDb, blockNumber); + } + public static async new( db: LevelUp, merkleTreeDb: MerkleTrees, diff --git a/yarn-project/world-state/src/synchronizer/world_state_synchronizer.ts b/yarn-project/world-state/src/synchronizer/world_state_synchronizer.ts index 39e75cb91cb..21e84b903ea 100644 --- a/yarn-project/world-state/src/synchronizer/world_state_synchronizer.ts +++ b/yarn-project/world-state/src/synchronizer/world_state_synchronizer.ts @@ -63,4 +63,11 @@ export interface WorldStateSynchronizer { * @returns An instance of MerkleTreeOperations that will not include uncommitted data. */ getCommitted(): MerkleTreeOperations; + + /** + * Returns a readonly instance of MerkleTreeOperations where the state is as it was at the given block number + * @param block - The block number to look at + * @returns An instance of MerkleTreeOperations + */ + getSnapshot(block: number): MerkleTreeOperations; } diff --git a/yarn-project/world-state/src/world-state-db/merkle_tree_db.ts b/yarn-project/world-state/src/world-state-db/merkle_tree_db.ts index b9dc631ff1f..13c6617513d 100644 --- a/yarn-project/world-state/src/world-state-db/merkle_tree_db.ts +++ b/yarn-project/world-state/src/world-state-db/merkle_tree_db.ts @@ -1,7 +1,7 @@ import { MAX_NEW_NULLIFIERS_PER_TX } from '@aztec/circuits.js'; import { Fr } from '@aztec/foundation/fields'; import { createDebugLogger } from '@aztec/foundation/log'; -import { BatchInsertionResult } from '@aztec/merkle-tree'; +import { BatchInsertionResult, IndexedTreeSnapshot, TreeSnapshot } from '@aztec/merkle-tree'; import { L2Block, LeafData, MerkleTreeId, SiblingPath } from '@aztec/types'; /** @@ -91,7 +91,13 @@ export type MerkleTreeDb = { [Property in keyof MerkleTreeOperations as Exclude]: WithIncludeUncommitted< MerkleTreeOperations[Property] >; -} & Pick; +} & Pick & { + /** + * Returns a snapshot of the current state of the trees. + * @param block - The block number to take the snapshot at. + */ + getSnapshot(block: number): Promise>; + }; /** * Defines the interface for operations on a set of Merkle Trees. diff --git a/yarn-project/world-state/src/world-state-db/merkle_trees.ts b/yarn-project/world-state/src/world-state-db/merkle_trees.ts index 3d167bf4a9c..311c071d8b1 100644 --- a/yarn-project/world-state/src/world-state-db/merkle_trees.ts +++ b/yarn-project/world-state/src/world-state-db/merkle_trees.ts @@ -521,6 +521,16 @@ export class MerkleTrees implements MerkleTreeDb { this.latestGlobalVariablesHash.rollback(); } + public getSnapshot(blockNumber: number) { + return Promise.all(this.trees.map(tree => tree.getSnapshot(blockNumber))); + } + + private async _snapshot(blockNumber: number): Promise { + for (const tree of this.trees) { + await tree.snapshot(blockNumber); + } + } + /** * Handles a single L2 block (i.e. Inserts the new commitments into the merkle tree). * @param l2Block - The L2 block to handle. @@ -599,6 +609,8 @@ export class MerkleTrees implements MerkleTreeDb { } } + await this._snapshot(l2Block.number); + return { isBlockOurs: ourBlock }; } }