From 63b9cdc5823df540c73b3e53d8e3c4117deb3b02 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=81lvaro=20Rodr=C3=ADguez?= Date: Tue, 5 Dec 2023 13:58:04 +0100 Subject: [PATCH] feat: refactor StandardIndexedTree for abstract leaves and preimages and optimized it (#3530) This PR prepares the standard indexed tree to be used for the public data tree. Changes: - Make StandardIndexedTree abstract in leaves in preparation to accomodate other trees than the nullifier tree - Removed the leaves in memory in the StandardIndexedTree and switch to db queries for low leaf finding (O(1)) - FindLeafIndex is now O(1) too --- .../aztec-node/src/aztec-node/server.ts | 26 +- .../src/structs/rollup/base_rollup.ts | 93 +++- .../circuits.js/src/tests/factories.ts | 2 +- yarn-project/end-to-end/tsconfig.json | 4 +- yarn-project/foundation/package.json | 1 + yarn-project/foundation/src/index.ts | 1 + yarn-project/foundation/src/trees/index.ts | 48 ++ yarn-project/merkle-tree/package.json | 1 + yarn-project/merkle-tree/src/index.ts | 2 +- .../src/interfaces/indexed_tree.ts | 55 ++- .../merkle-tree/src/interfaces/merkle_tree.ts | 8 + .../src/interfaces/update_only_tree.ts | 5 +- yarn-project/merkle-tree/src/load_tree.ts | 4 +- yarn-project/merkle-tree/src/new_tree.ts | 4 +- .../src/snapshots/append_only_snapshot.ts | 11 + .../src/snapshots/base_full_snapshot.ts | 15 +- .../snapshots/indexed_tree_snapshot.test.ts | 67 +-- .../src/snapshots/indexed_tree_snapshot.ts | 58 ++- .../src/snapshots/snapshot_builder.ts | 17 +- .../snapshots/snapshot_builder_test_suite.ts | 21 + .../src/sparse_tree/sparse_tree.ts | 4 + .../standard_indexed_tree.ts | 454 ++++++++++-------- .../test/standard_indexed_tree.test.ts | 67 ++- .../test/standard_indexed_tree_with_append.ts | 43 +- .../src/standard_tree/standard_tree.test.ts | 16 + .../src/standard_tree/standard_tree.ts | 10 + yarn-project/merkle-tree/src/tree_base.ts | 15 +- yarn-project/merkle-tree/tsconfig.json | 3 + .../src/type_conversion.ts | 6 +- .../src/block_builder/solo_block_builder.ts | 21 +- yarn-project/types/src/interfaces/index.ts | 3 +- .../types/src/interfaces/leaf_data.ts | 17 - ...nullifier_witness.ts => nullifier_tree.ts} | 11 +- .../src/interfaces/state_info_provider.ts | 2 +- .../merkle_tree_operations_facade.ts | 39 +- .../merkle_tree_snapshot_operations_facade.ts | 42 +- .../src/world-state-db/merkle_tree_db.ts | 32 +- .../src/world-state-db/merkle_trees.ts | 69 +-- yarn-project/yarn.lock | 1 + 39 files changed, 811 insertions(+), 487 deletions(-) create mode 100644 yarn-project/foundation/src/trees/index.ts delete mode 100644 yarn-project/types/src/interfaces/leaf_data.ts rename yarn-project/types/src/interfaces/{nullifier_witness.ts => nullifier_tree.ts} (79%) diff --git a/yarn-project/aztec-node/src/aztec-node/server.ts b/yarn-project/aztec-node/src/aztec-node/server.ts index 8ec02fe8668..bc415839d7b 100644 --- a/yarn-project/aztec-node/src/aztec-node/server.ts +++ b/yarn-project/aztec-node/src/aztec-node/server.ts @@ -8,6 +8,7 @@ import { L1_TO_L2_MSG_TREE_HEIGHT, NOTE_HASH_TREE_HEIGHT, NULLIFIER_TREE_HEIGHT, + NullifierLeafPreimage, PUBLIC_DATA_TREE_HEIGHT, } from '@aztec/circuits.js'; import { computeGlobalsHash, computePublicDataTreeIndex } from '@aztec/circuits.js/abis'; @@ -429,19 +430,19 @@ export class AztecNodeService implements AztecNode { return undefined; } - const leafDataPromise = db.getLeafData(MerkleTreeId.NULLIFIER_TREE, Number(index)); + const leafPreimagePromise = db.getLeafPreimage(MerkleTreeId.NULLIFIER_TREE, index); const siblingPathPromise = db.getSiblingPath( MerkleTreeId.NULLIFIER_TREE, BigInt(index), ); - const [leafData, siblingPath] = await Promise.all([leafDataPromise, siblingPathPromise]); + const [leafPreimage, siblingPath] = await Promise.all([leafPreimagePromise, siblingPathPromise]); - if (!leafData) { + if (!leafPreimage) { return undefined; } - return new NullifierMembershipWitness(BigInt(index), leafData, siblingPath); + return new NullifierMembershipWitness(BigInt(index), leafPreimage as NullifierLeafPreimage, siblingPath); } /** @@ -463,22 +464,21 @@ export class AztecNodeService implements AztecNode { nullifier: Fr, ): Promise { const committedDb = await this.#getWorldState(blockNumber); - const { index, alreadyPresent } = await committedDb.getPreviousValueIndex( - MerkleTreeId.NULLIFIER_TREE, - nullifier.toBigInt(), - ); + const findResult = await committedDb.getPreviousValueIndex(MerkleTreeId.NULLIFIER_TREE, nullifier.toBigInt()); + if (!findResult) { + return undefined; + } + const { index, alreadyPresent } = findResult; if (alreadyPresent) { this.log.warn(`Nullifier ${nullifier.toBigInt()} already exists in the tree`); } - const leafData = await committedDb.getLeafData(MerkleTreeId.NULLIFIER_TREE, index); - if (!leafData) { - return undefined; - } + const preimageData = (await committedDb.getLeafPreimage(MerkleTreeId.NULLIFIER_TREE, index))!; + const siblingPath = await committedDb.getSiblingPath( MerkleTreeId.NULLIFIER_TREE, BigInt(index), ); - return new NullifierMembershipWitness(BigInt(index), leafData, siblingPath); + return new NullifierMembershipWitness(BigInt(index), preimageData as NullifierLeafPreimage, siblingPath); } /** diff --git a/yarn-project/circuits.js/src/structs/rollup/base_rollup.ts b/yarn-project/circuits.js/src/structs/rollup/base_rollup.ts index 1d6babb0ebe..0ff77419842 100644 --- a/yarn-project/circuits.js/src/structs/rollup/base_rollup.ts +++ b/yarn-project/circuits.js/src/structs/rollup/base_rollup.ts @@ -1,5 +1,7 @@ +import { toBigIntBE, toBufferBE } from '@aztec/foundation/bigint-buffer'; import { Fr } from '@aztec/foundation/fields'; import { BufferReader, Tuple } from '@aztec/foundation/serialize'; +import { IndexedTreeLeaf, IndexedTreeLeafPreimage } from '@aztec/foundation/trees'; import { BLOCKS_TREE_HEIGHT, @@ -25,28 +27,103 @@ import { AppendOnlyTreeSnapshot } from './append_only_tree_snapshot.js'; * Class containing the data of a preimage of a single leaf in the nullifier tree. * Note: It's called preimage because this data gets hashed before being inserted as a node into the `IndexedTree`. */ -export class NullifierLeafPreimage { +export class NullifierLeafPreimage implements IndexedTreeLeafPreimage { constructor( /** * Leaf value inside the indexed tree's linked list. */ - public leafValue: Fr, + public nullifier: Fr, /** * Next value inside the indexed tree's linked list. */ - public nextValue: Fr, + public nextNullifier: Fr, /** * Index of the next leaf in the indexed tree's linked list. */ - public nextIndex: UInt32, + public nextIndex: bigint, ) {} - toBuffer() { - return serializeToBuffer(this.leafValue, this.nextValue, this.nextIndex); + getKey(): bigint { + return this.nullifier.toBigInt(); + } + + getNextKey(): bigint { + return this.nextNullifier.toBigInt(); + } + + getNextIndex(): bigint { + return this.nextIndex; + } + + asLeaf(): NullifierLeaf { + return new NullifierLeaf(this.nullifier); + } + + toBuffer(): Buffer { + return Buffer.concat(this.toHashInputs()); + } + + toHashInputs(): Buffer[] { + return [ + Buffer.from(this.nullifier.toBuffer()), + Buffer.from(toBufferBE(this.nextIndex, 32)), + Buffer.from(this.nextNullifier.toBuffer()), + ]; + } + + clone(): NullifierLeafPreimage { + return new NullifierLeafPreimage(this.nullifier, this.nextNullifier, this.nextIndex); + } + + static empty(): NullifierLeafPreimage { + return new NullifierLeafPreimage(Fr.ZERO, Fr.ZERO, 0n); + } + + static fromBuffer(buf: Buffer): NullifierLeafPreimage { + const nullifier = Fr.fromBuffer(buf.subarray(0, 32)); + const nextIndex = toBigIntBE(buf.subarray(32, 64)); + const nextNullifier = Fr.fromBuffer(buf.subarray(64, 96)); + return new NullifierLeafPreimage(nullifier, nextNullifier, nextIndex); + } + + static fromLeaf(leaf: NullifierLeaf, nextKey: bigint, nextIndex: bigint): NullifierLeafPreimage { + return new NullifierLeafPreimage(leaf.nullifier, new Fr(nextKey), nextIndex); + } + + static clone(preimage: NullifierLeafPreimage): NullifierLeafPreimage { + return new NullifierLeafPreimage(preimage.nullifier, preimage.nextNullifier, preimage.nextIndex); + } +} + +/** + * A nullifier to be inserted in the nullifier tree. + */ +export class NullifierLeaf implements IndexedTreeLeaf { + constructor( + /** + * Nullifier value. + */ + public nullifier: Fr, + ) {} + + getKey(): bigint { + return this.nullifier.toBigInt(); + } + + toBuffer(): Buffer { + return this.nullifier.toBuffer(); + } + + isEmpty(): boolean { + return this.nullifier.isZero(); + } + + static buildDummy(key: bigint): NullifierLeaf { + return new NullifierLeaf(new Fr(key)); } - static empty() { - return new NullifierLeafPreimage(Fr.ZERO, Fr.ZERO, 0); + static fromBuffer(buf: Buffer): NullifierLeaf { + return new NullifierLeaf(Fr.fromBuffer(buf)); } } diff --git a/yarn-project/circuits.js/src/tests/factories.ts b/yarn-project/circuits.js/src/tests/factories.ts index 078393f7ba8..789d84c7708 100644 --- a/yarn-project/circuits.js/src/tests/factories.ts +++ b/yarn-project/circuits.js/src/tests/factories.ts @@ -900,7 +900,7 @@ export function makeBaseRollupInputs(seed = 0): BaseRollupInputs { const lowNullifierLeafPreimages = makeTuple( MAX_NEW_NULLIFIERS_PER_BASE_ROLLUP, - x => new NullifierLeafPreimage(fr(x), fr(x + 0x100), x + 0x200), + x => new NullifierLeafPreimage(fr(x), fr(x + 0x100), BigInt(x + 0x200)), seed + 0x1000, ); diff --git a/yarn-project/end-to-end/tsconfig.json b/yarn-project/end-to-end/tsconfig.json index 8792db4c3c5..c44790b98d8 100644 --- a/yarn-project/end-to-end/tsconfig.json +++ b/yarn-project/end-to-end/tsconfig.json @@ -55,7 +55,5 @@ "path": "../world-state" } ], - "include": [ - "src" - ] + "include": ["src"] } diff --git a/yarn-project/foundation/package.json b/yarn-project/foundation/package.json index 372c031ef39..e6e8c8e66c3 100644 --- a/yarn-project/foundation/package.json +++ b/yarn-project/foundation/package.json @@ -28,6 +28,7 @@ "./sleep": "./dest/sleep/index.js", "./timer": "./dest/timer/index.js", "./transport": "./dest/transport/index.js", + "./trees": "./dest/trees/index.js", "./wasm": "./dest/wasm/index.js", "./worker": "./dest/worker/index.js", "./bigint-buffer": "./dest/bigint-buffer/index.js", diff --git a/yarn-project/foundation/src/index.ts b/yarn-project/foundation/src/index.ts index c75ef10e563..7e06583f10d 100644 --- a/yarn-project/foundation/src/index.ts +++ b/yarn-project/foundation/src/index.ts @@ -21,6 +21,7 @@ export * as serialize from './serialize/index.js'; export * as sleep from './sleep/index.js'; export * as timer from './timer/index.js'; export * as transport from './transport/index.js'; +export * as trees from './trees/index.js'; export * as types from './types/index.js'; export * as url from './url/index.js'; export * as wasm from './wasm/index.js'; diff --git a/yarn-project/foundation/src/trees/index.ts b/yarn-project/foundation/src/trees/index.ts new file mode 100644 index 00000000000..030a59f2570 --- /dev/null +++ b/yarn-project/foundation/src/trees/index.ts @@ -0,0 +1,48 @@ +/** + * A leaf of an indexed merkle tree. + */ +export interface IndexedTreeLeaf { + /** + * Returns key of the leaf. It's used for indexing. + */ + getKey(): bigint; + /** + * Serializes the leaf into a buffer. + */ + toBuffer(): Buffer; + /** + * Returns true if the leaf is empty. + */ + isEmpty(): boolean; +} + +/** + * Preimage of an indexed merkle tree leaf. + */ +export interface IndexedTreeLeafPreimage { + /** + * Returns key of the leaf corresponding to this preimage. + */ + getKey(): bigint; + /** + * Returns the key of the next leaf. + */ + getNextKey(): bigint; + /** + * Returns the index of the next leaf. + */ + getNextIndex(): bigint; + + /** + * Returns the preimage as a leaf. + */ + asLeaf(): IndexedTreeLeaf; + /** + * Serializes the preimage into a buffer. + */ + toBuffer(): Buffer; + /** + * Serializes the preimage to an array of buffers for hashing. + */ + toHashInputs(): Buffer[]; +} diff --git a/yarn-project/merkle-tree/package.json b/yarn-project/merkle-tree/package.json index 6418d219963..4c7b53d8f42 100644 --- a/yarn-project/merkle-tree/package.json +++ b/yarn-project/merkle-tree/package.json @@ -40,6 +40,7 @@ "tslib": "^2.4.0" }, "devDependencies": { + "@aztec/circuits.js": "workspace:^", "@jest/globals": "^29.5.0", "@types/jest": "^29.5.0", "@types/levelup": "^5.1.2", diff --git a/yarn-project/merkle-tree/src/index.ts b/yarn-project/merkle-tree/src/index.ts index 5181cecfc15..68826f44e42 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 { LowLeafWitnessData, StandardIndexedTree } from './standard_indexed_tree/standard_indexed_tree.js'; +export { 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'; diff --git a/yarn-project/merkle-tree/src/interfaces/indexed_tree.ts b/yarn-project/merkle-tree/src/interfaces/indexed_tree.ts index 46c13f49bd9..eee22a3ee2a 100644 --- a/yarn-project/merkle-tree/src/interfaces/indexed_tree.ts +++ b/yarn-project/merkle-tree/src/interfaces/indexed_tree.ts @@ -1,8 +1,26 @@ -import { LeafData, SiblingPath } from '@aztec/types'; +import { IndexedTreeLeafPreimage } from '@aztec/foundation/trees'; +import { SiblingPath } from '@aztec/types'; -import { LowLeafWitnessData } from '../index.js'; import { AppendOnlyTree } from './append_only_tree.js'; +/** + * All of the data to be return during batch insertion. + */ +export interface LowLeafWitnessData { + /** + * Preimage of the low nullifier that proves non membership. + */ + leafPreimage: IndexedTreeLeafPreimage; + /** + * Sibling path to prove membership of low nullifier. + */ + siblingPath: SiblingPath; + /** + * The index of low nullifier. + */ + index: bigint; +} + /** * The result of a batch insertion in an indexed merkle tree. */ @@ -35,27 +53,30 @@ export interface IndexedTree extends AppendOnlyTree { * @param includeUncommitted - If true, the uncommitted changes are included in the search. * @returns The found leaf index and a flag indicating if the corresponding leaf's value is equal to `newValue`. */ - findIndexOfPreviousValue( + findIndexOfPreviousKey( newValue: bigint, includeUncommitted: boolean, - ): { - /** - * The index of the found leaf. - */ - index: number; - /** - * A flag indicating if the corresponding leaf's value is equal to `newValue`. - */ - alreadyPresent: boolean; - }; + ): Promise< + | { + /** + * The index of the found leaf. + */ + index: bigint; + /** + * A flag indicating if the corresponding leaf's value is equal to `newValue`. + */ + alreadyPresent: boolean; + } + | undefined + >; /** - * Gets the latest LeafData copy. - * @param index - Index of the leaf of which to obtain the LeafData copy. + * Gets the latest LeafPreimage copy. + * @param index - Index of the leaf of which to obtain the LeafPreimage copy. * @param includeUncommitted - If true, the uncommitted changes are included in the search. - * @returns A copy of the leaf data at the given index or undefined if the leaf was not found. + * @returns A copy of the leaf preimage at the given index or undefined if the leaf was not found. */ - getLatestLeafDataCopy(index: number, includeUncommitted: boolean): LeafData | undefined; + getLatestLeafPreimageCopy(index: bigint, includeUncommitted: boolean): Promise; /** * Batch insert multiple leaves into the tree. diff --git a/yarn-project/merkle-tree/src/interfaces/merkle_tree.ts b/yarn-project/merkle-tree/src/interfaces/merkle_tree.ts index e4f65b326a2..ba3ffb4309b 100644 --- a/yarn-project/merkle-tree/src/interfaces/merkle_tree.ts +++ b/yarn-project/merkle-tree/src/interfaces/merkle_tree.ts @@ -49,4 +49,12 @@ export interface MerkleTree extends SiblingPathSource { * @param includeUncommitted - Set to true to include uncommitted updates in the data set. */ getLeafValue(index: bigint, includeUncommitted: boolean): Promise; + + /** + * Returns the index of a leaf given its value, or undefined if no leaf with that value is found. + * @param leaf - The leaf value to look for. + * @param includeUncommitted - Indicates whether to include uncommitted data. + * @returns The index of the first leaf found with a given value (undefined if not found). + */ + findLeafIndex(leaf: Buffer, includeUncommitted: boolean): Promise; } 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 6bd5c024d0c..06ce3a24096 100644 --- a/yarn-project/merkle-tree/src/interfaces/update_only_tree.ts +++ b/yarn-project/merkle-tree/src/interfaces/update_only_tree.ts @@ -1,5 +1,3 @@ -import { LeafData } from '@aztec/types'; - import { TreeSnapshotBuilder } from '../snapshots/snapshot_builder.js'; import { MerkleTree } from './merkle_tree.js'; @@ -12,6 +10,5 @@ export interface UpdateOnlyTree extends MerkleTree, TreeSnapshotBuilder { * @param leaf - The leaf value to be updated. * @param index - The leaf to be updated. */ - // TODO: Make this strictly a Buffer - updateLeaf(leaf: Buffer | LeafData, index: bigint): Promise; + updateLeaf(leaf: Buffer, index: bigint): Promise; } diff --git a/yarn-project/merkle-tree/src/load_tree.ts b/yarn-project/merkle-tree/src/load_tree.ts index baabe852735..9753a2d528d 100644 --- a/yarn-project/merkle-tree/src/load_tree.ts +++ b/yarn-project/merkle-tree/src/load_tree.ts @@ -13,14 +13,14 @@ import { TreeBase, decodeMeta } from './tree_base.js'; * @returns The newly created tree. */ export async function loadTree( - c: new (...args: any[]) => T, + c: new (db: LevelUp, hasher: Hasher, name: string, depth: number, size: bigint, root: Buffer) => T, db: LevelUp, hasher: Hasher, name: string, ): Promise { const meta: Buffer = await db.get(name); const { root, depth, size } = decodeMeta(meta); + const tree = new c(db, hasher, name, depth, size, root); - await tree.initFromDb(); return tree; } diff --git a/yarn-project/merkle-tree/src/new_tree.ts b/yarn-project/merkle-tree/src/new_tree.ts index f1cd4c2d3b5..1395d012d25 100644 --- a/yarn-project/merkle-tree/src/new_tree.ts +++ b/yarn-project/merkle-tree/src/new_tree.ts @@ -15,14 +15,14 @@ import { TreeBase } from './tree_base.js'; * @returns The newly created tree. */ export async function newTree( - c: new (...args: any[]) => T, + c: new (db: LevelUp, hasher: Hasher, name: string, depth: number, size: bigint) => T, db: LevelUp, hasher: Hasher, name: string, depth: number, prefilledSize = 1, ): Promise { - const tree = new c(db, hasher, name, depth, 0n, undefined); + const tree = new c(db, hasher, name, depth, 0n); await tree.init(prefilledSize); return tree; } 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 b530e981b27..46361bd5913 100644 --- a/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts +++ b/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.ts @@ -229,4 +229,15 @@ class AppendOnlySnapshot implements TreeSnapshot { return undefined; } } + + async findLeafIndex(value: Buffer): Promise { + const numLeaves = this.getNumLeaves(); + for (let i = 0n; i < numLeaves; i++) { + const currentValue = await this.getLeafValue(i); + if (currentValue && currentValue.equals(value)) { + return i; + } + } + return undefined; + } } diff --git a/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts b/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts index d77204beafa..b1157cf9d44 100644 --- a/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts +++ b/yarn-project/merkle-tree/src/snapshots/base_full_snapshot.ts @@ -65,7 +65,7 @@ export abstract class BaseFullTreeSnapshotBuilder 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); + await this.handleLeaf(i, node, batch); continue; } @@ -98,7 +98,7 @@ export abstract class BaseFullTreeSnapshotBuilder { @@ -218,4 +218,15 @@ export class BaseFullTreeSnapshot implements TreeSnapshot { path.reverse(); return path; } + + async findLeafIndex(value: Buffer): Promise { + const numLeaves = this.getNumLeaves(); + for (let i = 0n; i < numLeaves; i++) { + const currentValue = await this.getLeafValue(i); + if (currentValue && currentValue.equals(value)) { + return i; + } + } + return undefined; + } } 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 631531fcfc2..3846bbcc21d 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 @@ -1,3 +1,6 @@ +import { Fr, NullifierLeaf, NullifierLeafPreimage } from '@aztec/circuits.js'; +import { Hasher } from '@aztec/types'; + import levelup, { LevelUp } from 'levelup'; import { Pedersen, newTree } from '../index.js'; @@ -6,6 +9,12 @@ import { createMemDown } from '../test/utils/create_mem_down.js'; import { IndexedTreeSnapshotBuilder } from './indexed_tree_snapshot.js'; import { describeSnapshotBuilderTestSuite } from './snapshot_builder_test_suite.js'; +class NullifierTree extends StandardIndexedTreeWithAppend { + constructor(db: levelup.LevelUp, hasher: Hasher, name: string, depth: number, size: bigint = 0n, root?: Buffer) { + super(db, hasher, name, depth, size, NullifierLeafPreimage, NullifierLeaf, root); + } +} + describe('IndexedTreeSnapshotBuilder', () => { let db: LevelUp; let tree: StandardIndexedTreeWithAppend; @@ -13,15 +22,15 @@ describe('IndexedTreeSnapshotBuilder', () => { beforeEach(async () => { db = levelup(createMemDown()); - tree = await newTree(StandardIndexedTreeWithAppend, db, new Pedersen(), 'test', 4); - snapshotBuilder = new IndexedTreeSnapshotBuilder(db, tree); + tree = await newTree(NullifierTree, db, new Pedersen(), 'test', 4); + snapshotBuilder = new IndexedTreeSnapshotBuilder(db, tree, NullifierLeafPreimage); }); describeSnapshotBuilderTestSuite( () => tree, () => snapshotBuilder, async () => { - const newLeaves = Array.from({ length: 2 }).map(() => Buffer.from(Math.random().toString())); + const newLeaves = Array.from({ length: 2 }).map(() => new NullifierLeaf(Fr.random()).toBuffer()); await tree.appendLeaves(newLeaves); }, ); @@ -31,14 +40,14 @@ describe('IndexedTreeSnapshotBuilder', () => { 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), + tree.getLatestLeafPreimageCopy(0n, false), + tree.getLatestLeafPreimageCopy(1n, false), + tree.getLatestLeafPreimageCopy(2n, 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), + tree.getLatestLeafPreimageCopy(3n, false), + tree.getLatestLeafPreimageCopy(4n, false), + tree.getLatestLeafPreimageCopy(5n, false), ]); await snapshotBuilder.snapshot(1); @@ -46,35 +55,35 @@ describe('IndexedTreeSnapshotBuilder', () => { 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), + tree.getLatestLeafPreimageCopy(0n, false), + tree.getLatestLeafPreimageCopy(1n, false), + tree.getLatestLeafPreimageCopy(2n, false), + tree.getLatestLeafPreimageCopy(3n, false), + tree.getLatestLeafPreimageCopy(4n, false), + tree.getLatestLeafPreimageCopy(5n, 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), + snapshot1.getLatestLeafPreimageCopy(0n), + snapshot1.getLatestLeafPreimageCopy(1n), + snapshot1.getLatestLeafPreimageCopy(2n), + snapshot1.getLatestLeafPreimageCopy(3n), + snapshot1.getLatestLeafPreimageCopy(4n), + snapshot1.getLatestLeafPreimageCopy(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), + snapshot2.getLatestLeafPreimageCopy(0n), + snapshot2.getLatestLeafPreimageCopy(1n), + snapshot2.getLatestLeafPreimageCopy(2n), + snapshot2.getLatestLeafPreimageCopy(3n), + snapshot2.getLatestLeafPreimageCopy(4n), + snapshot2.getLatestLeafPreimageCopy(5n), ]); expect(actualLeavesAtBlock2).toEqual(expectedLeavesAtBlock2); }); @@ -85,12 +94,12 @@ describe('IndexedTreeSnapshotBuilder', () => { 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); + const historicalPrevValue = await tree.findIndexOfPreviousKey(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); + await expect(snapshot.findIndexOfPreviousKey(2n)).resolves.toEqual(historicalPrevValue); }); }); }); diff --git a/yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.ts b/yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.ts index 6725bd394e5..28aeefdc953 100644 --- a/yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.ts +++ b/yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.ts @@ -1,10 +1,9 @@ -import { toBufferBE } from '@aztec/foundation/bigint-buffer'; -import { LeafData } from '@aztec/types'; +import { IndexedTreeLeafPreimage } from '@aztec/foundation/trees'; import { LevelUp, LevelUpChain } from 'levelup'; import { IndexedTree } from '../interfaces/indexed_tree.js'; -import { decodeTreeValue, encodeTreeValue } from '../standard_indexed_tree/standard_indexed_tree.js'; +import { PreimageFactory } 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'; @@ -17,44 +16,54 @@ export class IndexedTreeSnapshotBuilder extends BaseFullTreeSnapshotBuilder implements TreeSnapshotBuilder { - constructor(db: LevelUp, tree: IndexedTree & TreeBase) { + constructor(db: LevelUp, tree: IndexedTree & TreeBase, private leafPreimageBuilder: PreimageFactory) { super(db, tree); } protected openSnapshot(root: Buffer, numLeaves: bigint): IndexedTreeSnapshot { - return new IndexedTreeSnapshotImpl(this.db, root, numLeaves, this.tree); + return new IndexedTreeSnapshotImpl(this.db, root, numLeaves, this.tree, this.leafPreimageBuilder); } - 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)); + protected async handleLeaf(index: bigint, node: Buffer, batch: LevelUpChain) { + const leafPreimage = await this.tree.getLatestLeafPreimageCopy(index, false); + if (leafPreimage) { + batch.put(snapshotLeafValue(node, index), leafPreimage.toBuffer()); } } } /** A snapshot of an indexed tree at a particular point in time */ class IndexedTreeSnapshotImpl extends BaseFullTreeSnapshot implements IndexedTreeSnapshot { + constructor( + db: LevelUp, + historicRoot: Buffer, + numLeaves: bigint, + tree: IndexedTree & TreeBase, + private leafPreimageBuilder: PreimageFactory, + ) { + super(db, historicRoot, numLeaves, tree); + } + async getLeafValue(index: bigint): Promise { - const leafData = await this.getLatestLeafDataCopy(index); - return leafData ? toBufferBE(leafData.value, 32) : undefined; + const leafPreimage = await this.getLatestLeafPreimageCopy(index); + return leafPreimage?.toBuffer(); } - async getLatestLeafDataCopy(index: bigint): Promise { + async getLatestLeafPreimageCopy(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); + return this.leafPreimageBuilder.fromBuffer(leafValue); } else { return undefined; } } - async findIndexOfPreviousValue(newValue: bigint): Promise<{ + async findIndexOfPreviousKey(newValue: bigint): Promise<{ /** * The index of the found leaf. */ - index: number; + index: bigint; /** * A flag indicating if the corresponding leaf's value is equal to `newValue`. */ @@ -65,18 +74,18 @@ class IndexedTreeSnapshotImpl extends BaseFullTreeSnapshot implements IndexedTre for (let i = 0; i < numLeaves; i++) { // this is very inefficient - const storedLeaf = await this.getLatestLeafDataCopy(BigInt(i))!; + const storedLeaf = await this.getLatestLeafPreimageCopy(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) { + } else if (storedLeaf.getKey() > newValue) { diff.push(newValue); - } else if (storedLeaf.value === newValue) { - return { index: i, alreadyPresent: true }; + } else if (storedLeaf.getKey() === newValue) { + return { index: BigInt(i), alreadyPresent: true }; } else { - diff.push(newValue - storedLeaf.value); + diff.push(newValue - storedLeaf.getKey()); } } @@ -87,6 +96,13 @@ class IndexedTreeSnapshotImpl extends BaseFullTreeSnapshot implements IndexedTre } } - return { index: minIndex, alreadyPresent: false }; + return { index: BigInt(minIndex), alreadyPresent: false }; + } + + async findLeafIndex(value: Buffer): Promise { + const index = await this.tree.findLeafIndex(value, false); + if (index !== undefined && index < this.getNumLeaves()) { + return index; + } } } diff --git a/yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts b/yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts index a6722306301..b1fd74f9bdc 100644 --- a/yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts +++ b/yarn-project/merkle-tree/src/snapshots/snapshot_builder.ts @@ -1,4 +1,5 @@ -import { LeafData, SiblingPath } from '@aztec/types'; +import { IndexedTreeLeafPreimage } from '@aztec/foundation/trees'; +import { SiblingPath } from '@aztec/types'; /** * An interface for a tree that can record snapshots of its contents. @@ -47,6 +48,14 @@ export interface TreeSnapshot { * @param index - The index of the leaf for which a sibling path is required. */ getSiblingPath(index: bigint): Promise>; + + /** + * Returns the index of a leaf given its value, or undefined if no leaf with that value is found. + * @param treeId - The ID of the tree. + * @param value - The leaf value to look for. + * @returns The index of the first leaf found with a given value (undefined if not found). + */ + findLeafIndex(value: Buffer): Promise; } /** A snapshot of an indexed tree */ @@ -55,18 +64,18 @@ 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; + getLatestLeafPreimageCopy(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<{ + findIndexOfPreviousKey(newValue: bigint): Promise<{ /** * The index of the found leaf. */ - index: number; + index: bigint; /** * A flag indicating if the corresponding leaf's value is equal to `newValue`. */ 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 3b66c36164c..f50ff1d69ae 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 @@ -186,6 +186,7 @@ export function describeSnapshotBuilderTestSuite { + it('returns the historical leaf index when the snapshot was taken', async () => { + await modifyTree(tree); + await tree.commit(); + const snapshot = await snapshotBuilder.snapshot(1); + + const initialLastLeafIndex = tree.getNumLeaves(false) - 1n; + let lastLeaf = await tree.getLeafValue(initialLastLeafIndex, false); + expect(await snapshot.findLeafIndex(lastLeaf!)).toBe(initialLastLeafIndex); + + await modifyTree(tree); + await tree.commit(); + + const newLastLeafIndex = tree.getNumLeaves(false) - 1n; + lastLeaf = await tree.getLeafValue(newLastLeafIndex, false); + + expect(await snapshot.findLeafIndex(lastLeaf!)).toBe(undefined); + }); + }); }); } 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 92cdc4152fc..138ca8f21e7 100644 --- a/yarn-project/merkle-tree/src/sparse_tree/sparse_tree.ts +++ b/yarn-project/merkle-tree/src/sparse_tree/sparse_tree.ts @@ -41,4 +41,8 @@ export class SparseTree extends TreeBase implements UpdateOnlyTree { public getSnapshot(block: number): Promise { return this.#snapshotBuilder.getSnapshot(block); } + + public findLeafIndex(_value: Buffer, _includeUncommitted: boolean): Promise { + throw new Error('Finding leaf index is not supported for sparse trees'); + } } 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 ebbf3a3d0ee..325b438a0f1 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 @@ -1,87 +1,107 @@ import { toBigIntBE, toBufferBE } from '@aztec/foundation/bigint-buffer'; -import { Fr } from '@aztec/foundation/fields'; import { createDebugLogger } from '@aztec/foundation/log'; -import { LeafData, SiblingPath } from '@aztec/types'; - -import { BatchInsertionResult, IndexedTree } from '../interfaces/indexed_tree.js'; -import { IndexedTreeSnapshotBuilder } from '../snapshots/indexed_tree_snapshot.js'; -import { IndexedTreeSnapshot } from '../snapshots/snapshot_builder.js'; +import { IndexedTreeLeaf, IndexedTreeLeafPreimage } from '@aztec/foundation/trees'; +import { Hasher, SiblingPath } from '@aztec/types'; + +import { LevelUp } from 'levelup'; + +import { + BatchInsertionResult, + IndexedTree, + IndexedTreeSnapshot, + IndexedTreeSnapshotBuilder, + LowLeafWitnessData, +} from '../index.js'; import { TreeBase } from '../tree_base.js'; const log = createDebugLogger('aztec:standard-indexed-tree'); -const indexToKeyLeaf = (name: string, index: bigint) => { - return `${name}:leaf:${toBufferBE(index, 32).toString('hex')}`; -}; - -const keyLeafToIndex = (key: string): bigint => { - const index = key.split(':')[2]; - return toBigIntBE(Buffer.from(index, 'hex')); -}; - -const zeroLeaf: LeafData = { - value: 0n, - nextValue: 0n, - nextIndex: 0n, -}; - /** - * All of the data to be return during batch insertion. + * Factory for creating leaf preimages. */ -export interface LowLeafWitnessData { +export interface PreimageFactory { + /** + * Creates a new preimage from a leaf. + * @param leaf - Leaf to create a preimage from. + * @param nextKey - Next key of the leaf. + * @param nextIndex - Next index of the leaf. + */ + fromLeaf(leaf: IndexedTreeLeaf, nextKey: bigint, nextIndex: bigint): IndexedTreeLeafPreimage; + /** + * Creates a new preimage from a buffer. + * @param buffer - Buffer to create a preimage from. + */ + fromBuffer(buffer: Buffer): IndexedTreeLeafPreimage; /** - * Preimage of the low nullifier that proves non membership. + * Creates an empty preimage. */ - leafData: LeafData; + empty(): IndexedTreeLeafPreimage; /** - * Sibling path to prove membership of low nullifier. + * Creates a copy of a preimage. + * @param preimage - Preimage to be cloned. */ - siblingPath: SiblingPath; + clone(preimage: IndexedTreeLeafPreimage): IndexedTreeLeafPreimage; +} + +/** + * Factory for creating leaves. + */ +export interface LeafFactory { + /** + * Creates a new leaf from a buffer. + * @param key - Key of the leaf. + */ + buildDummy(key: bigint): IndexedTreeLeaf; /** - * The index of low nullifier. + * Creates a new leaf from a buffer. + * @param buffer - Buffer to create a leaf from. */ - index: bigint; + fromBuffer(buffer: Buffer): IndexedTreeLeaf; } +export const buildDbKeyForPreimage = (name: string, index: bigint) => { + return `${name}:leaf_by_index:${toBufferBE(index, 32).toString('hex')}`; +}; + +export const buildDbKeyForLeafIndex = (name: string, key: bigint) => { + return `${name}:leaf_index_by_leaf_key:${toBufferBE(key, 32).toString('hex')}`; +}; + /** * Pre-compute empty witness. * @param treeHeight - Height of tree for sibling path. * @returns An empty witness. */ -function getEmptyLowLeafWitness(treeHeight: N): LowLeafWitnessData { +function getEmptyLowLeafWitness( + treeHeight: N, + leafPreimageFactory: PreimageFactory, +): LowLeafWitnessData { return { - leafData: zeroLeaf, + leafPreimage: leafPreimageFactory.empty(), index: 0n, siblingPath: new SiblingPath(treeHeight, Array(treeHeight).fill(toBufferBE(0n, 32))), }; } -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]); -}; - -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)); - return { - value, - nextIndex, - nextValue, - } as LeafData; -}; - /** - * Indexed merkle tree. + * Standard implementation of an indexed tree. */ export class StandardIndexedTree extends TreeBase implements IndexedTree { - #snapshotBuilder = new IndexedTreeSnapshotBuilder(this.db, this); - - protected leaves: LeafData[] = []; - protected cachedLeaves: { [key: number]: LeafData } = {}; + #snapshotBuilder = new IndexedTreeSnapshotBuilder(this.db, this, this.leafPreimageFactory); + protected cachedLeafPreimages: { [key: string]: IndexedTreeLeafPreimage } = {}; + + public constructor( + db: LevelUp, + hasher: Hasher, + name: string, + depth: number, + size: bigint = 0n, + protected leafPreimageFactory: PreimageFactory, + protected leafFactory: LeafFactory, + root?: Buffer, + ) { + super(db, hasher, name, depth, size, root); + } /** * Appends the given leaves to the tree. @@ -89,7 +109,7 @@ export class StandardIndexedTree extends TreeBase implements IndexedTree { * @returns Empty promise. * @remarks Use batchInsert method instead. */ - public appendLeaves(_leaves: Buffer[]): Promise { + appendLeaves(_leaves: Buffer[]): Promise { throw new Error('Not implemented'); } @@ -117,88 +137,149 @@ export class StandardIndexedTree extends TreeBase implements IndexedTree { * @param includeUncommitted - Indicates whether to include uncommitted leaves in the computation. * @returns The value of the leaf at the given index or undefined if the leaf is empty. */ - public getLeafValue(index: bigint, includeUncommitted: boolean): Promise { - const leaf = this.getLatestLeafDataCopy(Number(index), includeUncommitted); - if (!leaf) { - return Promise.resolve(undefined); - } - return Promise.resolve(toBufferBE(leaf.value, 32)); + public async getLeafValue(index: bigint, includeUncommitted: boolean): Promise { + const preimage = await this.getLatestLeafPreimageCopy(index, includeUncommitted); + return preimage && preimage.toBuffer(); } /** * 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. + * @param newKey - The new key to be inserted into the tree. * @param includeUncommitted - If true, the uncommitted changes are included in the search. * @returns The found leaf index and a flag indicating if the corresponding leaf's value is equal to `newValue`. */ - findIndexOfPreviousValue( - newValue: bigint, + async findIndexOfPreviousKey( + newKey: bigint, includeUncommitted: boolean, - ): { - /** - * 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(includeUncommitted); - const diff: bigint[] = []; - - for (let i = 0; i < numLeaves; i++) { - const storedLeaf = this.getLatestLeafDataCopy(i, includeUncommitted)!; - - // 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); + ): Promise< + | { + /** + * The index of the found leaf. + */ + index: bigint; + /** + * A flag indicating if the corresponding leaf's value is equal to `newValue`. + */ + alreadyPresent: boolean; + } + | undefined + > { + let lowLeafIndex = await this.getDbLowLeafIndex(newKey); + let lowLeafPreimage = lowLeafIndex !== undefined ? await this.getDbPreimage(lowLeafIndex) : undefined; + + if (includeUncommitted) { + const cachedLowLeafIndex = this.getCachedLowLeafIndex(newKey); + if (cachedLowLeafIndex !== undefined) { + const cachedLowLeafPreimage = this.getCachedPreimage(cachedLowLeafIndex)!; + if (!lowLeafPreimage || cachedLowLeafPreimage.getKey() > lowLeafPreimage.getKey()) { + lowLeafIndex = cachedLowLeafIndex; + lowLeafPreimage = cachedLowLeafPreimage; + } } } - const minIndex = this.findMinIndex(diff); - return { index: minIndex, alreadyPresent: false }; + + if (lowLeafIndex === undefined || !lowLeafPreimage) { + return undefined; + } + + return { + index: lowLeafIndex, + alreadyPresent: lowLeafPreimage.getKey() === newKey, + }; + } + + private getCachedLowLeafIndex(key: bigint): bigint | undefined { + const indexes = Object.getOwnPropertyNames(this.cachedLeafPreimages); + const lowLeafIndexes = indexes + .map(index => ({ + index: BigInt(index), + key: this.cachedLeafPreimages[index].getKey(), + })) + .filter(({ key: candidateKey }) => candidateKey <= key) + .sort((a, b) => Number(b.key - a.key)); + return lowLeafIndexes[0]?.index; + } + + private getCachedLeafIndex(key: bigint): bigint | undefined { + const index = Object.keys(this.cachedLeafPreimages).find(index => { + return this.cachedLeafPreimages[index].getKey() === key; + }); + if (index) { + return BigInt(index); + } + return undefined; + } + + private async getDbLowLeafIndex(key: bigint): Promise { + return await new Promise((resolve, reject) => { + let lowLeafIndex: bigint | undefined; + this.db + .createReadStream({ + gte: buildDbKeyForLeafIndex(this.getName(), 0n), + lte: buildDbKeyForLeafIndex(this.getName(), key), + limit: 1, + reverse: true, + }) + .on('data', data => { + lowLeafIndex = toBigIntBE(data.value); + }) + .on('close', function () {}) + .on('end', function () { + resolve(lowLeafIndex); + }) + .on('error', function () { + log.error('stream error'); + reject(); + }); + }); + } + + private async getDbPreimage(index: bigint): Promise { + const dbPreimage = await this.db + .get(buildDbKeyForPreimage(this.getName(), index)) + .then(data => this.leafPreimageFactory.fromBuffer(data)) + .catch(() => undefined); + return dbPreimage; + } + + private getCachedPreimage(index: bigint): IndexedTreeLeafPreimage | undefined { + return this.cachedLeafPreimages[index.toString()]; } /** - * Gets the latest LeafData copy. - * @param index - Index of the leaf of which to obtain the LeafData copy. + * Gets the latest LeafPreimage copy. + * @param index - Index of the leaf of which to obtain the LeafPreimage copy. * @param includeUncommitted - If true, the uncommitted changes are included in the search. - * @returns A copy of the leaf data at the given index or undefined if the leaf was not found. + * @returns A copy of the leaf preimage at the given index or undefined if the leaf was not found. */ - public getLatestLeafDataCopy(index: number, includeUncommitted: boolean): LeafData | undefined { - const leaf = !includeUncommitted ? this.leaves[index] : this.cachedLeaves[index] ?? this.leaves[index]; - return leaf - ? ({ - value: leaf.value, - nextIndex: leaf.nextIndex, - nextValue: leaf.nextValue, - } as LeafData) - : undefined; + public async getLatestLeafPreimageCopy( + index: bigint, + includeUncommitted: boolean, + ): Promise { + const preimage = !includeUncommitted + ? await this.getDbPreimage(index) + : this.getCachedPreimage(index) ?? (await this.getDbPreimage(index)); + return preimage && this.leafPreimageFactory.clone(preimage); } /** - * Finds the index of the minimum value in an array. - * @param values - The collection of values to be searched. - * @returns The index of the minimum value in the array. + * Returns the index of a leaf given its value, or undefined if no leaf with that value is found. + * @param value - The leaf value to look for. + * @param includeUncommitted - Indicates whether to include uncommitted data. + * @returns The index of the first leaf found with a given value (undefined if not found). */ - private findMinIndex(values: bigint[]) { - if (!values.length) { - return 0; + public async findLeafIndex(value: Buffer, includeUncommitted: boolean): Promise { + const leaf = this.leafFactory.fromBuffer(value); + let index = await this.db + .get(buildDbKeyForLeafIndex(this.getName(), leaf.getKey())) + .then(data => toBigIntBE(data)) + .catch(() => undefined); + + if (includeUncommitted && index === undefined) { + const cachedIndex = this.getCachedLeafIndex(leaf.getKey()); + index = cachedIndex; } - let minIndex = 0; - for (let i = 1; i < values.length; i++) { - if (values[minIndex] > values[i]) { - minIndex = i; - } - } - return minIndex; + return index; } /** @@ -220,66 +301,31 @@ export class StandardIndexedTree extends TreeBase implements IndexedTree { throw new Error(`Prefilled size must be at least 1!`); } - const leaves: LeafData[] = []; + const leaves: IndexedTreeLeafPreimage[] = []; for (let i = 0n; i < prefilledSize; i++) { - const newLeaf = { - value: toBigIntBE(Buffer.from([Number(i)])), - nextIndex: i + 1n, - nextValue: i + 1n, - }; - leaves.push(newLeaf); + const newLeaf = this.leafFactory.buildDummy(i); + const newLeafPreimage = this.leafPreimageFactory.fromLeaf(newLeaf, i + 1n, i + 1n); + leaves.push(newLeafPreimage); } - // Make the first leaf have 0 value - leaves[0].value = 0n; - // Make the last leaf point to the first leaf - leaves[prefilledSize - 1].nextIndex = 0n; - leaves[prefilledSize - 1].nextValue = 0n; + leaves[prefilledSize - 1] = this.leafPreimageFactory.fromLeaf(leaves[prefilledSize - 1].asLeaf(), 0n, 0n); await this.encodeAndAppendLeaves(leaves, true); await this.commit(); } - /** - * Loads Merkle tree data from a database and assigns them to this object. - */ - public async initFromDb(): Promise { - const startingIndex = 0n; - const values: LeafData[] = []; - const promise = new Promise((resolve, reject) => { - this.db - .createReadStream({ - gte: indexToKeyLeaf(this.getName(), startingIndex), - lte: indexToKeyLeaf(this.getName(), 2n ** BigInt(this.getDepth())), - }) - .on('data', function (data) { - const index = keyLeafToIndex(data.key.toString('utf-8')); - values[Number(index)] = decodeTreeValue(data.value); - }) - .on('close', function () {}) - .on('end', function () { - resolve(); - }) - .on('error', function () { - log.error('stream error'); - reject(); - }); - }); - await promise; - this.leaves = values; - } - /** * Commits all the leaves to the database and removes them from a cache. */ private async commitLeaves(): Promise { const batch = this.db.batch(); - const keys = Object.getOwnPropertyNames(this.cachedLeaves); + const keys = Object.getOwnPropertyNames(this.cachedLeafPreimages); for (const key of keys) { - const index = Number(key); - batch.put(indexToKeyLeaf(this.getName(), BigInt(index)), encodeTreeValue(this.cachedLeaves[index])); - this.leaves[index] = this.cachedLeaves[index]; + const leaf = this.cachedLeafPreimages[key]; + const index = BigInt(key); + batch.put(buildDbKeyForPreimage(this.getName(), index), leaf.toBuffer()); + batch.put(buildDbKeyForLeafIndex(this.getName(), leaf.getKey()), toBufferBE(index, 32)); } await batch.write(); this.clearCachedLeaves(); @@ -289,20 +335,21 @@ export class StandardIndexedTree extends TreeBase implements IndexedTree { * Clears the cache. */ private clearCachedLeaves() { - this.cachedLeaves = {}; + this.cachedLeafPreimages = {}; } /** * Updates a leaf in the tree. - * @param leaf - New contents of the leaf. + * @param preimage - New contents of the leaf. * @param index - Index of the leaf to be updated. */ - protected async updateLeaf(leaf: LeafData, index: bigint) { + protected async updateLeaf(preimage: IndexedTreeLeafPreimage, index: bigint) { if (index > this.maxIndex) { throw Error(`Index out of bounds. Index ${index}, max index: ${this.maxIndex}.`); } - const encodedLeaf = this.encodeLeaf(leaf, true); + this.cachedLeafPreimages[index.toString()] = preimage; + const encodedLeaf = this.encodeLeaf(preimage, true); await this.addLeafToCacheAndHashToRoot(encodedLeaf, index); const numLeaves = this.getNumLeaves(true); if (index >= numLeaves) { @@ -426,45 +473,45 @@ export class StandardIndexedTree extends TreeBase implements IndexedTree { leaves: Buffer[], subtreeHeight: SubtreeHeight, ): Promise> { - const emptyLowLeafWitness = getEmptyLowLeafWitness(this.getDepth() as TreeHeight); + const emptyLowLeafWitness = getEmptyLowLeafWitness(this.getDepth() as TreeHeight, this.leafPreimageFactory); // Accumulators const lowLeavesWitnesses: LowLeafWitnessData[] = leaves.map(() => emptyLowLeafWitness); - const pendingInsertionSubtree: LeafData[] = leaves.map(() => zeroLeaf); + const pendingInsertionSubtree: IndexedTreeLeafPreimage[] = leaves.map(() => this.leafPreimageFactory.empty()); // Start info const startInsertionIndex = this.getNumLeaves(true); - const leavesToInsert = leaves.map(leaf => toBigIntBE(leaf)); + const leavesToInsert = leaves.map(leaf => this.leafFactory.fromBuffer(leaf)); const sortedDescendingLeafTuples = leavesToInsert .map((leaf, index) => ({ leaf, index })) - .sort((a, b) => Number(b.leaf - a.leaf)); + .sort((a, b) => Number(b.leaf.getKey() - a.leaf.getKey())); const sortedDescendingLeaves = sortedDescendingLeafTuples.map(leafTuple => leafTuple.leaf); // Get insertion path for each leaf for (let i = 0; i < leavesToInsert.length; i++) { - const newValue = sortedDescendingLeaves[i]; - const originalIndex = leavesToInsert.indexOf(newValue); + const newLeaf = sortedDescendingLeaves[i]; + const originalIndex = leavesToInsert.indexOf(newLeaf); - if (newValue === 0n) { + if (newLeaf.isEmpty()) { continue; } - const indexOfPrevious = this.findIndexOfPreviousValue(newValue, true); - - // get the low leaf - const lowLeaf = this.getLatestLeafDataCopy(indexOfPrevious.index, true); - if (lowLeaf === undefined) { + const indexOfPrevious = await this.findIndexOfPreviousKey(newLeaf.getKey(), true); + if (indexOfPrevious === undefined) { return { lowLeavesWitnessData: undefined, - sortedNewLeaves: sortedDescendingLeafTuples.map(leafTuple => new Fr(leafTuple.leaf).toBuffer()), + sortedNewLeaves: sortedDescendingLeafTuples.map(leafTuple => leafTuple.leaf.toBuffer()), sortedNewLeavesIndexes: sortedDescendingLeafTuples.map(leafTuple => leafTuple.index), newSubtreeSiblingPath: await this.getSubtreeSiblingPath(subtreeHeight, true), }; } + + // get the low leaf (existence checked in getting index) + const lowLeafPreimage = (await this.getLatestLeafPreimageCopy(indexOfPrevious.index, true))!; const siblingPath = await this.getSiblingPath(BigInt(indexOfPrevious.index), true); const witness: LowLeafWitnessData = { - leafData: { ...lowLeaf }, + leafPreimage: lowLeafPreimage, index: BigInt(indexOfPrevious.index), siblingPath, }; @@ -472,20 +519,23 @@ export class StandardIndexedTree extends TreeBase implements IndexedTree { // Update the running paths lowLeavesWitnesses[i] = witness; - const currentPendingLeaf: LeafData = { - value: newValue, - nextValue: lowLeaf.nextValue, - nextIndex: lowLeaf.nextIndex, - }; + const currentPendingPreimageLeaf = this.leafPreimageFactory.fromLeaf( + newLeaf, + lowLeafPreimage.getNextKey(), + lowLeafPreimage.getNextIndex(), + ); - pendingInsertionSubtree[originalIndex] = currentPendingLeaf; + pendingInsertionSubtree[originalIndex] = currentPendingPreimageLeaf; - lowLeaf.nextValue = newValue; - lowLeaf.nextIndex = startInsertionIndex + BigInt(originalIndex); + const newLowLeafPreimage = this.leafPreimageFactory.fromLeaf( + lowLeafPreimage.asLeaf(), + newLeaf.getKey(), + startInsertionIndex + BigInt(originalIndex), + ); const lowLeafIndex = indexOfPrevious.index; - this.cachedLeaves[lowLeafIndex] = lowLeaf; - await this.updateLeaf(lowLeaf, BigInt(lowLeafIndex)); + this.cachedLeafPreimages[lowLeafIndex.toString()] = newLowLeafPreimage; + await this.updateLeaf(newLowLeafPreimage, lowLeafIndex); } const newSubtreeSiblingPath = await this.getSubtreeSiblingPath( @@ -500,7 +550,7 @@ export class StandardIndexedTree extends TreeBase implements IndexedTree { return { lowLeavesWitnessData: lowLeavesWitnesses, - sortedNewLeaves: sortedDescendingLeafTuples.map(leafTuple => Buffer.from(new Fr(leafTuple.leaf).toBuffer())), + sortedNewLeaves: sortedDescendingLeafTuples.map(leafTuple => leafTuple.leaf.toBuffer()), sortedNewLeavesIndexes: sortedDescendingLeafTuples.map(leafTuple => leafTuple.index), newSubtreeSiblingPath, }; @@ -527,19 +577,19 @@ export class StandardIndexedTree extends TreeBase implements IndexedTree { /** * Encodes leaves and appends them to a tree. - * @param leaves - Leaves to encode. + * @param preimages - Leaves to encode. * @param hash0Leaf - Indicates whether 0 value leaf should be hashed. See {@link encodeLeaf}. * @returns Empty promise */ - private async encodeAndAppendLeaves(leaves: LeafData[], hash0Leaf: boolean): Promise { - const startInsertionIndex = Number(this.getNumLeaves(true)); + private async encodeAndAppendLeaves(preimages: IndexedTreeLeafPreimage[], hash0Leaf: boolean): Promise { + const startInsertionIndex = this.getNumLeaves(true); - const serializedLeaves = leaves.map((leaf, i) => { - this.cachedLeaves[startInsertionIndex + i] = leaf; - return this.encodeLeaf(leaf, hash0Leaf); + const hashedLeaves = preimages.map((preimage, i) => { + this.cachedLeafPreimages[(startInsertionIndex + BigInt(i)).toString()] = preimage; + return this.encodeLeaf(preimage, hash0Leaf); }); - await super.appendLeaves(serializedLeaves); + await super.appendLeaves(hashedLeaves); } /** @@ -550,14 +600,12 @@ export class StandardIndexedTree extends TreeBase implements IndexedTree { * nullifier it is improbable that a valid nullifier would be 0. * @returns Leaf encoded in a buffer. */ - private encodeLeaf(leaf: LeafData, hash0Leaf: boolean): Buffer { + private encodeLeaf(leaf: IndexedTreeLeafPreimage, hash0Leaf: boolean): Buffer { let encodedLeaf; - if (!hash0Leaf && leaf.value == 0n) { + if (!hash0Leaf && leaf.getKey() == 0n) { encodedLeaf = toBufferBE(0n, 32); } else { - encodedLeaf = this.hasher.hashInputs( - [leaf.value, leaf.nextIndex, leaf.nextValue].map(val => toBufferBE(val, 32)), - ); + encodedLeaf = this.hasher.hashInputs(leaf.toHashInputs()); } return encodedLeaf; } diff --git a/yarn-project/merkle-tree/src/standard_indexed_tree/test/standard_indexed_tree.test.ts b/yarn-project/merkle-tree/src/standard_indexed_tree/test/standard_indexed_tree.test.ts index 2f6db0b1ac1..9ebc8c30472 100644 --- a/yarn-project/merkle-tree/src/standard_indexed_tree/test/standard_indexed_tree.test.ts +++ b/yarn-project/merkle-tree/src/standard_indexed_tree/test/standard_indexed_tree.test.ts @@ -1,3 +1,4 @@ +import { Fr, NullifierLeaf, NullifierLeafPreimage } from '@aztec/circuits.js'; import { toBufferBE } from '@aztec/foundation/bigint-buffer'; import { Hasher, SiblingPath } from '@aztec/types'; @@ -8,16 +9,22 @@ import { treeTestSuite } from '../../test/test_suite.js'; import { createMemDown } from '../../test/utils/create_mem_down.js'; import { StandardIndexedTreeWithAppend } from './standard_indexed_tree_with_append.js'; +class NullifierTree extends StandardIndexedTreeWithAppend { + constructor(db: levelup.LevelUp, hasher: Hasher, name: string, depth: number, size: bigint = 0n, root?: Buffer) { + super(db, hasher, name, depth, size, NullifierLeafPreimage, NullifierLeaf, root); + } +} + const createDb = async (levelUp: levelup.LevelUp, hasher: Hasher, name: string, depth: number, prefilledSize = 1) => { - return await newTree(StandardIndexedTreeWithAppend, levelUp, hasher, name, depth, prefilledSize); + return await newTree(NullifierTree, levelUp, hasher, name, depth, prefilledSize); }; const createFromName = async (levelUp: levelup.LevelUp, hasher: Hasher, name: string) => { - return await loadTree(StandardIndexedTreeWithAppend, levelUp, hasher, name); + return await loadTree(NullifierTree, levelUp, hasher, name); }; -const createIndexedTreeLeaf = (value: number, nextIndex: number, nextValue: number) => { - return [toBufferBE(BigInt(value), 32), toBufferBE(BigInt(nextIndex), 32), toBufferBE(BigInt(nextValue), 32)]; +const createIndexedTreeLeafHashInputs = (value: number, nextIndex: number, nextValue: number) => { + return new NullifierLeafPreimage(new Fr(value), new Fr(nextValue), BigInt(nextIndex)).toHashInputs(); }; const verifyCommittedState = async ( @@ -57,7 +64,7 @@ describe('StandardIndexedTreeSpecific', () => { * nextVal 0 0 0 0 0 0 0 0. */ - const initialLeafHash = pedersen.hashInputs(createIndexedTreeLeaf(0, 0, 0)); + const initialLeafHash = pedersen.hashInputs(createIndexedTreeLeafHashInputs(0, 0, 0)); const level1ZeroHash = pedersen.hash(INITIAL_LEAF, INITIAL_LEAF); const level2ZeroHash = pedersen.hash(level1ZeroHash, level1ZeroHash); @@ -91,8 +98,8 @@ describe('StandardIndexedTreeSpecific', () => { * nextIdx 1 0 0 0 0 0 0 0 * nextVal 30 0 0 0 0 0 0 0. */ - index0Hash = pedersen.hashInputs(createIndexedTreeLeaf(0, 1, 30)); - let index1Hash = pedersen.hashInputs(createIndexedTreeLeaf(30, 0, 0)); + index0Hash = pedersen.hashInputs(createIndexedTreeLeafHashInputs(0, 1, 30)); + let index1Hash = pedersen.hashInputs(createIndexedTreeLeafHashInputs(30, 0, 0)); e10 = pedersen.hash(index0Hash, index1Hash); e20 = pedersen.hash(e10, level1ZeroHash); root = pedersen.hash(e20, level2ZeroHash); @@ -118,8 +125,8 @@ describe('StandardIndexedTreeSpecific', () => { * nextIdx 2 0 1 0 0 0 0 0 * nextVal 10 0 30 0 0 0 0 0. */ - index0Hash = pedersen.hashInputs(createIndexedTreeLeaf(0, 2, 10)); - let index2Hash = pedersen.hashInputs(createIndexedTreeLeaf(10, 1, 30)); + index0Hash = pedersen.hashInputs(createIndexedTreeLeafHashInputs(0, 2, 10)); + let index2Hash = pedersen.hashInputs(createIndexedTreeLeafHashInputs(10, 1, 30)); e10 = pedersen.hash(index0Hash, index1Hash); let e11 = pedersen.hash(index2Hash, INITIAL_LEAF); e20 = pedersen.hash(e10, e11); @@ -151,8 +158,8 @@ describe('StandardIndexedTreeSpecific', () => { * nextVal 10 0 20 30 0 0 0 0. */ e10 = pedersen.hash(index0Hash, index1Hash); - index2Hash = pedersen.hashInputs(createIndexedTreeLeaf(10, 3, 20)); - const index3Hash = pedersen.hashInputs(createIndexedTreeLeaf(20, 1, 30)); + index2Hash = pedersen.hashInputs(createIndexedTreeLeafHashInputs(10, 3, 20)); + const index3Hash = pedersen.hashInputs(createIndexedTreeLeafHashInputs(20, 1, 30)); e11 = pedersen.hash(index2Hash, index3Hash); e20 = pedersen.hash(e10, e11); root = pedersen.hash(e20, level2ZeroHash); @@ -182,8 +189,8 @@ describe('StandardIndexedTreeSpecific', () => { * nextIdx 2 4 3 1 0 0 0 0 * nextVal 10 50 20 30 0 0 0 0. */ - index1Hash = pedersen.hashInputs(createIndexedTreeLeaf(30, 4, 50)); - const index4Hash = pedersen.hashInputs(createIndexedTreeLeaf(50, 0, 0)); + index1Hash = pedersen.hashInputs(createIndexedTreeLeafHashInputs(30, 4, 50)); + const index4Hash = pedersen.hashInputs(createIndexedTreeLeafHashInputs(50, 0, 0)); e10 = pedersen.hash(index0Hash, index1Hash); e20 = pedersen.hash(e10, e11); const e12 = pedersen.hash(index4Hash, INITIAL_LEAF); @@ -255,7 +262,7 @@ describe('StandardIndexedTreeSpecific', () => { */ const INITIAL_LEAF = toBufferBE(0n, 32); - const initialLeafHash = pedersen.hashInputs(createIndexedTreeLeaf(0, 0, 0)); + const initialLeafHash = pedersen.hashInputs(createIndexedTreeLeafHashInputs(0, 0, 0)); const level1ZeroHash = pedersen.hash(INITIAL_LEAF, INITIAL_LEAF); const level2ZeroHash = pedersen.hash(level1ZeroHash, level1ZeroHash); let index0Hash = initialLeafHash; @@ -289,8 +296,8 @@ describe('StandardIndexedTreeSpecific', () => { * nextIdx 1 0 0 0 0 0 0 0 * nextVal 30 0 0 0 0 0 0 0. */ - index0Hash = pedersen.hashInputs(createIndexedTreeLeaf(0, 1, 30)); - let index1Hash = pedersen.hashInputs(createIndexedTreeLeaf(30, 0, 0)); + index0Hash = pedersen.hashInputs(createIndexedTreeLeafHashInputs(0, 1, 30)); + let index1Hash = pedersen.hashInputs(createIndexedTreeLeafHashInputs(30, 0, 0)); e10 = pedersen.hash(index0Hash, index1Hash); e20 = pedersen.hash(e10, level1ZeroHash); root = pedersen.hash(e20, level2ZeroHash); @@ -315,8 +322,8 @@ describe('StandardIndexedTreeSpecific', () => { * nextIdx 2 0 1 0 0 0 0 0 * nextVal 10 0 30 0 0 0 0 0. */ - index0Hash = pedersen.hashInputs(createIndexedTreeLeaf(0, 2, 10)); - let index2Hash = pedersen.hashInputs(createIndexedTreeLeaf(10, 1, 30)); + index0Hash = pedersen.hashInputs(createIndexedTreeLeafHashInputs(0, 2, 10)); + let index2Hash = pedersen.hashInputs(createIndexedTreeLeafHashInputs(10, 1, 30)); e10 = pedersen.hash(index0Hash, index1Hash); let e11 = pedersen.hash(index2Hash, INITIAL_LEAF); e20 = pedersen.hash(e10, e11); @@ -348,8 +355,8 @@ describe('StandardIndexedTreeSpecific', () => { * nextVal 10 0 20 30 0 0 0 0. */ e10 = pedersen.hash(index0Hash, index1Hash); - index2Hash = pedersen.hashInputs(createIndexedTreeLeaf(10, 3, 20)); - const index3Hash = pedersen.hashInputs(createIndexedTreeLeaf(20, 1, 30)); + index2Hash = pedersen.hashInputs(createIndexedTreeLeafHashInputs(10, 3, 20)); + const index3Hash = pedersen.hashInputs(createIndexedTreeLeafHashInputs(20, 1, 30)); e11 = pedersen.hash(index2Hash, index3Hash); e20 = pedersen.hash(e10, e11); root = pedersen.hash(e20, level2ZeroHash); @@ -387,8 +394,8 @@ describe('StandardIndexedTreeSpecific', () => { * nextIdx 2 6 3 1 0 0 0 0 * nextVal 10 50 20 30 0 0 0 0. */ - index1Hash = pedersen.hashInputs(createIndexedTreeLeaf(30, 6, 50)); - const index6Hash = pedersen.hashInputs(createIndexedTreeLeaf(50, 0, 0)); + index1Hash = pedersen.hashInputs(createIndexedTreeLeafHashInputs(30, 6, 50)); + const index6Hash = pedersen.hashInputs(createIndexedTreeLeafHashInputs(50, 0, 0)); e10 = pedersen.hash(index0Hash, index1Hash); e20 = pedersen.hash(e10, e11); const e13 = pedersen.hash(index6Hash, INITIAL_LEAF); @@ -469,4 +476,20 @@ describe('StandardIndexedTreeSpecific', () => { const actualRoot = insertTree.getRoot(true); expect(actualRoot).toEqual(expectedRoot); }); + + it('should be able to find indexes of leaves', async () => { + const db = levelup(createMemDown()); + const tree = await createDb(db, pedersen, 'test', 3); + const values = [Buffer.alloc(32, 1), Buffer.alloc(32, 2)]; + + await tree.appendLeaves([values[0]]); + + expect(await tree.findLeafIndex(values[0], true)).toBeDefined(); + expect(await tree.findLeafIndex(values[0], false)).toBe(undefined); + expect(await tree.findLeafIndex(values[1], true)).toBe(undefined); + + await tree.commit(); + + expect(await tree.findLeafIndex(values[0], false)).toBeDefined(); + }); }); diff --git a/yarn-project/merkle-tree/src/standard_indexed_tree/test/standard_indexed_tree_with_append.ts b/yarn-project/merkle-tree/src/standard_indexed_tree/test/standard_indexed_tree_with_append.ts index 49a90e611f1..990f4e6ef5f 100644 --- a/yarn-project/merkle-tree/src/standard_indexed_tree/test/standard_indexed_tree_with_append.ts +++ b/yarn-project/merkle-tree/src/standard_indexed_tree/test/standard_indexed_tree_with_append.ts @@ -1,6 +1,3 @@ -import { toBigIntBE } from '@aztec/foundation/bigint-buffer'; -import { LeafData } from '@aztec/types'; - import { StandardIndexedTree } from '../../index.js'; /** @@ -27,10 +24,10 @@ export class StandardIndexedTreeWithAppend extends StandardIndexedTree { * @returns Empty promise. */ private async appendLeaf(leaf: Buffer): Promise { - const newValue = toBigIntBE(leaf); + const newLeaf = this.leafFactory.fromBuffer(leaf); // Special case when appending zero - if (newValue === 0n) { + if (newLeaf.getKey() === 0n) { const newSize = (this.cachedSize ?? this.size) + 1n; if (newSize - 1n > this.maxIndex) { throw Error(`Can't append beyond max index. Max index: ${this.maxIndex}`); @@ -39,27 +36,31 @@ export class StandardIndexedTreeWithAppend extends StandardIndexedTree { return; } - const indexOfPrevious = this.findIndexOfPreviousValue(newValue, true); - const previousLeafCopy = this.getLatestLeafDataCopy(indexOfPrevious.index, true); - - if (previousLeafCopy === undefined) { + const lowLeafIndex = await this.findIndexOfPreviousKey(newLeaf.getKey(), true); + if (lowLeafIndex === undefined) { throw new Error(`Previous leaf not found!`); } - const newLeaf = { - value: newValue, - nextIndex: previousLeafCopy.nextIndex, - nextValue: previousLeafCopy.nextValue, - } as LeafData; - if (indexOfPrevious.alreadyPresent) { + const lowLeafPreimage = (await this.getLatestLeafPreimageCopy(lowLeafIndex.index, true))!; + + const newLeafPreimage = this.leafPreimageFactory.fromLeaf( + newLeaf, + lowLeafPreimage.getNextKey(), + lowLeafPreimage.getNextIndex(), + ); + + if (lowLeafIndex.alreadyPresent) { return; } // insert a new leaf at the highest index and update the values of our previous leaf copy const currentSize = this.getNumLeaves(true); - previousLeafCopy.nextIndex = BigInt(currentSize); - previousLeafCopy.nextValue = newLeaf.value; - this.cachedLeaves[Number(currentSize)] = newLeaf; - this.cachedLeaves[Number(indexOfPrevious.index)] = previousLeafCopy; - await this.updateLeaf(previousLeafCopy, BigInt(indexOfPrevious.index)); - await this.updateLeaf(newLeaf, this.getNumLeaves(true)); + const newLowLeafPreimage = this.leafPreimageFactory.fromLeaf( + lowLeafPreimage.asLeaf(), + newLeaf.getKey(), + BigInt(currentSize), + ); + this.cachedLeafPreimages[Number(currentSize)] = newLeafPreimage; + this.cachedLeafPreimages[Number(lowLeafIndex.index)] = newLowLeafPreimage; + await this.updateLeaf(newLowLeafPreimage, BigInt(lowLeafIndex.index)); + await this.updateLeaf(newLeafPreimage, this.getNumLeaves(true)); } } diff --git a/yarn-project/merkle-tree/src/standard_tree/standard_tree.test.ts b/yarn-project/merkle-tree/src/standard_tree/standard_tree.test.ts index ee3191f42ff..b211017d851 100644 --- a/yarn-project/merkle-tree/src/standard_tree/standard_tree.test.ts +++ b/yarn-project/merkle-tree/src/standard_tree/standard_tree.test.ts @@ -69,4 +69,20 @@ describe('StandardTree_batchAppend', () => { expect(tree.getRoot(true)).toEqual(root); }); + + it('should be able to find indexes of leaves', async () => { + const db = levelup(createMemDown()); + const tree = await createDb(db, pedersen, 'test', 3); + const values = [Buffer.alloc(32, 1), Buffer.alloc(32, 2)]; + + await tree.appendLeaves([values[0]]); + + expect(await tree.findLeafIndex(values[0], true)).toBeDefined(); + expect(await tree.findLeafIndex(values[0], false)).toBe(undefined); + expect(await tree.findLeafIndex(values[1], true)).toBe(undefined); + + await tree.commit(); + + expect(await tree.findLeafIndex(values[0], false)).toBeDefined(); + }); }); 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 0b92572a4b8..55b4f532469 100644 --- a/yarn-project/merkle-tree/src/standard_tree/standard_tree.ts +++ b/yarn-project/merkle-tree/src/standard_tree/standard_tree.ts @@ -24,4 +24,14 @@ export class StandardTree extends TreeBase implements AppendOnlyTree { public getSnapshot(block: number): Promise { return this.#snapshotBuilder.getSnapshot(block); } + + public async findLeafIndex(value: Buffer, includeUncommitted: boolean): Promise { + for (let i = 0n; i < this.getNumLeaves(includeUncommitted); i++) { + const currentValue = await this.getLeafValue(i, includeUncommitted); + if (currentValue && currentValue.equals(value)) { + return i; + } + } + return undefined; + } } diff --git a/yarn-project/merkle-tree/src/tree_base.ts b/yarn-project/merkle-tree/src/tree_base.ts index c57a0499171..291ac258082 100644 --- a/yarn-project/merkle-tree/src/tree_base.ts +++ b/yarn-project/merkle-tree/src/tree_base.ts @@ -241,13 +241,6 @@ export abstract class TreeBase implements MerkleTree { await this.writeMeta(); } - /** - * Initializes the tree from the database. - */ - public async initFromDb(): Promise { - // Implemented only by Indexed Tree to populate the leaf cache. - } - /** * Writes meta data to the provided batch. * @param batch - The batch to which to write the meta data. @@ -307,4 +300,12 @@ export abstract class TreeBase implements MerkleTree { } this.cachedSize = numLeaves + BigInt(leaves.length); } + + /** + * Returns the index of a leaf given its value, or undefined if no leaf with that value is found. + * @param value - The leaf value to look for. + * @param includeUncommitted - Indicates whether to include uncommitted data. + * @returns The index of the first leaf found with a given value (undefined if not found). + */ + abstract findLeafIndex(value: Buffer, includeUncommitted: boolean): Promise; } diff --git a/yarn-project/merkle-tree/tsconfig.json b/yarn-project/merkle-tree/tsconfig.json index 831130c7c84..35f81f8b801 100644 --- a/yarn-project/merkle-tree/tsconfig.json +++ b/yarn-project/merkle-tree/tsconfig.json @@ -11,6 +11,9 @@ }, { "path": "../types" + }, + { + "path": "../circuits.js" } ], "include": ["src"] diff --git a/yarn-project/noir-protocol-circuits/src/type_conversion.ts b/yarn-project/noir-protocol-circuits/src/type_conversion.ts index 9eb01a54516..cc054063d02 100644 --- a/yarn-project/noir-protocol-circuits/src/type_conversion.ts +++ b/yarn-project/noir-protocol-circuits/src/type_conversion.ts @@ -1293,9 +1293,9 @@ export function mapNullifierLeafPreimageToNoir( nullifierLeafPreimage: NullifierLeafPreimage, ): NullifierLeafPreimageNoir { return { - leaf_value: mapFieldToNoir(nullifierLeafPreimage.leafValue), - next_value: mapFieldToNoir(nullifierLeafPreimage.nextValue), - next_index: mapFieldToNoir(new Fr(nullifierLeafPreimage.nextIndex)), + leaf_value: mapFieldToNoir(nullifierLeafPreimage.nullifier), + next_value: mapFieldToNoir(nullifierLeafPreimage.nextNullifier), + next_index: mapNumberToNoir(Number(nullifierLeafPreimage.nextIndex)), }; } diff --git a/yarn-project/sequencer-client/src/block_builder/solo_block_builder.ts b/yarn-project/sequencer-client/src/block_builder/solo_block_builder.ts index 7f70eb98b1c..b8350d44731 100644 --- a/yarn-project/sequencer-client/src/block_builder/solo_block_builder.ts +++ b/yarn-project/sequencer-client/src/block_builder/solo_block_builder.ts @@ -61,8 +61,6 @@ import { BlockBuilder } from './index.js'; import { AllowedTreeNames, OutputWithTreeSnapshot } from './types.js'; const frToBigInt = (fr: Fr) => toBigIntBE(fr.toBuffer()); -const bigintToFr = (num: bigint) => new Fr(num); -const bigintToNum = (num: bigint) => Number(num); // Denotes fields that are not used now, but will be in the future const FUTURE_FR = new Fr(0n); @@ -566,19 +564,16 @@ export class SoloBlockBuilder implements BlockBuilder { const tree = MerkleTreeId.NULLIFIER_TREE; const prevValueIndex = await this.db.getPreviousValueIndex(tree, frToBigInt(nullifier)); - const prevValueInfo = await this.db.getLeafData(tree, prevValueIndex.index); - if (!prevValueInfo) { + if (!prevValueIndex) { throw new Error(`Nullifier tree should have one initial leaf`); } + const prevValuePreimage = (await this.db.getLeafPreimage(tree, prevValueIndex.index))!; + const prevValueSiblingPath = await this.db.getSiblingPath(tree, BigInt(prevValueIndex.index)); return { index: prevValueIndex, - leafPreimage: new NullifierLeafPreimage( - bigintToFr(prevValueInfo.value), - bigintToFr(prevValueInfo.nextValue), - bigintToNum(prevValueInfo.nextIndex), - ), + leafPreimage: prevValuePreimage, witness: new MembershipWitness( NULLIFIER_TREE_HEIGHT, BigInt(prevValueIndex.index), @@ -737,12 +732,8 @@ export class SoloBlockBuilder implements BlockBuilder { newPublicDataReadsSiblingPaths, lowNullifierLeafPreimages: makeTuple(MAX_NEW_NULLIFIERS_PER_BASE_ROLLUP, i => i < nullifierWitnessLeaves.length - ? new NullifierLeafPreimage( - new Fr(nullifierWitnessLeaves[i].leafData.value), - new Fr(nullifierWitnessLeaves[i].leafData.nextValue), - Number(nullifierWitnessLeaves[i].leafData.nextIndex), - ) - : new NullifierLeafPreimage(Fr.ZERO, Fr.ZERO, 0), + ? (nullifierWitnessLeaves[i].leafPreimage as NullifierLeafPreimage) + : NullifierLeafPreimage.empty(), ), lowNullifierMembershipWitness: makeTuple(MAX_NEW_NULLIFIERS_PER_BASE_ROLLUP, i => i < lowNullifierMembershipWitnesses.length diff --git a/yarn-project/types/src/interfaces/index.ts b/yarn-project/types/src/interfaces/index.ts index e0cbc82677c..44ed98bbed4 100644 --- a/yarn-project/types/src/interfaces/index.ts +++ b/yarn-project/types/src/interfaces/index.ts @@ -6,5 +6,4 @@ export * from './deployed-contract.js'; export * from './node-info.js'; export * from './sync-status.js'; export * from './configs.js'; -export * from './leaf_data.js'; -export * from './nullifier_witness.js'; +export * from './nullifier_tree.js'; diff --git a/yarn-project/types/src/interfaces/leaf_data.ts b/yarn-project/types/src/interfaces/leaf_data.ts deleted file mode 100644 index 2edc8e09818..00000000000 --- a/yarn-project/types/src/interfaces/leaf_data.ts +++ /dev/null @@ -1,17 +0,0 @@ -/** - * A leaf of a tree. - */ -export interface LeafData { - /** - * A value of the leaf. - */ - value: bigint; - /** - * An index of the next leaf. - */ - nextIndex: bigint; - /** - * A value of the next leaf. - */ - nextValue: bigint; -} diff --git a/yarn-project/types/src/interfaces/nullifier_witness.ts b/yarn-project/types/src/interfaces/nullifier_tree.ts similarity index 79% rename from yarn-project/types/src/interfaces/nullifier_witness.ts rename to yarn-project/types/src/interfaces/nullifier_tree.ts index 90dc6d9a1c7..14fdf426b8d 100644 --- a/yarn-project/types/src/interfaces/nullifier_witness.ts +++ b/yarn-project/types/src/interfaces/nullifier_tree.ts @@ -1,7 +1,6 @@ -import { Fr, NULLIFIER_TREE_HEIGHT } from '@aztec/circuits.js'; +import { Fr, NULLIFIER_TREE_HEIGHT, NullifierLeafPreimage } from '@aztec/circuits.js'; import { SiblingPath } from '../sibling_path.js'; -import { LeafData } from './leaf_data.js'; /** * Nullifier membership witness. @@ -18,7 +17,7 @@ export class NullifierMembershipWitness { /** * Preimage of the nullifier. */ - public readonly leafData: LeafData, + public readonly leafPreimage: NullifierLeafPreimage, /** * Sibling path to prove membership of the nullifier. */ @@ -32,9 +31,9 @@ export class NullifierMembershipWitness { public toFieldArray(): Fr[] { return [ new Fr(this.index), - new Fr(this.leafData.value), - new Fr(this.leafData.nextIndex), - new Fr(this.leafData.nextValue), + new Fr(this.leafPreimage.nullifier), + new Fr(this.leafPreimage.nextIndex), + new Fr(this.leafPreimage.nextNullifier), ...this.siblingPath.toFieldArray(), ]; } diff --git a/yarn-project/types/src/interfaces/state_info_provider.ts b/yarn-project/types/src/interfaces/state_info_provider.ts index cec3f6fed55..acbfbf72fc4 100644 --- a/yarn-project/types/src/interfaces/state_info_provider.ts +++ b/yarn-project/types/src/interfaces/state_info_provider.ts @@ -12,7 +12,7 @@ import { L1ToL2MessageAndIndex } from '../l1_to_l2_message.js'; import { L2Block } from '../l2_block.js'; import { MerkleTreeId } from '../merkle_tree_id.js'; import { SiblingPath } from '../sibling_path.js'; -import { NullifierMembershipWitness } from './nullifier_witness.js'; +import { NullifierMembershipWitness } from './nullifier_tree.js'; /** Helper type for a specific L2 block number or the latest block number */ type BlockNumber = number | 'latest'; diff --git a/yarn-project/world-state/src/merkle-tree/merkle_tree_operations_facade.ts b/yarn-project/world-state/src/merkle-tree/merkle_tree_operations_facade.ts index 29ba293736d..494040e381e 100644 --- a/yarn-project/world-state/src/merkle-tree/merkle_tree_operations_facade.ts +++ b/yarn-project/world-state/src/merkle-tree/merkle_tree_operations_facade.ts @@ -1,6 +1,8 @@ +import { NullifierLeafPreimage } from '@aztec/circuits.js'; import { Fr } from '@aztec/foundation/fields'; +import { IndexedTreeLeafPreimage } from '@aztec/foundation/trees'; import { BatchInsertionResult } from '@aztec/merkle-tree'; -import { L2Block, LeafData, MerkleTreeId, SiblingPath } from '@aztec/types'; +import { L2Block, MerkleTreeId, SiblingPath } from '@aztec/types'; import { CurrentTreeRoots, HandleL2BlockResult, MerkleTreeDb, MerkleTreeOperations, TreeInfo } from '../index.js'; @@ -59,16 +61,19 @@ export class MerkleTreeOperationsFacade implements MerkleTreeOperations { 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; - }> { + ): Promise< + | { + /** + * The index of the found leaf. + */ + index: bigint; + /** + * A flag indicating if the corresponding leaf's value is equal to `newValue`. + */ + alreadyPresent: boolean; + } + | undefined + > { return this.trees.getPreviousValueIndex(treeId, value, this.includeUncommitted); } @@ -79,7 +84,7 @@ export class MerkleTreeOperationsFacade implements MerkleTreeOperations { * @param index - The index to insert into. * @returns Empty promise. */ - updateLeaf(treeId: MerkleTreeId.NULLIFIER_TREE, leaf: LeafData, index: bigint): Promise { + updateLeaf(treeId: MerkleTreeId.NULLIFIER_TREE, leaf: NullifierLeafPreimage, index: bigint): Promise { return this.trees.updateLeaf(treeId, leaf, index); } @@ -87,10 +92,14 @@ export class MerkleTreeOperationsFacade implements MerkleTreeOperations { * Gets the leaf data at a given index and tree. * @param treeId - The ID of the tree get the leaf from. * @param index - The index of the leaf to get. - * @returns Leaf data. + * @returns Leaf preimage. */ - getLeafData(treeId: MerkleTreeId.NULLIFIER_TREE, index: number): Promise { - return this.trees.getLeafData(treeId, index, this.includeUncommitted); + async getLeafPreimage( + treeId: MerkleTreeId.NULLIFIER_TREE, + index: bigint, + ): Promise { + const preimage = await this.trees.getLeafPreimage(treeId, index, this.includeUncommitted); + return preimage as IndexedTreeLeafPreimage | undefined; } /** 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 93c94d19163..ef112c30e35 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,6 +1,7 @@ import { Fr } from '@aztec/circuits.js'; +import { IndexedTreeLeafPreimage } from '@aztec/foundation/trees'; import { BatchInsertionResult, IndexedTreeSnapshot, TreeSnapshot } from '@aztec/merkle-tree'; -import { LeafData, MerkleTreeId, SiblingPath } from '@aztec/types'; +import { MerkleTreeId, SiblingPath } from '@aztec/types'; import { CurrentTreeRoots, HandleL2BlockResult, MerkleTreeDb, MerkleTreeOperations, TreeInfo } from '../index.js'; @@ -28,23 +29,19 @@ export class MerkleTreeSnapshotOperationsFacade implements MerkleTreeOperations 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; + return tree.findLeafIndex(value); } getLatestGlobalVariablesHash(): Promise { return Promise.reject(new Error('not implemented')); } - async getLeafData(treeId: MerkleTreeId.NULLIFIER_TREE, index: number): Promise { + async getLeafPreimage( + treeId: MerkleTreeId.NULLIFIER_TREE, + index: bigint, + ): Promise { const snapshot = (await this.#getTreeSnapshot(treeId)) as IndexedTreeSnapshot; - return snapshot.getLatestLeafDataCopy(BigInt(index)); + return snapshot.getLatestLeafPreimageCopy(BigInt(index)); } async getLeafValue(treeId: MerkleTreeId, index: bigint): Promise { @@ -55,16 +52,19 @@ export class MerkleTreeSnapshotOperationsFacade implements MerkleTreeOperations 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; - }> { + ): Promise< + | { + /** + * The index of the found leaf. + */ + index: bigint; + /** + * A flag indicating if the corresponding leaf's value is equal to `newValue`. + */ + alreadyPresent: boolean; + } + | undefined + > { return Promise.reject(new Error('not implemented')); } 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 13c6617513d..40b8ab668f2 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,8 +1,9 @@ -import { MAX_NEW_NULLIFIERS_PER_TX } from '@aztec/circuits.js'; +import { MAX_NEW_NULLIFIERS_PER_TX, NullifierLeafPreimage } from '@aztec/circuits.js'; import { Fr } from '@aztec/foundation/fields'; import { createDebugLogger } from '@aztec/foundation/log'; +import { IndexedTreeLeafPreimage } from '@aztec/foundation/trees'; import { BatchInsertionResult, IndexedTreeSnapshot, TreeSnapshot } from '@aztec/merkle-tree'; -import { L2Block, LeafData, MerkleTreeId, SiblingPath } from '@aztec/types'; +import { L2Block, MerkleTreeId, SiblingPath } from '@aztec/types'; /** * Type alias for the nullifier tree ID. @@ -136,23 +137,26 @@ export interface MerkleTreeOperations { getPreviousValueIndex( treeId: IndexedTreeId, 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; - }>; + ): Promise< + | { + /** + * The index of the found leaf. + */ + index: bigint; + /** + * A flag indicating if the corresponding leaf's value is equal to `newValue`. + */ + alreadyPresent: boolean; + } + | undefined + >; /** * Returns the data at a specific leaf. * @param treeId - The tree for which leaf data should be returned. * @param index - The index of the leaf required. */ - getLeafData(treeId: IndexedTreeId, index: number): Promise; + getLeafPreimage(treeId: IndexedTreeId, index: bigint): Promise; /** * Update the leaf data at the given index. @@ -160,7 +164,7 @@ export interface MerkleTreeOperations { * @param leaf - The updated leaf value. * @param index - The index of the leaf to be updated. */ - updateLeaf(treeId: IndexedTreeId | PublicTreeId, leaf: LeafData | Buffer, index: bigint): Promise; + updateLeaf(treeId: IndexedTreeId | PublicTreeId, leaf: NullifierLeafPreimage | Buffer, index: bigint): Promise; /** * Returns the index containing a leaf value. 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 311c071d8b1..56e6ee96972 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 @@ -7,12 +7,15 @@ import { NOTE_HASH_TREE_HEIGHT, NULLIFIER_SUBTREE_HEIGHT, NULLIFIER_TREE_HEIGHT, + NullifierLeaf, + NullifierLeafPreimage, PUBLIC_DATA_TREE_HEIGHT, } from '@aztec/circuits.js'; import { computeBlockHash, computeGlobalsHash } from '@aztec/circuits.js/abis'; import { Committable } from '@aztec/foundation/committable'; import { SerialQueue } from '@aztec/foundation/fifo'; import { createDebugLogger } from '@aztec/foundation/log'; +import { IndexedTreeLeafPreimage } from '@aztec/foundation/trees'; import { AppendOnlyTree, BatchInsertionResult, @@ -25,7 +28,7 @@ import { loadTree, newTree, } from '@aztec/merkle-tree'; -import { L2Block, LeafData, MerkleTreeId, SiblingPath } from '@aztec/types'; +import { Hasher, L2Block, MerkleTreeId, SiblingPath } from '@aztec/types'; import { default as levelup } from 'levelup'; @@ -53,6 +56,15 @@ interface FromDbOptions { const LAST_GLOBAL_VARS_HASH = 'lastGlobalVarsHash'; +/** + * The nullifier tree is an indexed tree. + */ +class NullifierTree extends StandardIndexedTree { + constructor(db: levelup.LevelUp, hasher: Hasher, name: string, depth: number, size: bigint = 0n, root?: Buffer) { + super(db, hasher, name, depth, size, NullifierLeafPreimage, NullifierLeaf, root); + } +} + /** * A convenience class for managing multiple merkle trees. */ @@ -82,7 +94,7 @@ export class MerkleTrees implements MerkleTreeDb { CONTRACT_TREE_HEIGHT, ); const nullifierTree = await initializeTree( - StandardIndexedTree, + NullifierTree, this.db, hasher, `${MerkleTreeId[MerkleTreeId.NULLIFIER_TREE]}`, @@ -310,19 +322,20 @@ export class MerkleTrees implements MerkleTreeDb { treeId: IndexedTreeId, value: bigint, includeUncommitted: boolean, - ): 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 await this.synchronize(() => - Promise.resolve(this._getIndexedTree(treeId).findIndexOfPreviousValue(value, includeUncommitted)), - ); + ): Promise< + | { + /** + * The index of the found leaf. + */ + index: bigint; + /** + * A flag indicating if the corresponding leaf's value is equal to `newValue`. + */ + alreadyPresent: boolean; + } + | undefined + > { + return await this.synchronize(() => this._getIndexedTree(treeId).findIndexOfPreviousKey(value, includeUncommitted)); } /** @@ -330,15 +343,15 @@ export class MerkleTrees implements MerkleTreeDb { * @param treeId - The ID of the tree get the leaf from. * @param index - The index of the leaf to get. * @param includeUncommitted - Indicates whether to include uncommitted data. - * @returns Leaf data. + * @returns Leaf preimage. */ - public async getLeafData( + public async getLeafPreimage( treeId: IndexedTreeId, - index: number, + index: bigint, includeUncommitted: boolean, - ): Promise { + ): Promise { return await this.synchronize(() => - Promise.resolve(this._getIndexedTree(treeId).getLatestLeafDataCopy(index, includeUncommitted)), + this._getIndexedTree(treeId).getLatestLeafPreimageCopy(index, includeUncommitted), ); } @@ -356,13 +369,7 @@ export class MerkleTrees implements MerkleTreeDb { ): Promise { return await this.synchronize(async () => { const tree = this.trees[treeId]; - for (let i = 0n; i < tree.getNumLeaves(includeUncommitted); i++) { - const currentValue = await tree.getLeafValue(i, includeUncommitted); - if (currentValue && currentValue.equals(value)) { - return i; - } - } - return undefined; + return await tree.findLeafIndex(value, includeUncommitted); }); } @@ -373,7 +380,7 @@ export class MerkleTrees implements MerkleTreeDb { * @param index - The index to insert into. * @returns Empty promise. */ - public async updateLeaf(treeId: IndexedTreeId | PublicTreeId, leaf: LeafData | Buffer, index: bigint): Promise { + public async updateLeaf(treeId: IndexedTreeId | PublicTreeId, leaf: Buffer, index: bigint): Promise { return await this.synchronize(() => this._updateLeaf(treeId, leaf, index)); } @@ -486,11 +493,7 @@ export class MerkleTrees implements MerkleTreeDb { return await tree.appendLeaves(leaves); } - private async _updateLeaf( - treeId: IndexedTreeId | PublicTreeId, - leaf: LeafData | Buffer, - index: bigint, - ): Promise { + private async _updateLeaf(treeId: IndexedTreeId | PublicTreeId, leaf: Buffer, index: bigint): Promise { const tree = this.trees[treeId]; if (!('updateLeaf' in tree)) { throw new Error('Tree does not support `updateLeaf` method'); diff --git a/yarn-project/yarn.lock b/yarn-project/yarn.lock index b0696b5e8ae..6dc27eb04c8 100644 --- a/yarn-project/yarn.lock +++ b/yarn-project/yarn.lock @@ -549,6 +549,7 @@ __metadata: version: 0.0.0-use.local resolution: "@aztec/merkle-tree@workspace:merkle-tree" dependencies: + "@aztec/circuits.js": "workspace:^" "@aztec/foundation": "workspace:^" "@aztec/types": "workspace:^" "@jest/globals": ^29.5.0