Skip to content

Commit

Permalink
feat: add indexed tree snapshots
Browse files Browse the repository at this point in the history
  • Loading branch information
alexghr committed Nov 29, 2023
1 parent 955b8f2 commit 51c52dd
Show file tree
Hide file tree
Showing 10 changed files with 349 additions and 158 deletions.
3 changes: 2 additions & 1 deletion yarn-project/merkle-tree/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,12 @@ export * from './interfaces/merkle_tree.js';
export * from './interfaces/update_only_tree.js';
export * from './pedersen.js';
export * from './sparse_tree/sparse_tree.js';
export * from './standard_indexed_tree/standard_indexed_tree.js';
export { LowLeafWitnessData, StandardIndexedTree } from './standard_indexed_tree/standard_indexed_tree.js';
export * from './standard_tree/standard_tree.js';
export { INITIAL_LEAF } from './tree_base.js';
export { newTree } from './new_tree.js';
export { loadTree } from './load_tree.js';
export * from './snapshots/snapshot_builder.js';
export * from './snapshots/full_snapshot.js';
export * from './snapshots/append_only_snapshot.js';
export * from './snapshots/indexed_tree_snapshot.js';
Original file line number Diff line number Diff line change
Expand Up @@ -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) =>
Expand Down Expand Up @@ -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<SiblingPathSource> {
const leafCount = await this.#getLeafCountAtBlock(block);
Expand Down
182 changes: 182 additions & 0 deletions yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts
Original file line number Diff line number Diff line change
@@ -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<T extends TreeBase, S extends SiblingPathSource>
implements TreeSnapshotBuilder<S>
{
constructor(protected db: LevelUp, protected tree: T) {}

async snapshot(block: number): Promise<S> {
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
// <node hash>:0 -> <left child's hash>
// <node hash>:1 -> <right child's hash>
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<S> {
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<Buffer | undefined> {
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<N extends number>(index: bigint): Promise<SiblingPath<N>> {
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<N>(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;
}
}
6 changes: 3 additions & 3 deletions yarn-project/merkle-tree/src/snapshots/full_snapshot.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
Loading

0 comments on commit 51c52dd

Please sign in to comment.