diff --git a/yarn-project/kv-store/src/index.ts b/yarn-project/kv-store/src/index.ts index 2a71333f9e6b..d0959987625a 100644 --- a/yarn-project/kv-store/src/index.ts +++ b/yarn-project/kv-store/src/index.ts @@ -2,4 +2,5 @@ export * from './interfaces/array.js'; export * from './interfaces/map.js'; export * from './interfaces/singleton.js'; export * from './interfaces/store.js'; +export * from './interfaces/collection.js'; export * from './lmdb/store.js'; diff --git a/yarn-project/kv-store/src/interfaces/collection.ts b/yarn-project/kv-store/src/interfaces/collection.ts new file mode 100644 index 000000000000..09f130d1fc35 --- /dev/null +++ b/yarn-project/kv-store/src/interfaces/collection.ts @@ -0,0 +1,49 @@ +/** + * A collection of items + */ +export interface AztecCollection { + /** + * The size of the collection + */ + size: number; + + /** + * Adds values to the collection + * @param vals - The values to push to the end of the array + */ + insert(...vals: T[]): Promise; + + /** + * Deletes items from the collection + * @param ids - The ids of the items to delete + */ + delete(...ids: number[]): Promise; + + /** + * Gets an item by id + * @param id - The id of the item to get + */ + get(id: number): T | undefined; + + /** + * Gets a subset of items by an index + * @param indexName - The name of the index + * @param key - The key to get the item by + */ + entriesByIndex(indexName: K, key: string): IterableIterator<[number, T]>; + + /** + * Iterates over the array with indexes. + */ + entries(): IterableIterator<[number, T]>; + + /** + * Iterates over the array. + */ + values(): IterableIterator; + + /** + * Iterates over the array. + */ + [Symbol.iterator](): IterableIterator; +} diff --git a/yarn-project/kv-store/src/interfaces/store.ts b/yarn-project/kv-store/src/interfaces/store.ts index d7ccfa3cd29a..51354bde1b5b 100644 --- a/yarn-project/kv-store/src/interfaces/store.ts +++ b/yarn-project/kv-store/src/interfaces/store.ts @@ -1,4 +1,5 @@ import { AztecArray } from './array.js'; +import { AztecCollection } from './collection.js'; import { AztecMap, AztecMultiMap } from './map.js'; import { AztecSingleton } from './singleton.js'; @@ -32,6 +33,13 @@ export interface AztecKVStore { */ createSingleton(name: string): AztecSingleton; + /** + * Creates a new collection + * @param name - The name of the collection + * @param indexFns - Functions to index items by + */ + createCollection(name: string, indexFns: Record string>): AztecCollection; + /** * Starts a transaction. All calls to read/write data while in a transaction are queued and executed atomically. * @param callback - The callback to execute in a transaction diff --git a/yarn-project/kv-store/src/lmdb/collection.ts b/yarn-project/kv-store/src/lmdb/collection.ts new file mode 100644 index 000000000000..8357700bbd3a --- /dev/null +++ b/yarn-project/kv-store/src/lmdb/collection.ts @@ -0,0 +1,127 @@ +import { Database, Key } from 'lmdb'; + +import { AztecArray } from '../interfaces/array.js'; +import { AztecCollection } from '../interfaces/collection.js'; +import { AztecMultiMap } from '../interfaces/map.js'; +import { AztecSingleton } from '../interfaces/singleton.js'; +import { AztecKVStore } from '../interfaces/store.js'; + +/** Maps an object to a value to be indexed by */ +type IndexFn = (item: T) => string; + +/** Internal structure to reference indexes */ +type Index = { + /** Mapping function */ + fn: IndexFn; + /** Map */ + map: AztecMultiMap; +}; + +/** + * A collection + */ +export class LmdbAztecCollection implements AztecCollection { + #db: Database; + #items: AztecArray; + #deletedCount: AztecSingleton; + #indexes = new Map>(); + + constructor(store: AztecKVStore, db: Database, name: string, indexFns: Record string>) { + this.#db = db; + this.#items = store.createArray(name + ':' + 'items'); + this.#deletedCount = store.createSingleton(name + ':' + 'deletedCount'); + + for (const [name, fn] of Object.entries<(item: T) => string>(indexFns)) { + this.#indexes.set(name as K, { + fn, + map: store.createMultiMap(name + ':index:' + name), + }); + } + } + + get size(): number { + return this.#items.length - (this.#deletedCount.get() ?? 0); + } + + get(id: number): T | undefined { + return this.#items.at(id); + } + + *entriesByIndex(indexName: K, key: string): IterableIterator<[number, T]> { + const index = this.#indexes.get(indexName); + + if (typeof index === 'undefined') { + throw new Error(`Index ${indexName} does not exist`); + } + + const ids = index.map.getValues(key); + for (const id of ids) { + const item = this.#items.at(id); + if (typeof item === 'undefined') { + continue; + } + + yield [id, item]; + } + } + + insert(...items: T[]): Promise { + return this.#db.transaction(() => { + const length = this.#items.length; + const ids: number[] = []; + void this.#items.push(...items); + + for (const [index, item] of items.entries()) { + const id = length + index; + ids.push(id); + for (const index of this.#indexes.values()) { + const key = index.fn(item); + void index.map.set(key, id); + } + } + + return ids; + }); + } + + delete(...ids: number[]): Promise { + return this.#db.transaction(() => { + for (const id of ids) { + const item = this.#items.at(id); + + if (typeof item === 'undefined') { + continue; + } + + for (const index of this.#indexes.values()) { + const key = index.fn(item); + void index.map.delete(key); + } + + void this.#items.setAt(id, undefined); + } + + void this.#deletedCount.set((this.#deletedCount.get() ?? 0) + ids.length); + }); + } + + *entries(): IterableIterator<[number, T]> { + for (const entry of this.#items.entries()) { + if (entry[1] === null) { + continue; + } + + yield entry as [number, T]; + } + } + + *values(): IterableIterator { + for (const [_, item] of this.entries()) { + yield item; + } + } + + [Symbol.iterator](): IterableIterator { + return this.values(); + } +} diff --git a/yarn-project/kv-store/src/lmdb/store.ts b/yarn-project/kv-store/src/lmdb/store.ts index 05b7d391ca65..965fd29826f9 100644 --- a/yarn-project/kv-store/src/lmdb/store.ts +++ b/yarn-project/kv-store/src/lmdb/store.ts @@ -3,10 +3,12 @@ import { EthAddress } from '@aztec/foundation/eth-address'; import { Database, Key, RootDatabase, open } from 'lmdb'; import { AztecArray } from '../interfaces/array.js'; +import { AztecCollection } from '../interfaces/collection.js'; import { AztecMap, AztecMultiMap } from '../interfaces/map.js'; import { AztecSingleton } from '../interfaces/singleton.js'; import { AztecKVStore } from '../interfaces/store.js'; import { LmdbAztecArray } from './array.js'; +import { LmdbAztecCollection } from './collection.js'; import { LmdbAztecMap } from './map.js'; import { LmdbAztecSingleton } from './singleton.js'; @@ -96,6 +98,13 @@ export class AztecLmdbStore implements AztecKVStore { return new LmdbAztecSingleton(this.#data, name); } + createCollection( + name: string, + indexFns: Record string>, + ): AztecCollection { + return new LmdbAztecCollection(this, this.#data, name, indexFns); + } + /** * Runs a callback in a transaction. * @param callback - Function to execute in a transaction diff --git a/yarn-project/pxe/src/database/kv_pxe_database.ts b/yarn-project/pxe/src/database/kv_pxe_database.ts index 645f3a8edc10..a4beca7059ce 100644 --- a/yarn-project/pxe/src/database/kv_pxe_database.ts +++ b/yarn-project/pxe/src/database/kv_pxe_database.ts @@ -1,6 +1,6 @@ import { AztecAddress, BlockHeader, CompleteAddress } from '@aztec/circuits.js'; import { Fr } from '@aztec/foundation/fields'; -import { AztecArray, AztecKVStore, AztecMap, AztecMultiMap, AztecSingleton } from '@aztec/kv-store'; +import { AztecArray, AztecCollection, AztecKVStore, AztecMap, AztecMultiMap, AztecSingleton } from '@aztec/kv-store'; import { ContractDao, MerkleTreeId, NoteFilter, PublicKey } from '@aztec/types'; import { NoteDao } from './note_dao.js'; @@ -19,7 +19,7 @@ type SerializedBlockHeader = { */ export class KVPxeDatabase implements PxeDatabase { #blockHeader: AztecSingleton; - #addresses: AztecMap; + #addresses: AztecCollection; #authWitnesses: AztecMap; #capsules: AztecArray; #contracts: AztecMap; @@ -33,7 +33,9 @@ export class KVPxeDatabase implements PxeDatabase { constructor(db: AztecKVStore) { this.#db = db; - this.#addresses = db.createMap('addresses'); + this.#addresses = db.createCollection('addresses', { + address: item => CompleteAddress.fromBuffer(item).address.toString(), + }); this.#authWitnesses = db.createMap('auth_witnesses'); this.#capsules = db.createArray('capsules'); this.#blockHeader = db.createSingleton('block_header'); @@ -221,13 +223,13 @@ export class KVPxeDatabase implements PxeDatabase { return this.#db.transaction(() => { const addressString = completeAddress.address.toString(); const buffer = completeAddress.toBuffer(); - const existing = this.#addresses.get(addressString); - if (!existing) { - void this.#addresses.set(addressString, buffer); + const existing = Array.from(this.#addresses.entriesByIndex('address', addressString)); + if (existing.length === 0) { + void this.#addresses.insert(buffer); return true; } - if (existing.equals(buffer)) { + if (existing[0][1]?.equals(buffer)) { return false; } @@ -238,8 +240,8 @@ export class KVPxeDatabase implements PxeDatabase { } getCompleteAddress(address: AztecAddress): Promise { - const value = this.#addresses.get(address.toString()); - return Promise.resolve(value ? CompleteAddress.fromBuffer(value) : undefined); + const [value] = Array.from(this.#addresses.entriesByIndex('address', address.toString())); + return Promise.resolve(value ? CompleteAddress.fromBuffer(value[1]) : undefined); } getCompleteAddresses(): Promise {