wip: add collection [ci rebuild]

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';

yarn-project/kv-store/src/interfaces/collection.ts

@@ -0,0 +1,49 @@
+/**
+ * A collection of items
+ */
+export interface AztecCollection<T, K extends string> {
+  /**
+   * 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<number[]>;
+
+  /**
+   * Deletes items from the collection
+   * @param ids - The ids of the items to delete
+   */
+  delete(...ids: number[]): Promise<void>;
+
+  /**
+   * 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<T>;
+
+  /**
+   * Iterates over the array.
+   */
+  [Symbol.iterator](): IterableIterator<T>;
+}

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<T>(name: string): AztecSingleton<T>;
 
+  /**
+   * Creates a new collection
+   * @param name - The name of the collection
+   * @param indexFns - Functions to index items by
+   */
+  createCollection<T, K extends string>(name: string, indexFns: Record<K, (item: T) => string>): AztecCollection<T, K>;
+
   /**
    * 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

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<T> = (item: T) => string;
+
+/** Internal structure to reference indexes */
+type Index<T> = {
+  /** Mapping function */
+  fn: IndexFn<T>;
+  /** Map */
+  map: AztecMultiMap<string, number>;
+};
+
+/**
+ * A collection
+ */
+export class LmdbAztecCollection<T, K extends string> implements AztecCollection<T, K> {
+  #db: Database;
+  #items: AztecArray<T | undefined>;
+  #deletedCount: AztecSingleton<number>;
+  #indexes = new Map<K, Index<T>>();
+
+  constructor(store: AztecKVStore, db: Database<unknown, Key>, name: string, indexFns: Record<K, (item: T) => 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<number[]> {
+    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<void> {
+    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<T> {
+    for (const [_, item] of this.entries()) {
+      yield item;
+    }
+  }
+
+  [Symbol.iterator](): IterableIterator<T> {
+    return this.values();
+  }
+}

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<T, K extends string = string>(
+    name: string,
+    indexFns: Record<K, (item: T) => string>,
+  ): AztecCollection<T, K> {
+    return new LmdbAztecCollection(this, this.#data, name, indexFns);
+  }
+
   /**
    * Runs a callback in a transaction.
    * @param callback - Function to execute in a transaction

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<SerializedBlockHeader>;
-  #addresses: AztecMap<string, Buffer>;
+  #addresses: AztecCollection<Buffer, 'address'>;
   #authWitnesses: AztecMap<string, Buffer[]>;
   #capsules: AztecArray<Buffer[]>;
   #contracts: AztecMap<string, Buffer>;
@@ -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<CompleteAddress | undefined> {
-    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<CompleteAddress[]> {