From 3b169dd713e1e57b686f81e3d2beb08c7a86e1d3 Mon Sep 17 00:00:00 2001 From: Alex Gherghisan Date: Wed, 29 Nov 2023 11:33:43 +0000 Subject: [PATCH 1/8] feat: add tree snapshots --- yarn-project/merkle-tree/src/index.ts | 3 + .../snapshots/append_only_snapshot.test.ts | 47 +++++ .../src/snapshots/append_only_snapshot.ts | 179 ++++++++++++++++++ .../snapshots/incremental_snapshot.test.ts | 27 +++ .../src/snapshots/incremental_snapshot.ts | 156 +++++++++++++++ .../src/snapshots/snapshot_builder.ts | 18 ++ .../snapshots/snapshot_builder_test_suite.ts | 138 ++++++++++++++ yarn-project/merkle-tree/src/tree_base.ts | 20 ++ 8 files changed, 588 insertions(+) create mode 100644 yarn-project/merkle-tree/src/snapshots/append_only_snapshot.test.ts create mode 100644 yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts create mode 100644 yarn-project/merkle-tree/src/snapshots/incremental_snapshot.test.ts create mode 100644 yarn-project/merkle-tree/src/snapshots/incremental_snapshot.ts create mode 100644 yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts create mode 100644 yarn-project/merkle-tree/src/snapshots/snapshot_builder_test_suite.ts diff --git a/yarn-project/merkle-tree/src/index.ts b/yarn-project/merkle-tree/src/index.ts index de19e295bfd..9fa9f0e7aba 100644 --- a/yarn-project/merkle-tree/src/index.ts +++ b/yarn-project/merkle-tree/src/index.ts @@ -9,3 +9,6 @@ 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/incremental_snapshot.js'; +export * from './snapshots/append_only_snapshot.js'; 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..6b2cffe6e09 --- /dev/null +++ b/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.test.ts @@ -0,0 +1,47 @@ +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); + }); + + it('takes snapshots', async () => { + await tree.appendLeaves([Buffer.from('a'), Buffer.from('b'), Buffer.from('c')]); + await tree.commit(); + + const expectedPathAtSnapshot1 = await tree.getSiblingPath(1n, false); + + const snapshot1 = await snapshotBuilder.snapshot(1); + + await tree.appendLeaves([Buffer.from('d'), Buffer.from('e'), Buffer.from('f')]); + await tree.commit(); + + const expectedPathAtSnapshot2 = await tree.getSiblingPath(1n, false); + + const snapshot2 = await snapshotBuilder.snapshot(2); + + await expect(snapshot1.getSiblingPath(1n, false)).resolves.toEqual(expectedPathAtSnapshot1); + await expect(snapshot2.getSiblingPath(1n, false)).resolves.toEqual(expectedPathAtSnapshot2); + }); + + 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..844e3ca9aed --- /dev/null +++ b/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts @@ -0,0 +1,179 @@ +import { Hasher, SiblingPath } from '@aztec/types'; + +import { LevelUp } from 'levelup'; + +import { AppendOnlyTree } from '../interfaces/append_only_tree.js'; +import { SiblingPathSource } from '../interfaces/merkle_tree.js'; +import { TreeBase } from '../tree_base.js'; +import { SnapshotBuilder } from './snapshot_builder.js'; + +const nodeVersionKey = (name: string, level: number, index: bigint) => + `snapshot:${name}:node:${level}:${index}:version`; +const nodePreviousValueKey = (name: string, level: number, index: bigint) => + `snapshot:${name}:node:${level}:${index}:value`; +const snapshotMetaKey = (name: string, version: number) => `snapshot:${name}:${version}`; + +/** + * 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) (stores the previous value for each node and at which snapshot it was last modified) + * 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 SnapshotBuilder { + constructor(private db: LevelUp, private tree: TreeBase & AppendOnlyTree, private hasher: Hasher) {} + async getSnapshot(version: number): Promise { + const filledLeavesAtVersion = await this.#getLeafCountAtVersion(version); + + if (typeof filledLeavesAtVersion === 'undefined') { + throw new Error(`Version ${version} does not exist for tree ${this.tree.getName()}`); + } + + return new AppendOnlySnapshot(this.db, version, filledLeavesAtVersion, this.tree, this.hasher); + } + + async snapshot(version: number): Promise { + const leafCountAtVersion = await this.#getLeafCountAtVersion(version); + if (typeof leafCountAtVersion !== 'undefined') { + throw new Error(`Version ${version} of tree ${this.tree.getName()} already exists`); + } + + 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 BF and update latest values + while (queue.length > 0) { + const [node, level, index] = queue.shift()!; + + const previousValue = await this.db.get(nodePreviousValueKey(treeName, level, index)).catch(() => undefined); + if (!previousValue || !node.equals(previousValue)) { + // console.log(`Node at ${level}:${index} has changed`); + batch.put(nodeVersionKey(treeName, level, index), String(version)); + batch.put(nodePreviousValueKey(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; + } + + 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 leafCount = this.tree.getNumLeaves(false); + batch.put(snapshotMetaKey(treeName, version), leafCount); + await batch.write(); + + return new AppendOnlySnapshot(this.db, version, leafCount, this.tree, this.hasher); + } + + async #getLeafCountAtVersion(version: number): Promise { + const filledLeavesAtVersion = await this.db + .get(snapshotMetaKey(this.tree.getName(), version)) + .then(x => BigInt(x.toString())) + .catch(() => undefined); + return filledLeavesAtVersion; + } +} + +/** + * a + */ +class AppendOnlySnapshot implements SiblingPathSource { + constructor( + private db: LevelUp, + private version: number, + private leafCountAtVersion: bigint, + private tree: TreeBase & AppendOnlyTree, + private hasher: Hasher, + ) {} + + public async getSiblingPath(index: bigint, _: boolean): 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.#getHistoricNodeValue(level, siblingIndex); + path.push(sibling); + + level -= 1; + index >>= 1n; + } + + return new SiblingPath(this.tree.getDepth() as N, path); + } + + async #getHistoricNodeValue(level: number, index: bigint): Promise { + const lastNodeVersion = await this.#getNodeVersion(level, index); + + // node has never been set + if (typeof lastNodeVersion === 'undefined') { + // console.log(`node ${level}:${index} not found, returning zero hash`); + return this.tree.getZeroHash(level); + } + + // node was set some time in the past + if (lastNodeVersion <= this.version) { + // console.log(`node ${level}:${index} unchanged ${lastNodeVersion} <= ${this.version}`); + return this.db.get(nodePreviousValueKey(this.tree.getName(), level, index)); + } + + // the node has been modified since this snapshot was taken + // because we're working with an AppendOnly tree, historic 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-1) hashes + const depth = this.tree.getDepth(); + const leafStart = index * 2n ** BigInt(depth - level); + if (leafStart >= this.leafCountAtVersion) { + // console.log(`subtree rooted at ${level}:${index} outside of snapshot, returning zero hash`); + return this.tree.getZeroHash(level); + } + + const [lhs, rhs] = await Promise.all([ + this.#getHistoricNodeValue(level + 1, 2n * index), + this.#getHistoricNodeValue(level + 1, 2n * index + 1n), + ]); + + // console.log(`recreating node ${level}:${index}`); + return this.hasher.hash(lhs, rhs); + } + + async #getNodeVersion(level: number, index: bigint): Promise { + try { + const value: Buffer | string = await this.db.get(nodeVersionKey(this.tree.getName(), level, index)); + return parseInt(value.toString(), 10); + } catch (err) { + return undefined; + } + } +} diff --git a/yarn-project/merkle-tree/src/snapshots/incremental_snapshot.test.ts b/yarn-project/merkle-tree/src/snapshots/incremental_snapshot.test.ts new file mode 100644 index 00000000000..69308e4f022 --- /dev/null +++ b/yarn-project/merkle-tree/src/snapshots/incremental_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 { IncrementalSnapshotBuilder } from './incremental_snapshot.js'; +import { describeSnapshotBuilderTestSuite } from './snapshot_builder_test_suite.js'; + +describe('FullSnapshotBuilder', () => { + let tree: StandardTree; + let snapshotBuilder: IncrementalSnapshotBuilder; + let db: LevelUp; + + beforeEach(async () => { + db = levelup(createMemDown()); + tree = await newTree(StandardTree, db, new Pedersen(), 'test', 4); + snapshotBuilder = new IncrementalSnapshotBuilder(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/incremental_snapshot.ts b/yarn-project/merkle-tree/src/snapshots/incremental_snapshot.ts new file mode 100644 index 00000000000..43e2c7096ae --- /dev/null +++ b/yarn-project/merkle-tree/src/snapshots/incremental_snapshot.ts @@ -0,0 +1,156 @@ +import { SiblingPath } from '@aztec/types'; + +import { LevelUp } from 'levelup'; + +import { SiblingPathSource } from '../interfaces/merkle_tree.js'; +import { TreeBase } from '../tree_base.js'; +import { SnapshotBuilder } from './snapshot_builder.js'; + +const snapshotChildKey = (node: Buffer, child: 0 | 1) => + Buffer.concat([Buffer.from('snapshot'), node, Buffer.from(String(child))]); + +const snapshotRootKey = (treeName: string, version: number) => `snapshot:${treeName}:${version}`; + +/** + * 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 IncrementalSnapshotBuilder implements SnapshotBuilder { + constructor(private db: LevelUp, private tree: TreeBase) {} + + async snapshot(version: number): Promise { + const existingRoot = await this.#getRootAtVersion(version); + + if (existingRoot) { + throw new Error(`Version ${version} for tree ${this.tree.getName()} already exists`); + } + + const batch = this.db.batch(); + const root = this.tree.getRoot(false); + const depth = this.tree.getDepth(); + const queue: [Buffer, number, bigint][] = [[root, 0, 0n]]; + + // walk the tree BF 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 + 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 + 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(), version), root); + await batch.write(); + + return new IncrementalSnapshot(this.db, root, this.tree); + } + + async getSnapshot(version: number): Promise { + const historicRoot = await this.#getRootAtVersion(version); + + if (!historicRoot) { + throw new Error(`Version ${version} does not exist for tree ${this.tree.getName()}`); + } + + return new IncrementalSnapshot(this.db, historicRoot, this.tree); + } + + #getRootAtVersion(version: number): Promise { + return this.db.get(snapshotRootKey(this.tree.getName(), version)).catch(() => undefined); + } +} + +/** + * A source of sibling paths from a snapshot tree + */ +class IncrementalSnapshot implements SiblingPathSource { + constructor(private db: LevelUp, private historicRoot: Buffer, private tree: TreeBase) {} + + async getSiblingPath(index: bigint, _includeUncommitted: boolean): Promise> { + const root = this.historicRoot; + const pathFromRoot = this.#getPathFromRoot(index); + const siblings: Buffer[] = []; + + 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]; + + siblings.push(sibling); + node = next; + } + + // 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); + } + + /** + * 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/snapshot_builder.ts b/yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts new file mode 100644 index 00000000000..9f154358e07 --- /dev/null +++ b/yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts @@ -0,0 +1,18 @@ +import { SiblingPathSource } from '../interfaces/merkle_tree.js'; + +/** + * An interface for a tree that can record snapshots of its contents. + */ +export interface SnapshotBuilder { + /** + * Creates a snapshot of the tree at the given version. + * @param version - The version to snapshot the tree at. + */ + snapshot(version: number): Promise; + + /** + * Returns a snapshot of the tree at the given version. + * @param version - The version of the snapshot to return. + */ + getSnapshot(version: number): Promise; +} 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..7638182e4f1 --- /dev/null +++ b/yarn-project/merkle-tree/src/snapshots/snapshot_builder_test_suite.ts @@ -0,0 +1,138 @@ +import { TreeBase } from '../tree_base.js'; +import { SnapshotBuilder } 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('refuses to take the same snapshot twice', async () => { + await modifyTree(tree); + await tree.commit(); + + const version = 1; + await snapshotBuilder.snapshot(version); + await expect(snapshotBuilder.snapshot(version)).rejects.toThrow(); + }); + + it('returns the same path if tree has not advanced', async () => { + await modifyTree(tree); + await tree.commit(); + const snapshot = await snapshotBuilder.snapshot(1); + + const historicPaths = await Promise.all(leaves.map(leaf => snapshot.getSiblingPath(leaf, false))); + 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, false))); + + for (const [index, path] of historicPaths.entries()) { + expect(path).toEqual(expectedPaths[index]); + } + }); + + it('returns historic paths at old snapshots', async () => { + await modifyTree(tree); + await tree.commit(); + const snapshot = await snapshotBuilder.snapshot(1); + + const historicPaths = await Promise.all(leaves.map(leaf => snapshot.getSiblingPath(leaf, false))); + const expectedPaths = await Promise.all(leaves.map(leaf => tree.getSiblingPath(leaf, false))); + + await modifyTree(tree); + await tree.commit(); + + await snapshotBuilder.snapshot(2); + + 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, false))); + + 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, false))); + + 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(); + }); + }); + }); +} 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. */ From 955b8f2e77c6a7f7f14ec8d93d88c60cf1d3387f Mon Sep 17 00:00:00 2001 From: Alex Gherghisan Date: Wed, 29 Nov 2023 15:11:52 +0000 Subject: [PATCH 2/8] refactor: code review --- yarn-project/merkle-tree/src/index.ts | 2 +- .../src/snapshots/append_only_snapshot.ts | 105 ++++++++++-------- ...snapshot.test.ts => full_snapshot.test.ts} | 6 +- ...cremental_snapshot.ts => full_snapshot.ts} | 39 ++++--- .../src/snapshots/snapshot_builder.ts | 8 +- .../snapshots/snapshot_builder_test_suite.ts | 25 +++-- 6 files changed, 104 insertions(+), 81 deletions(-) rename yarn-project/merkle-tree/src/snapshots/{incremental_snapshot.test.ts => full_snapshot.test.ts} (79%) rename yarn-project/merkle-tree/src/snapshots/{incremental_snapshot.ts => full_snapshot.ts} (76%) diff --git a/yarn-project/merkle-tree/src/index.ts b/yarn-project/merkle-tree/src/index.ts index 9fa9f0e7aba..be11945c0be 100644 --- a/yarn-project/merkle-tree/src/index.ts +++ b/yarn-project/merkle-tree/src/index.ts @@ -10,5 +10,5 @@ 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/incremental_snapshot.js'; +export * from './snapshots/full_snapshot.js'; export * from './snapshots/append_only_snapshot.js'; diff --git a/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts b/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts index 844e3ca9aed..c05241fd524 100644 --- a/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts +++ b/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts @@ -7,11 +7,16 @@ import { SiblingPathSource } from '../interfaces/merkle_tree.js'; import { TreeBase } from '../tree_base.js'; import { SnapshotBuilder } from './snapshot_builder.js'; -const nodeVersionKey = (name: string, level: number, index: bigint) => - `snapshot:${name}:node:${level}:${index}:version`; -const nodePreviousValueKey = (name: string, level: number, index: bigint) => - `snapshot:${name}:node:${level}:${index}:value`; -const snapshotMetaKey = (name: string, version: number) => `snapshot:${name}:${version}`; +// 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 - currently just the count of leaves +const snapshotLeafCountKey = (treeName: string, block: number) => `snapshot:leafCount:${treeName}:${block}`; /** * A more space-efficient way of storing snapshots of AppendOnlyTrees that trades space need for slower @@ -23,27 +28,28 @@ const snapshotMetaKey = (name: string, version: number) => `snapshot:${name}:${v * M - count of snapshots * H - tree height * - * Space complexity: O(N) (stores the previous value for each node and at which snapshot it was last modified) + * 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 SnapshotBuilder { constructor(private db: LevelUp, private tree: TreeBase & AppendOnlyTree, private hasher: Hasher) {} - async getSnapshot(version: number): Promise { - const filledLeavesAtVersion = await this.#getLeafCountAtVersion(version); + async getSnapshot(block: number): Promise { + const leafCount = await this.#getLeafCountAtBlock(block); - if (typeof filledLeavesAtVersion === 'undefined') { - throw new Error(`Version ${version} does not exist for tree ${this.tree.getName()}`); + if (typeof leafCount === 'undefined') { + throw new Error(`Snapshot for tree ${this.tree.getName()} at block ${block} does not exist`); } - return new AppendOnlySnapshot(this.db, version, filledLeavesAtVersion, this.tree, this.hasher); + return new AppendOnlySnapshot(this.db, block, leafCount, this.tree, this.hasher); } - async snapshot(version: number): Promise { - const leafCountAtVersion = await this.#getLeafCountAtVersion(version); - if (typeof leafCountAtVersion !== 'undefined') { - throw new Error(`Version ${version} of tree ${this.tree.getName()} already exists`); + async snapshot(block: number): Promise { + const leafCountAtBlock = await this.#getLeafCountAtBlock(block); + if (typeof leafCountAtBlock !== 'undefined') { + // no-op, we already have a snapshot + return new AppendOnlySnapshot(this.db, block, leafCountAtBlock, this.tree, this.hasher); } const batch = this.db.batch(); @@ -52,15 +58,16 @@ export class AppendOnlySnapshotBuilder implements SnapshotBuilder { const treeName = this.tree.getName(); const queue: [Buffer, number, bigint][] = [[root, 0, 0n]]; - // walk the BF and update latest values + // walk the tree in BF and store latest nodes while (queue.length > 0) { const [node, level, index] = queue.shift()!; - const previousValue = await this.db.get(nodePreviousValueKey(treeName, level, index)).catch(() => undefined); - if (!previousValue || !node.equals(previousValue)) { - // console.log(`Node at ${level}:${index} has changed`); - batch.put(nodeVersionKey(treeName, level, index), String(version)); - batch.put(nodePreviousValueKey(treeName, level, index), node); + 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; @@ -72,6 +79,7 @@ export class AppendOnlySnapshotBuilder implements SnapshotBuilder { 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), @@ -87,18 +95,18 @@ export class AppendOnlySnapshotBuilder implements SnapshotBuilder { } const leafCount = this.tree.getNumLeaves(false); - batch.put(snapshotMetaKey(treeName, version), leafCount); + batch.put(snapshotLeafCountKey(treeName, block), String(leafCount)); await batch.write(); - return new AppendOnlySnapshot(this.db, version, leafCount, this.tree, this.hasher); + return new AppendOnlySnapshot(this.db, block, leafCount, this.tree, this.hasher); } - async #getLeafCountAtVersion(version: number): Promise { - const filledLeavesAtVersion = await this.db - .get(snapshotMetaKey(this.tree.getName(), version)) + async #getLeafCountAtBlock(block: number): Promise { + const leafCount = await this.db + .get(snapshotLeafCountKey(this.tree.getName(), block)) .then(x => BigInt(x.toString())) .catch(() => undefined); - return filledLeavesAtVersion; + return leafCount; } } @@ -108,13 +116,13 @@ export class AppendOnlySnapshotBuilder implements SnapshotBuilder { class AppendOnlySnapshot implements SiblingPathSource { constructor( private db: LevelUp, - private version: number, - private leafCountAtVersion: bigint, + private block: number, + private leafCount: bigint, private tree: TreeBase & AppendOnlyTree, private hasher: Hasher, ) {} - public async getSiblingPath(index: bigint, _: boolean): Promise> { + public async getSiblingPath(index: bigint): Promise> { const path: Buffer[] = []; const depth = this.tree.getDepth(); let level = depth; @@ -123,54 +131,55 @@ class AppendOnlySnapshot implements SiblingPathSource { const isRight = index & 0x01n; const siblingIndex = isRight ? index - 1n : index + 1n; - const sibling = await this.#getHistoricNodeValue(level, siblingIndex); + const sibling = await this.#getHistoricalNodeValue(level, siblingIndex); path.push(sibling); level -= 1; index >>= 1n; } - return new SiblingPath(this.tree.getDepth() as N, path); + return new SiblingPath(depth as N, path); } - async #getHistoricNodeValue(level: number, index: bigint): Promise { - const lastNodeVersion = await this.#getNodeVersion(level, index); + async #getHistoricalNodeValue(level: number, index: bigint): Promise { + const blockNumber = await this.#getBlockNumberThatModifiedNode(level, index); // node has never been set - if (typeof lastNodeVersion === 'undefined') { - // console.log(`node ${level}:${index} not found, returning zero hash`); + if (typeof blockNumber === 'undefined') { return this.tree.getZeroHash(level); } // node was set some time in the past - if (lastNodeVersion <= this.version) { - // console.log(`node ${level}:${index} unchanged ${lastNodeVersion} <= ${this.version}`); - return this.db.get(nodePreviousValueKey(this.tree.getName(), level, index)); + 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, historic leaves never change + // 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-1) hashes + // 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.leafCountAtVersion) { - // console.log(`subtree rooted at ${level}:${index} outside of snapshot, returning zero hash`); + if (leafStart >= this.leafCount) { return this.tree.getZeroHash(level); } const [lhs, rhs] = await Promise.all([ - this.#getHistoricNodeValue(level + 1, 2n * index), - this.#getHistoricNodeValue(level + 1, 2n * index + 1n), + this.#getHistoricalNodeValue(level + 1, 2n * index), + this.#getHistoricalNodeValue(level + 1, 2n * index + 1n), ]); - // console.log(`recreating node ${level}:${index}`); return this.hasher.hash(lhs, rhs); } - async #getNodeVersion(level: number, index: bigint): Promise { + async #getBlockNumberThatModifiedNode(level: number, index: bigint): Promise { try { - const value: Buffer | string = await this.db.get(nodeVersionKey(this.tree.getName(), level, index)); + 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/incremental_snapshot.test.ts b/yarn-project/merkle-tree/src/snapshots/full_snapshot.test.ts similarity index 79% rename from yarn-project/merkle-tree/src/snapshots/incremental_snapshot.test.ts rename to yarn-project/merkle-tree/src/snapshots/full_snapshot.test.ts index 69308e4f022..25d0680afbb 100644 --- a/yarn-project/merkle-tree/src/snapshots/incremental_snapshot.test.ts +++ b/yarn-project/merkle-tree/src/snapshots/full_snapshot.test.ts @@ -2,18 +2,18 @@ import levelup, { LevelUp } from 'levelup'; import { Pedersen, StandardTree, newTree } from '../index.js'; import { createMemDown } from '../test/utils/create_mem_down.js'; -import { IncrementalSnapshotBuilder } from './incremental_snapshot.js'; +import { FullSnapshotBuilder } from './full_snapshot.js'; import { describeSnapshotBuilderTestSuite } from './snapshot_builder_test_suite.js'; describe('FullSnapshotBuilder', () => { let tree: StandardTree; - let snapshotBuilder: IncrementalSnapshotBuilder; + let snapshotBuilder: FullSnapshotBuilder; let db: LevelUp; beforeEach(async () => { db = levelup(createMemDown()); tree = await newTree(StandardTree, db, new Pedersen(), 'test', 4); - snapshotBuilder = new IncrementalSnapshotBuilder(db, tree); + snapshotBuilder = new FullSnapshotBuilder(db, tree); }); describeSnapshotBuilderTestSuite( diff --git a/yarn-project/merkle-tree/src/snapshots/incremental_snapshot.ts b/yarn-project/merkle-tree/src/snapshots/full_snapshot.ts similarity index 76% rename from yarn-project/merkle-tree/src/snapshots/incremental_snapshot.ts rename to yarn-project/merkle-tree/src/snapshots/full_snapshot.ts index 43e2c7096ae..52af1ed18d6 100644 --- a/yarn-project/merkle-tree/src/snapshots/incremental_snapshot.ts +++ b/yarn-project/merkle-tree/src/snapshots/full_snapshot.ts @@ -6,10 +6,12 @@ import { SiblingPathSource } from '../interfaces/merkle_tree.js'; import { TreeBase } from '../tree_base.js'; import { SnapshotBuilder } from './snapshot_builder.js'; +// ket for a node's children const snapshotChildKey = (node: Buffer, child: 0 | 1) => - Buffer.concat([Buffer.from('snapshot'), node, Buffer.from(String(child))]); + Buffer.concat([Buffer.from('snapshot:node:'), node, Buffer.from(':' + child)]); -const snapshotRootKey = (treeName: string, version: number) => `snapshot:${treeName}:${version}`; +// metadata for a snapshot - the root of the historical tree +const snapshotRootKey = (treeName: string, block: number) => `snapshot:root:${treeName}:${block}`; /** * Builds a full snapshot of a tree. This implementation works for any Merkle tree and stores @@ -25,14 +27,14 @@ const snapshotRootKey = (treeName: string, version: number) => `snapshot:${treeN * Worst case space complexity: O(N * M) * Sibling path access: O(H) database reads */ -export class IncrementalSnapshotBuilder implements SnapshotBuilder { +export class FullSnapshotBuilder implements SnapshotBuilder { constructor(private db: LevelUp, private tree: TreeBase) {} - async snapshot(version: number): Promise { - const existingRoot = await this.#getRootAtVersion(version); + async snapshot(block: number): Promise { + const historicalRoot = await this.#getRootAtBlock(block); - if (existingRoot) { - throw new Error(`Version ${version} for tree ${this.tree.getName()} already exists`); + if (historicalRoot) { + return new FullSnapshot(this.db, historicalRoot, this.tree); } const batch = this.db.batch(); @@ -40,7 +42,7 @@ export class IncrementalSnapshotBuilder implements SnapshotBuilder { const depth = this.tree.getDepth(); const queue: [Buffer, number, bigint][] = [[root, 0, 0n]]; - // walk the tree BF and store each of its nodes in the database + // walk the tree breadth-first and store each of its nodes in the database // for each node we save two keys // :0 -> // :1 -> @@ -49,6 +51,7 @@ export class IncrementalSnapshotBuilder implements SnapshotBuilder { // 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; @@ -81,34 +84,38 @@ export class IncrementalSnapshotBuilder implements SnapshotBuilder { } } - batch.put(snapshotRootKey(this.tree.getName(), version), root); + batch.put(snapshotRootKey(this.tree.getName(), block), root); await batch.write(); - return new IncrementalSnapshot(this.db, root, this.tree); + return new FullSnapshot(this.db, root, this.tree); } async getSnapshot(version: number): Promise { - const historicRoot = await this.#getRootAtVersion(version); + const historicRoot = await this.#getRootAtBlock(version); if (!historicRoot) { throw new Error(`Version ${version} does not exist for tree ${this.tree.getName()}`); } - return new IncrementalSnapshot(this.db, historicRoot, this.tree); + return new FullSnapshot(this.db, historicRoot, this.tree); } - #getRootAtVersion(version: number): Promise { - return this.db.get(snapshotRootKey(this.tree.getName(), version)).catch(() => undefined); + async #getRootAtBlock(version: number): Promise { + try { + return await this.db.get(snapshotRootKey(this.tree.getName(), version)); + } catch (err) { + return undefined; + } } } /** * A source of sibling paths from a snapshot tree */ -class IncrementalSnapshot implements SiblingPathSource { +class FullSnapshot implements SiblingPathSource { constructor(private db: LevelUp, private historicRoot: Buffer, private tree: TreeBase) {} - async getSiblingPath(index: bigint, _includeUncommitted: boolean): Promise> { + async getSiblingPath(index: bigint): Promise> { const root = this.historicRoot; const pathFromRoot = this.#getPathFromRoot(index); const siblings: Buffer[] = []; diff --git a/yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts b/yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts index 9f154358e07..1fcde85e467 100644 --- a/yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts +++ b/yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts @@ -6,13 +6,13 @@ import { SiblingPathSource } from '../interfaces/merkle_tree.js'; export interface SnapshotBuilder { /** * Creates a snapshot of the tree at the given version. - * @param version - The version to snapshot the tree at. + * @param block - The version to snapshot the tree at. */ - snapshot(version: number): Promise; + snapshot(block: number): Promise; /** * Returns a snapshot of the tree at the given version. - * @param version - The version of the snapshot to return. + * @param block - The version of the snapshot to return. */ - getSnapshot(version: number): Promise; + getSnapshot(block: number): Promise; } 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 index 7638182e4f1..c4f2b2891ce 100644 --- a/yarn-project/merkle-tree/src/snapshots/snapshot_builder_test_suite.ts +++ b/yarn-project/merkle-tree/src/snapshots/snapshot_builder_test_suite.ts @@ -26,16 +26,16 @@ export function describeSnapshotBuilderTestSuite { + it('is idempotent', async () => { await modifyTree(tree); await tree.commit(); - const version = 1; - await snapshotBuilder.snapshot(version); - await expect(snapshotBuilder.snapshot(version)).rejects.toThrow(); + 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 advanced', async () => { + it('returns the same path if tree has not diverged', async () => { await modifyTree(tree); await tree.commit(); const snapshot = await snapshotBuilder.snapshot(1); @@ -65,25 +65,28 @@ export function describeSnapshotBuilderTestSuite { + it('retains old snapshots even if new one are created', async () => { await modifyTree(tree); await tree.commit(); - const snapshot = await snapshotBuilder.snapshot(1); - const historicPaths = await Promise.all(leaves.map(leaf => snapshot.getSiblingPath(leaf, false))); 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, false))); + for (const [index, path] of historicPaths.entries()) { expect(path).toEqual(expectedPaths[index]); } }); - it('retains old snapshots even if new one are created', async () => { + it('retains old snapshots even if new one are created and the tree diverges', async () => { await modifyTree(tree); await tree.commit(); @@ -96,7 +99,11 @@ export function describeSnapshotBuilderTestSuite snapshot.getSiblingPath(leaf, false))); for (const [index, path] of historicPaths.entries()) { From 51c52ddcee0c5bae9fd196fa8d8425621ba5db8e Mon Sep 17 00:00:00 2001 From: Alex Gherghisan Date: Wed, 29 Nov 2023 19:53:06 +0000 Subject: [PATCH 3/8] feat: add indexed tree snapshots --- yarn-project/merkle-tree/src/index.ts | 3 +- .../src/snapshots/append_only_snapshot.ts | 4 +- .../src/snapshots/base_full_snapshot.ts | 182 ++++++++++++++++++ .../src/snapshots/full_snapshot.test.ts | 6 +- .../src/snapshots/full_snapshot.ts | 152 +-------------- .../snapshots/indexed_tree_snapshot.test.ts | 82 ++++++++ .../src/snapshots/indexed_tree_snapshot.ts | 48 +++++ .../src/snapshots/snapshot_builder.ts | 21 +- .../snapshots/snapshot_builder_test_suite.ts | 4 +- .../standard_indexed_tree.ts | 5 +- 10 files changed, 349 insertions(+), 158 deletions(-) create mode 100644 yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts create mode 100644 yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.test.ts create mode 100644 yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.ts diff --git a/yarn-project/merkle-tree/src/index.ts b/yarn-project/merkle-tree/src/index.ts index be11945c0be..5181cecfc15 100644 --- a/yarn-project/merkle-tree/src/index.ts +++ b/yarn-project/merkle-tree/src/index.ts @@ -4,7 +4,7 @@ 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'; @@ -12,3 +12,4 @@ 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/snapshots/append_only_snapshot.ts b/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts index c05241fd524..b42250eaa0e 100644 --- a/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts +++ b/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts @@ -5,7 +5,7 @@ import { LevelUp } from 'levelup'; import { AppendOnlyTree } from '../interfaces/append_only_tree.js'; import { SiblingPathSource } from '../interfaces/merkle_tree.js'; import { TreeBase } from '../tree_base.js'; -import { SnapshotBuilder } from './snapshot_builder.js'; +import { TreeSnapshotBuilder } from './snapshot_builder.js'; // stores the last block that modified this node const nodeModifiedAtBlockKey = (treeName: string, level: number, index: bigint) => @@ -33,7 +33,7 @@ const snapshotLeafCountKey = (treeName: string, block: number) => `snapshot:leaf * Best case: O(H) database reads + O(1) hashes * Worst case: O(H) database reads + O(H) hashes */ -export class AppendOnlySnapshotBuilder implements SnapshotBuilder { +export class AppendOnlySnapshotBuilder implements TreeSnapshotBuilder { constructor(private db: LevelUp, private tree: TreeBase & AppendOnlyTree, private hasher: Hasher) {} async getSnapshot(block: number): Promise { const leafCount = await this.#getLeafCountAtBlock(block); 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..3f93daa9609 --- /dev/null +++ b/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts @@ -0,0 +1,182 @@ +import { SiblingPath } from '@aztec/types'; + +import { LevelUp, LevelUpChain } from 'levelup'; + +import { SiblingPathSource } from '../interfaces/merkle_tree.js'; +import { TreeBase } from '../tree_base.js'; +import { TreeSnapshot, TreeSnapshotBuilder } from './snapshot_builder.js'; + +// ket 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 - the root of the historical tree +const snapshotRootKey = (treeName: string, block: number) => `snapshot:root:${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 historicalRoot = await this.#getRootAtBlock(block); + + if (historicalRoot) { + return this.openSnapshot(historicalRoot); + } + + const batch = this.db.batch(); + const root = this.tree.getRoot(); + 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); + await batch.write(); + + return this.openSnapshot(root); + } + + protected handleLeaf(_index: bigint, _node: Buffer, _batch: LevelUpChain) { + return; + } + + async getSnapshot(version: number): Promise { + const historicRoot = await this.#getRootAtBlock(version); + + if (!historicRoot) { + throw new Error(`Version ${version} does not exist for tree ${this.tree.getName()}`); + } + + return this.openSnapshot(historicRoot); + } + + protected abstract openSnapshot(root: Buffer): S; + + async #getRootAtBlock(block: number): Promise { + try { + return await this.db.get(snapshotRootKey(this.tree.getName(), block)); + } 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 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); + } + + 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 index 25d0680afbb..3f2cc2af791 100644 --- a/yarn-project/merkle-tree/src/snapshots/full_snapshot.test.ts +++ b/yarn-project/merkle-tree/src/snapshots/full_snapshot.test.ts @@ -2,18 +2,18 @@ import levelup, { LevelUp } from 'levelup'; import { Pedersen, StandardTree, newTree } from '../index.js'; import { createMemDown } from '../test/utils/create_mem_down.js'; -import { FullSnapshotBuilder } from './full_snapshot.js'; +import { FullTreeSnapshotBuilder } from './full_snapshot.js'; import { describeSnapshotBuilderTestSuite } from './snapshot_builder_test_suite.js'; describe('FullSnapshotBuilder', () => { let tree: StandardTree; - let snapshotBuilder: FullSnapshotBuilder; + let snapshotBuilder: FullTreeSnapshotBuilder; let db: LevelUp; beforeEach(async () => { db = levelup(createMemDown()); tree = await newTree(StandardTree, db, new Pedersen(), 'test', 4); - snapshotBuilder = new FullSnapshotBuilder(db, tree); + snapshotBuilder = new FullTreeSnapshotBuilder(db, tree); }); describeSnapshotBuilderTestSuite( diff --git a/yarn-project/merkle-tree/src/snapshots/full_snapshot.ts b/yarn-project/merkle-tree/src/snapshots/full_snapshot.ts index 52af1ed18d6..83772a4368b 100644 --- a/yarn-project/merkle-tree/src/snapshots/full_snapshot.ts +++ b/yarn-project/merkle-tree/src/snapshots/full_snapshot.ts @@ -1,17 +1,7 @@ -import { SiblingPath } from '@aztec/types'; - -import { LevelUp } from 'levelup'; - import { SiblingPathSource } from '../interfaces/merkle_tree.js'; import { TreeBase } from '../tree_base.js'; -import { SnapshotBuilder } from './snapshot_builder.js'; - -// ket 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 - the root of the historical tree -const snapshotRootKey = (treeName: string, block: number) => `snapshot:root:${treeName}:${block}`; +import { BaseFullTreeSnapshot, BaseFullTreeSnapshotBuilder } from './base_full_snapshot.js'; +import { TreeSnapshotBuilder } from './snapshot_builder.js'; /** * Builds a full snapshot of a tree. This implementation works for any Merkle tree and stores @@ -27,137 +17,11 @@ const snapshotRootKey = (treeName: string, block: number) => `snapshot:root:${tr * Worst case space complexity: O(N * M) * Sibling path access: O(H) database reads */ -export class FullSnapshotBuilder implements SnapshotBuilder { - constructor(private db: LevelUp, private tree: TreeBase) {} - - async snapshot(block: number): Promise { - const historicalRoot = await this.#getRootAtBlock(block); - - if (historicalRoot) { - return new FullSnapshot(this.db, historicalRoot, this.tree); - } - - const batch = this.db.batch(); - const root = this.tree.getRoot(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 - 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); - await batch.write(); - - return new FullSnapshot(this.db, root, this.tree); - } - - async getSnapshot(version: number): Promise { - const historicRoot = await this.#getRootAtBlock(version); - - if (!historicRoot) { - throw new Error(`Version ${version} does not exist for tree ${this.tree.getName()}`); - } - - return new FullSnapshot(this.db, historicRoot, this.tree); - } - - async #getRootAtBlock(version: number): Promise { - try { - return await this.db.get(snapshotRootKey(this.tree.getName(), version)); - } catch (err) { - return undefined; - } - } -} - -/** - * A source of sibling paths from a snapshot tree - */ -class FullSnapshot implements SiblingPathSource { - constructor(private db: LevelUp, private historicRoot: Buffer, private tree: TreeBase) {} - - async getSiblingPath(index: bigint): Promise> { - const root = this.historicRoot; - const pathFromRoot = this.#getPathFromRoot(index); - const siblings: Buffer[] = []; - - 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]; - - siblings.push(sibling); - node = next; - } - - // 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); - } - - /** - * 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; +export class FullTreeSnapshotBuilder + extends BaseFullTreeSnapshotBuilder + implements TreeSnapshotBuilder +{ + protected openSnapshot(root: Buffer): SiblingPathSource { + return new BaseFullTreeSnapshot(this.db, root, 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..0ddb0ca7303 --- /dev/null +++ b/yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.test.ts @@ -0,0 +1,82 @@ +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); + }); + }); +}); 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..81151863cd4 --- /dev/null +++ b/yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.ts @@ -0,0 +1,48 @@ +import { LevelUp, LevelUpChain } from 'levelup'; + +import { IndexedTree, LeafData } 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): IndexedTreeSnapshot { + return new IndexedTreeSnapshotImpl(this.db, root, 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 getLatestLeafDataCopy(index: bigint): Promise { + let leafNode: Buffer; + for await (const [node, _sibling] of this.pathFromRootToLeaf(index)) { + leafNode = node; + } + + const leafValue = await this.db.get(snapshotLeafValue(leafNode!, index)).catch(() => undefined); + if (leafValue) { + return decodeTreeValue(leafValue); + } else { + return undefined; + } + } +} diff --git a/yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts b/yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts index 1fcde85e467..6735175bc64 100644 --- a/yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts +++ b/yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts @@ -1,18 +1,33 @@ +import { LeafData } from '../interfaces/indexed_tree.js'; import { SiblingPathSource } from '../interfaces/merkle_tree.js'; /** * An interface for a tree that can record snapshots of its contents. */ -export interface SnapshotBuilder { +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; + 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; + getSnapshot(block: number): Promise; +} + +/** + * A tree snapshot + */ +export interface TreeSnapshot extends SiblingPathSource {} + +/** A snapshot of an indexed tree */ +export interface IndexedTreeSnapshot extends SiblingPathSource { + /** + * Gets the historical data for a leaf + * @param index - The index of the leaf to get the data for + */ + getLatestLeafDataCopy(index: bigint): Promise; } 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 index c4f2b2891ce..f28bfea803a 100644 --- a/yarn-project/merkle-tree/src/snapshots/snapshot_builder_test_suite.ts +++ b/yarn-project/merkle-tree/src/snapshots/snapshot_builder_test_suite.ts @@ -1,8 +1,8 @@ import { TreeBase } from '../tree_base.js'; -import { SnapshotBuilder } from './snapshot_builder.js'; +import { TreeSnapshotBuilder } from './snapshot_builder.js'; /** Creates a test suit for snapshots */ -export function describeSnapshotBuilderTestSuite( +export function describeSnapshotBuilderTestSuite( getTree: () => T, getSnapshotBuilder: () => S, modifyTree: (tree: T) => Promise, 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 29c2bac627c..c8eca1e3181 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 @@ -53,15 +53,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)); From c6e508f2bf64b675dd72585ff8275e52a174bcfd Mon Sep 17 00:00:00 2001 From: Alex Gherghisan Date: Thu, 30 Nov 2023 09:56:03 +0000 Subject: [PATCH 4/8] fix: ts compile --- yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts | 2 +- .../merkle-tree/src/snapshots/indexed_tree_snapshot.ts | 4 +++- yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts | 3 ++- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts b/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts index 3f93daa9609..f4a007e9b86 100644 --- a/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts +++ b/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts @@ -42,7 +42,7 @@ export abstract class BaseFullTreeSnapshotBuilder Date: Thu, 30 Nov 2023 12:41:24 +0000 Subject: [PATCH 5/8] feat: add operations facade over merkle tree snapshots --- .../src/interfaces/append_only_tree.ts | 3 +- .../src/interfaces/update_only_tree.ts | 3 +- .../snapshots/append_only_snapshot.test.ts | 19 --- .../src/snapshots/append_only_snapshot.ts | 39 ++++- .../src/snapshots/base_full_snapshot.ts | 65 +++++++-- .../src/snapshots/full_snapshot.ts | 11 +- .../src/snapshots/indexed_tree_snapshot.ts | 56 ++++++- .../src/snapshots/snapshot_builder.ts | 51 ++++++- .../snapshots/snapshot_builder_test_suite.ts | 10 +- .../src/sparse_tree/sparse_tree.ts | 12 ++ .../standard_indexed_tree.ts | 12 ++ .../src/standard_tree/standard_tree.ts | 11 ++ .../merkle_tree_snapshot_operations_facade.ts | 137 ++++++++++++++++++ .../server_world_state_synchronizer.ts | 5 + .../synchronizer/world_state_synchronizer.ts | 7 + .../src/world-state-db/merkle_tree_db.ts | 10 +- .../src/world-state-db/merkle_trees.ts | 12 ++ 17 files changed, 399 insertions(+), 64 deletions(-) create mode 100644 yarn-project/world-state/src/merkle-tree/merkle_tree_snapshot_operations_facade.ts 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 index 6b2cffe6e09..b66eb2af22b 100644 --- a/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.test.ts +++ b/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.test.ts @@ -17,25 +17,6 @@ describe('AppendOnlySnapshot', () => { snapshotBuilder = new AppendOnlySnapshotBuilder(db, tree, hasher); }); - it('takes snapshots', async () => { - await tree.appendLeaves([Buffer.from('a'), Buffer.from('b'), Buffer.from('c')]); - await tree.commit(); - - const expectedPathAtSnapshot1 = await tree.getSiblingPath(1n, false); - - const snapshot1 = await snapshotBuilder.snapshot(1); - - await tree.appendLeaves([Buffer.from('d'), Buffer.from('e'), Buffer.from('f')]); - await tree.commit(); - - const expectedPathAtSnapshot2 = await tree.getSiblingPath(1n, false); - - const snapshot2 = await snapshotBuilder.snapshot(2); - - await expect(snapshot1.getSiblingPath(1n, false)).resolves.toEqual(expectedPathAtSnapshot1); - await expect(snapshot2.getSiblingPath(1n, false)).resolves.toEqual(expectedPathAtSnapshot2); - }); - describeSnapshotBuilderTestSuite( () => tree, () => snapshotBuilder, diff --git a/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts b/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts index b42250eaa0e..bce6675e86c 100644 --- a/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts +++ b/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts @@ -3,9 +3,8 @@ import { Hasher, SiblingPath } from '@aztec/types'; import { LevelUp } from 'levelup'; import { AppendOnlyTree } from '../interfaces/append_only_tree.js'; -import { SiblingPathSource } from '../interfaces/merkle_tree.js'; import { TreeBase } from '../tree_base.js'; -import { TreeSnapshotBuilder } from './snapshot_builder.js'; +import { TreeSnapshot, TreeSnapshotBuilder } from './snapshot_builder.js'; // stores the last block that modified this node const nodeModifiedAtBlockKey = (treeName: string, level: number, index: bigint) => @@ -35,7 +34,7 @@ const snapshotLeafCountKey = (treeName: string, block: number) => `snapshot:leaf */ export class AppendOnlySnapshotBuilder implements TreeSnapshotBuilder { constructor(private db: LevelUp, private tree: TreeBase & AppendOnlyTree, private hasher: Hasher) {} - async getSnapshot(block: number): Promise { + async getSnapshot(block: number): Promise { const leafCount = await this.#getLeafCountAtBlock(block); if (typeof leafCount === 'undefined') { @@ -45,7 +44,7 @@ export class AppendOnlySnapshotBuilder implements TreeSnapshotBuilder { return new AppendOnlySnapshot(this.db, block, leafCount, this.tree, this.hasher); } - async snapshot(block: number): Promise { + async snapshot(block: number): Promise { const leafCountAtBlock = await this.#getLeafCountAtBlock(block); if (typeof leafCountAtBlock !== 'undefined') { // no-op, we already have a snapshot @@ -113,7 +112,7 @@ export class AppendOnlySnapshotBuilder implements TreeSnapshotBuilder { /** * a */ -class AppendOnlySnapshot implements SiblingPathSource { +class AppendOnlySnapshot implements TreeSnapshot { constructor( private db: LevelUp, private block: number, @@ -141,6 +140,36 @@ class AppendOnlySnapshot implements SiblingPathSource { return new SiblingPath(depth as N, path); } + getDepth(): number { + return this.tree.getDepth(); + } + + getNumLeaves(): bigint { + return this.leafCount; + } + + getRoot(): Buffer { + return this.tree.getRoot(false); + } + + 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); diff --git a/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts b/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts index f4a007e9b86..2a838cea41a 100644 --- a/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts +++ b/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts @@ -2,7 +2,6 @@ import { SiblingPath } from '@aztec/types'; import { LevelUp, LevelUpChain } from 'levelup'; -import { SiblingPathSource } from '../interfaces/merkle_tree.js'; import { TreeBase } from '../tree_base.js'; import { TreeSnapshot, TreeSnapshotBuilder } from './snapshot_builder.js'; @@ -12,6 +11,7 @@ const snapshotChildKey = (node: Buffer, child: 0 | 1) => // metadata for a snapshot - the root of the historical tree const snapshotRootKey = (treeName: string, block: number) => `snapshot:root:${treeName}:${block}`; +const snapshotNumLeavesKey = (treeName: string, block: number) => `snapshot:leafCount:${treeName}:${block}`; /** * Builds a full snapshot of a tree. This implementation works for any Merkle tree and stores @@ -29,20 +29,21 @@ const snapshotRootKey = (treeName: string, block: number) => `snapshot:root:${tr * Worst case space complexity: O(N * M) * Sibling path access: O(H) database reads */ -export abstract class BaseFullTreeSnapshotBuilder +export abstract class BaseFullTreeSnapshotBuilder implements TreeSnapshotBuilder { constructor(protected db: LevelUp, protected tree: T) {} async snapshot(block: number): Promise { - const historicalRoot = await this.#getRootAtBlock(block); + const snapshotMetadata = await this.#getSnapshotMeta(block); - if (historicalRoot) { - return this.openSnapshot(historicalRoot); + 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]]; @@ -90,9 +91,10 @@ export abstract class BaseFullTreeSnapshotBuilder { - const historicRoot = await this.#getRootAtBlock(version); + const snapshotMetadata = await this.#getSnapshotMeta(version); - if (!historicRoot) { + if (!snapshotMetadata) { throw new Error(`Version ${version} does not exist for tree ${this.tree.getName()}`); } - return this.openSnapshot(historicRoot); + return this.openSnapshot(snapshotMetadata.root, snapshotMetadata.numLeaves); } - protected abstract openSnapshot(root: Buffer): S; + protected abstract openSnapshot(root: Buffer, numLeaves: bigint): S; - async #getRootAtBlock(block: number): Promise { + 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 { - return await this.db.get(snapshotRootKey(this.tree.getName(), block)); + 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; } @@ -124,7 +137,12 @@ export abstract class BaseFullTreeSnapshotBuilder(index: bigint): Promise> { const siblings: Buffer[] = []; @@ -140,6 +158,27 @@ export class BaseFullTreeSnapshot implements TreeSnapshot { 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); diff --git a/yarn-project/merkle-tree/src/snapshots/full_snapshot.ts b/yarn-project/merkle-tree/src/snapshots/full_snapshot.ts index 83772a4368b..c78d0ebb188 100644 --- a/yarn-project/merkle-tree/src/snapshots/full_snapshot.ts +++ b/yarn-project/merkle-tree/src/snapshots/full_snapshot.ts @@ -1,7 +1,6 @@ -import { SiblingPathSource } from '../interfaces/merkle_tree.js'; import { TreeBase } from '../tree_base.js'; import { BaseFullTreeSnapshot, BaseFullTreeSnapshotBuilder } from './base_full_snapshot.js'; -import { TreeSnapshotBuilder } from './snapshot_builder.js'; +import { TreeSnapshot, TreeSnapshotBuilder } from './snapshot_builder.js'; /** * Builds a full snapshot of a tree. This implementation works for any Merkle tree and stores @@ -18,10 +17,10 @@ import { TreeSnapshotBuilder } from './snapshot_builder.js'; * Sibling path access: O(H) database reads */ export class FullTreeSnapshotBuilder - extends BaseFullTreeSnapshotBuilder - implements TreeSnapshotBuilder + extends BaseFullTreeSnapshotBuilder + implements TreeSnapshotBuilder { - protected openSnapshot(root: Buffer): SiblingPathSource { - return new BaseFullTreeSnapshot(this.db, root, this.tree); + 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.ts b/yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.ts index f3c21b2d96c..6725bd394e5 100644 --- a/yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.ts +++ b/yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.ts @@ -1,3 +1,4 @@ +import { toBufferBE } from '@aztec/foundation/bigint-buffer'; import { LeafData } from '@aztec/types'; import { LevelUp, LevelUpChain } from 'levelup'; @@ -20,8 +21,8 @@ export class IndexedTreeSnapshotBuilder super(db, tree); } - protected openSnapshot(root: Buffer): IndexedTreeSnapshot { - return new IndexedTreeSnapshotImpl(this.db, root, this.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) { @@ -34,12 +35,13 @@ export class IndexedTreeSnapshotBuilder /** A snapshot of an indexed tree at a particular point in time */ class IndexedTreeSnapshotImpl extends BaseFullTreeSnapshot implements IndexedTreeSnapshot { - async getLatestLeafDataCopy(index: bigint): Promise { - let leafNode: Buffer; - for await (const [node, _sibling] of this.pathFromRootToLeaf(index)) { - leafNode = node; - } + 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); @@ -47,4 +49,44 @@ class IndexedTreeSnapshotImpl extends BaseFullTreeSnapshot implements IndexedTre 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 index af1299be68c..a6722306301 100644 --- a/yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts +++ b/yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts @@ -1,6 +1,4 @@ -import { LeafData } from '@aztec/types'; - -import { SiblingPathSource } from '../interfaces/merkle_tree.js'; +import { LeafData, SiblingPath } from '@aztec/types'; /** * An interface for a tree that can record snapshots of its contents. @@ -22,13 +20,56 @@ export interface TreeSnapshotBuilder { /** * A tree snapshot */ -export interface TreeSnapshot extends SiblingPathSource {} +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 SiblingPathSource { +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 index f28bfea803a..5a3e611c280 100644 --- a/yarn-project/merkle-tree/src/snapshots/snapshot_builder_test_suite.ts +++ b/yarn-project/merkle-tree/src/snapshots/snapshot_builder_test_suite.ts @@ -40,7 +40,7 @@ export function describeSnapshotBuilderTestSuite snapshot.getSiblingPath(leaf, false))); + 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()) { @@ -58,7 +58,7 @@ export function describeSnapshotBuilderTestSuite snapshot.getSiblingPath(leaf, false))); + const historicPaths = await Promise.all(leaves.map(leaf => snapshot.getSiblingPath(leaf))); for (const [index, path] of historicPaths.entries()) { expect(path).toEqual(expectedPaths[index]); @@ -79,7 +79,7 @@ export function describeSnapshotBuilderTestSuite snapshot.getSiblingPath(leaf, false))); + const historicPaths = await Promise.all(leaves.map(leaf => snapshot.getSiblingPath(leaf))); for (const [index, path] of historicPaths.entries()) { expect(path).toEqual(expectedPaths[index]); @@ -104,7 +104,7 @@ export function describeSnapshotBuilderTestSuite snapshot.getSiblingPath(leaf, false))); + const historicPaths = await Promise.all(leaves.map(leaf => snapshot.getSiblingPath(leaf))); for (const [index, path] of historicPaths.entries()) { expect(path).toEqual(expectedPaths[index]); @@ -126,7 +126,7 @@ export function describeSnapshotBuilderTestSuite firstSnapshot.getSiblingPath(leaf, false))); + const historicPaths = await Promise.all(leaves.map(leaf => firstSnapshot.getSiblingPath(leaf))); for (const [index, path] of historicPaths.entries()) { expect(path).toEqual(expectedPaths[index]); 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 e3ff2c0ec24..7ce50316a2c 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 @@ -3,6 +3,8 @@ import { createDebugLogger } from '@aztec/foundation/log'; import { LeafData, SiblingPath } from '@aztec/types'; import { 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'); @@ -75,6 +77,8 @@ export 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 } = {}; @@ -551,6 +555,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/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..fb4855204e3 --- /dev/null +++ b/yarn-project/world-state/src/merkle-tree/merkle_tree_snapshot_operations_facade.ts @@ -0,0 +1,137 @@ +import { Fr } from '@aztec/circuits.js'; +import { IndexedTreeSnapshot, LowLeafWitnessData, 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<[LowLeafWitnessData[], SiblingPath] | [undefined, SiblingPath]> { + 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 04f4d749fee..f754822e91f 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 { LowLeafWitnessData } from '@aztec/merkle-tree'; +import { IndexedTreeSnapshot, LowLeafWitnessData, 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 cdf1aa5c9bd..92558f2ec83 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 @@ -524,6 +524,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. @@ -602,6 +612,8 @@ export class MerkleTrees implements MerkleTreeDb { } } + await this._snapshot(l2Block.number); + return { isBlockOurs: ourBlock }; } } From 317b5d009cf2259acf6716d55da93d17c7e9670b Mon Sep 17 00:00:00 2001 From: Alex Gherghisan Date: Thu, 30 Nov 2023 15:50:30 +0000 Subject: [PATCH 6/8] test: add tests for new snapshot API --- .../src/snapshots/append_only_snapshot.ts | 51 +++++++++++------- .../src/snapshots/base_full_snapshot.ts | 4 +- .../snapshots/indexed_tree_snapshot.test.ts | 14 +++++ .../snapshots/snapshot_builder_test_suite.ts | 52 +++++++++++++++++++ 4 files changed, 101 insertions(+), 20 deletions(-) diff --git a/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts b/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts index bce6675e86c..b530e981b27 100644 --- a/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts +++ b/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts @@ -14,8 +14,9 @@ const nodeModifiedAtBlockKey = (treeName: string, level: number, index: bigint) const historicalNodeKey = (treeName: string, level: number, index: bigint) => `snapshot:node:${treeName}:${level}:${index}:value`; -// metadata for a snapshot - currently just the count of leaves -const snapshotLeafCountKey = (treeName: string, block: number) => `snapshot:leafCount:${treeName}:${block}`; +// 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 @@ -35,20 +36,20 @@ const snapshotLeafCountKey = (treeName: string, block: number) => `snapshot:leaf export class AppendOnlySnapshotBuilder implements TreeSnapshotBuilder { constructor(private db: LevelUp, private tree: TreeBase & AppendOnlyTree, private hasher: Hasher) {} async getSnapshot(block: number): Promise { - const leafCount = await this.#getLeafCountAtBlock(block); + const meta = await this.#getSnapshotMeta(block); - if (typeof leafCount === 'undefined') { + 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, leafCount, this.tree, this.hasher); + return new AppendOnlySnapshot(this.db, block, meta.numLeaves, meta.root, this.tree, this.hasher); } async snapshot(block: number): Promise { - const leafCountAtBlock = await this.#getLeafCountAtBlock(block); - if (typeof leafCountAtBlock !== 'undefined') { + const meta = await this.#getSnapshotMeta(block); + if (typeof meta !== 'undefined') { // no-op, we already have a snapshot - return new AppendOnlySnapshot(this.db, block, leafCountAtBlock, this.tree, this.hasher); + return new AppendOnlySnapshot(this.db, block, meta.numLeaves, meta.root, this.tree, this.hasher); } const batch = this.db.batch(); @@ -93,19 +94,31 @@ export class AppendOnlySnapshotBuilder implements TreeSnapshotBuilder { } } - const leafCount = this.tree.getNumLeaves(false); - batch.put(snapshotLeafCountKey(treeName, block), String(leafCount)); + 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, leafCount, this.tree, this.hasher); + return new AppendOnlySnapshot(this.db, block, numLeaves, root, this.tree, this.hasher); } - async #getLeafCountAtBlock(block: number): Promise { - const leafCount = await this.db - .get(snapshotLeafCountKey(this.tree.getName(), block)) - .then(x => BigInt(x.toString())) - .catch(() => undefined); - return leafCount; + 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; + } } } @@ -117,6 +130,7 @@ class AppendOnlySnapshot implements TreeSnapshot { private db: LevelUp, private block: number, private leafCount: bigint, + private historicalRoot: Buffer, private tree: TreeBase & AppendOnlyTree, private hasher: Hasher, ) {} @@ -149,7 +163,8 @@ class AppendOnlySnapshot implements TreeSnapshot { } getRoot(): Buffer { - return this.tree.getRoot(false); + // we could recompute it, but it's way cheaper to just store the root + return this.historicalRoot; } async getLeafValue(index: bigint): Promise { diff --git a/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts b/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts index 2a838cea41a..79a27725aa9 100644 --- a/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts +++ b/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts @@ -9,9 +9,9 @@ import { TreeSnapshot, TreeSnapshotBuilder } from './snapshot_builder.js'; const snapshotChildKey = (node: Buffer, child: 0 | 1) => Buffer.concat([Buffer.from('snapshot:node:'), node, Buffer.from(':' + child)]); -// metadata for a snapshot - the root of the historical tree +// metadata for a snapshot const snapshotRootKey = (treeName: string, block: number) => `snapshot:root:${treeName}:${block}`; -const snapshotNumLeavesKey = (treeName: string, block: number) => `snapshot:leafCount:${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 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 index 0ddb0ca7303..631531fcfc2 100644 --- a/yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.test.ts +++ b/yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.test.ts @@ -79,4 +79,18 @@ describe('IndexedTreeSnapshotBuilder', () => { 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/snapshot_builder_test_suite.ts b/yarn-project/merkle-tree/src/snapshots/snapshot_builder_test_suite.ts index 5a3e611c280..3b66c36164c 100644 --- a/yarn-project/merkle-tree/src/snapshots/snapshot_builder_test_suite.ts +++ b/yarn-project/merkle-tree/src/snapshots/snapshot_builder_test_suite.ts @@ -141,5 +141,57 @@ export function describeSnapshotBuilderTestSuite { + 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); + }); + }); }); } From 069a9925db74e62c89ba928c2985cc662aee799f Mon Sep 17 00:00:00 2001 From: Alex Gherghisan Date: Thu, 30 Nov 2023 16:42:10 +0000 Subject: [PATCH 7/8] fiX: update typings after merge --- .../merkle_tree_snapshot_operations_facade.ts | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) 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 index fb4855204e3..93c94d19163 100644 --- 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 @@ -1,5 +1,5 @@ import { Fr } from '@aztec/circuits.js'; -import { IndexedTreeSnapshot, LowLeafWitnessData, TreeSnapshot } from '@aztec/merkle-tree'; +import { BatchInsertionResult, IndexedTreeSnapshot, TreeSnapshot } from '@aztec/merkle-tree'; import { LeafData, MerkleTreeId, SiblingPath } from '@aztec/types'; import { CurrentTreeRoots, HandleL2BlockResult, MerkleTreeDb, MerkleTreeOperations, TreeInfo } from '../index.js'; @@ -107,7 +107,13 @@ export class MerkleTreeSnapshotOperationsFacade implements MerkleTreeOperations return Promise.reject(new Error('Tree snapshot operations are read-only')); } - batchInsert(): Promise<[LowLeafWitnessData[], SiblingPath] | [undefined, SiblingPath]> { + 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')); } From de97077c69f4f07c8c11b727a4b18aabb4c6fd04 Mon Sep 17 00:00:00 2001 From: PhilWindle <60546371+PhilWindle@users.noreply.github.com> Date: Fri, 1 Dec 2023 13:57:58 +0000 Subject: [PATCH 8/8] Update yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts --- yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts b/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts index 79a27725aa9..d77204beafa 100644 --- a/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts +++ b/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts @@ -5,7 +5,7 @@ import { LevelUp, LevelUpChain } from 'levelup'; import { TreeBase } from '../tree_base.js'; import { TreeSnapshot, TreeSnapshotBuilder } from './snapshot_builder.js'; -// ket for a node's children +// key for a node's children const snapshotChildKey = (node: Buffer, child: 0 | 1) => Buffer.concat([Buffer.from('snapshot:node:'), node, Buffer.from(':' + child)]);