feat: persist pxe state

yarn-project/end-to-end/src/e2e_p2p_network.test.ts

@@ -150,7 +150,7 @@ describe('e2e_p2p_network', () => {
     numTxs: number,
   ): Promise<NodeContext> => {
     const rpcConfig = getRpcConfig();
-    const pxeService = await createPXEService(node, rpcConfig, {}, true);
+    const pxeService = await createPXEService(node, rpcConfig, true);
 
     const keyPair = ConstantKeyPair.random(new Grumpkin());
     const completeAddress = CompleteAddress.fromPrivateKeyAndPartialAddress(keyPair.getPrivateKey(), Fr.random());

yarn-project/end-to-end/src/fixtures/utils.ts

@@ -151,7 +151,7 @@ export async function setupPXEService(
   logger: DebugLogger;
 }> {
   const pxeServiceConfig = getPXEServiceConfig();
-  const pxe = await createPXEService(aztecNode, pxeServiceConfig, {}, useLogSuffix);
+  const pxe = await createPXEService(aztecNode, pxeServiceConfig, useLogSuffix);
 
   const wallets = await createAccounts(pxe, numberOfAccounts);
 

yarn-project/foundation/src/serialize/buffer_reader.ts

@@ -31,8 +31,16 @@ export class BufferReader {
    * @param bufferOrReader - A Buffer or BufferReader to initialize the BufferReader.
    * @returns An instance of BufferReader.
    */
-  public static asReader(bufferOrReader: Buffer | BufferReader) {
-    return Buffer.isBuffer(bufferOrReader) ? new BufferReader(bufferOrReader) : bufferOrReader;
+  public static asReader(bufferOrReader: Uint8Array | Buffer | BufferReader): BufferReader {
+    if (bufferOrReader instanceof BufferReader) {
+      return bufferOrReader;
+    }
+
+    const buf = Buffer.isBuffer(bufferOrReader)
+      ? bufferOrReader
+      : Buffer.from(bufferOrReader.buffer, bufferOrReader.byteOffset, bufferOrReader.byteLength);
+
+    return new BufferReader(buf);
   }
 
   /**

yarn-project/key-store/package.json

@@ -32,6 +32,7 @@
   "dependencies": {
     "@aztec/circuits.js": "workspace:^",
     "@aztec/foundation": "workspace:^",
+    "@aztec/kv-store": "workspace:^",
     "@aztec/types": "workspace:^",
     "tslib": "^2.4.0"
   },

yarn-project/key-store/src/test_key_store.ts

@@ -1,5 +1,6 @@
-import { GrumpkinPrivateKey } from '@aztec/circuits.js';
+import { GrumpkinPrivateKey, GrumpkinScalar, Point } from '@aztec/circuits.js';
 import { Grumpkin } from '@aztec/circuits.js/barretenberg';
+import { AztecKVStore, AztecMap } from '@aztec/kv-store';
 import { KeyPair, KeyStore, PublicKey } from '@aztec/types';
 
 import { ConstantKeyPair } from './key_pair.js';
@@ -9,30 +10,27 @@ import { ConstantKeyPair } from './key_pair.js';
  * It should be utilized in testing scenarios where secure key management is not required, and ease-of-use is prioritized.
  */
 export class TestKeyStore implements KeyStore {
-  private accounts: KeyPair[] = [];
-  constructor(private curve: Grumpkin) {}
+  #keys: AztecMap<string, Buffer>;
 
-  public addAccount(privKey: GrumpkinPrivateKey): PublicKey {
-    const keyPair = ConstantKeyPair.fromPrivateKey(this.curve, privKey);
-
-    // check if private key has already been used
-    const account = this.accounts.find(a => a.getPublicKey().equals(keyPair.getPublicKey()));
-    if (account) {
-      return account.getPublicKey();
-    }
+  constructor(private curve: Grumpkin, database: AztecKVStore) {
+    this.#keys = database.createMap('key_store');
+  }
 
-    this.accounts.push(keyPair);
+  public async addAccount(privKey: GrumpkinPrivateKey): Promise<PublicKey> {
+    const keyPair = ConstantKeyPair.fromPrivateKey(this.curve, privKey);
+    await this.#keys.setIfNotExists(keyPair.getPublicKey().toString(), keyPair.getPrivateKey().toBuffer());
     return keyPair.getPublicKey();
   }
 
-  public createAccount(): Promise<PublicKey> {
+  public async createAccount(): Promise<PublicKey> {
     const keyPair = ConstantKeyPair.random(this.curve);
-    this.accounts.push(keyPair);
-    return Promise.resolve(keyPair.getPublicKey());
+    await this.#keys.set(keyPair.getPublicKey().toString(), keyPair.getPrivateKey().toBuffer());
+    return keyPair.getPublicKey();
   }
 
   public getAccounts(): Promise<PublicKey[]> {
-    return Promise.resolve(this.accounts.map(a => a.getPublicKey()));
+    const range = Array.from(this.#keys.keys());
+    return Promise.resolve(range.map(key => Point.fromString(key)));
   }
 
   public getAccountPrivateKey(pubKey: PublicKey): Promise<GrumpkinPrivateKey> {
@@ -48,13 +46,13 @@ export class TestKeyStore implements KeyStore {
    * @param pubKey - The public key of the account to retrieve.
    * @returns The KeyPair object associated with the provided key.
    */
-  private getAccount(pubKey: PublicKey) {
-    const account = this.accounts.find(a => a.getPublicKey().equals(pubKey));
-    if (!account) {
+  private getAccount(pubKey: PublicKey): KeyPair {
+    const privKey = this.#keys.get(pubKey.toString());
+    if (!privKey) {
       throw new Error(
         'Unknown account.\nSee docs for context: https://docs.aztec.network/dev_docs/contracts/common_errors#unknown-contract-error',
       );
     }
-    return account;
+    return ConstantKeyPair.fromPrivateKey(this.curve, GrumpkinScalar.fromBuffer(privKey));
   }
 }

yarn-project/key-store/tsconfig.json

@@ -6,6 +6,9 @@
     "tsBuildInfoFile": ".tsbuildinfo"
   },
   "references": [
+    {
+      "path": "../kv-store"
+    },
     {
       "path": "../circuits.js"
     },

yarn-project/kv-store/.eslintrc.cjs

@@ -0,0 +1 @@
+module.exports = require('@aztec/foundation/eslint');

yarn-project/kv-store/README.md

@@ -0,0 +1,10 @@
+# KV Store
+
+The Aztec KV store is an implementation of a durable key-value database with a pluggable backend. THe only supported backend right now is LMDB by using the [`lmdb-js` package](https://github.com/kriszyp/lmdb-js).
+
+This package exports a number of primitive data structures that can be used to build domain-specific databases in each node component (e.g. a PXE database or an Archiver database). The data structures supported:
+
+- singleton - holds a single value. Great for when a value needs to be stored but it's not a collection (e.g. the latest block header or the length of an array)
+- array - works like a normal in-memory JS array. It can't contain holes and it can be used as a stack (push-pop mechanics).
+- map - a hashmap where keys can be numbers or strings
+- multi-map - just like a map but each key holds multiple values. Can be used for indexing into other data structures

yarn-project/kv-store/package.json

@@ -0,0 +1,50 @@
+{
+  "name": "@aztec/kv-store",
+  "version": "0.1.0",
+  "type": "module",
+  "exports": "./dest/index.js",
+  "scripts": {
+    "build": "yarn clean && tsc -b",
+    "build:dev": "tsc -b --watch",
+    "clean": "rm -rf ./dest .tsbuildinfo",
+    "formatting": "run -T prettier --check ./src && run -T eslint ./src",
+    "formatting:fix": "run -T eslint --fix ./src && run -T prettier -w ./src",
+    "test": "NODE_NO_WARNINGS=1 node --experimental-vm-modules $(yarn bin jest) --passWithNoTests",
+    "start": "DEBUG='aztec:*' && node ./dest/bin/index.js"
+  },
+  "inherits": [
+    "../package.common.json"
+  ],
+  "jest": {
+    "preset": "ts-jest/presets/default-esm",
+    "moduleNameMapper": {
+      "^(\\.{1,2}/.*)\\.[cm]?js$": "$1"
+    },
+    "testRegex": "./src/.*\\.test\\.(js|mjs|ts)$",
+    "rootDir": "./src",
+    "workerThreads": true
+  },
+  "dependencies": {
+    "@aztec/foundation": "workspace:^",
+    "lmdb": "^2.9.1"
+  },
+  "devDependencies": {
+    "@jest/globals": "^29.5.0",
+    "@types/jest": "^29.5.0",
+    "@types/node": "^18.7.23",
+    "jest": "^29.5.0",
+    "jest-mock-extended": "^3.0.3",
+    "ts-jest": "^29.1.0",
+    "ts-node": "^10.9.1",
+    "typescript": "^5.0.4"
+  },
+  "files": [
+    "dest",
+    "src",
+    "!*.test.*"
+  ],
+  "types": "./dest/index.d.ts",
+  "engines": {
+    "node": ">=18"
+  }
+}

yarn-project/kv-store/src/index.ts

@@ -0,0 +1,5 @@
+export * from './interfaces/array.js';
+export * from './interfaces/map.js';
+export * from './interfaces/singleton.js';
+export * from './interfaces/store.js';
+export * from './lmdb/store.js';

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

@@ -0,0 +1,54 @@
+/**
+ * An array backed by a persistent store. Can not have any holes in it.
+ */
+export interface AztecArray<T> {
+  /**
+   * The size of the array
+   */
+  length: number;
+
+  /**
+   * Pushes values to the end of the array
+   * @param vals - The values to push to the end of the array
+   * @returns The new length of the array
+   */
+  push(...vals: T[]): Promise<number>;
+
+  /**
+   * Pops a value from the end of the array.
+   * @returns The value that was popped, or undefined if the array was empty
+   */
+  pop(): Promise<T | undefined>;
+
+  /**
+   * Gets the value at the given index. Index can be in the range [-length, length - 1).
+   * If the index is negative, it will be treated as an offset from the end of the array.
+   *
+   * @param index - The index to get the value from
+   * @returns The value at the given index or undefined if the index is out of bounds
+   */
+  at(index: number): T | undefined;
+
+  /**
+   * Updates the value at the given index. Index can be in the range [-length, length - 1).
+   * @param index - The index to set the value at
+   * @param val - The value to set
+   * @returns Whether the value was set
+   */
+  setAt(index: number, val: T): Promise<boolean>;
+
+  /**
+   * 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/map.ts

@@ -0,0 +1,70 @@
+/**
+ * A map backed by a persistent store.
+ */
+export interface AztecMap<K extends string | number, V> {
+  /**
+   * Gets the value at the given key.
+   * @param key - The key to get the value from
+   */
+  get(key: K): V | undefined;
+
+  /**
+   * Checks if a key exists in the map.
+   * @param key - The key to check
+   * @returns True if the key exists, false otherwise
+   */
+  has(key: K): boolean;
+
+  /**
+   * Sets the value at the given key.
+   * @param key - The key to set the value at
+   * @param val - The value to set
+   */
+  set(key: K, val: V): Promise<boolean>;
+
+  /**
+   * Sets the value at the given key if it does not already exist.
+   * @param key - The key to set the value at
+   * @param val - The value to set
+   */
+  setIfNotExists(key: K, val: V): Promise<boolean>;
+
+  /**
+   * Deletes the value at the given key.
+   * @param key - The key to delete the value at
+   */
+  delete(key: K): Promise<boolean>;
+
+  /**
+   * Iterates over the map's key-value entries
+   */
+  entries(): IterableIterator<[K, V]>;
+
+  /**
+   * Iterates over the map's values
+   */
+  values(): IterableIterator<V>;
+
+  /**
+   * Iterates over the map's keys
+   */
+  keys(): IterableIterator<K>;
+}
+
+/**
+ * A map backed by a persistent store that can have multiple values for a single key.
+ */
+export interface AztecMultiMap<K extends string | number, V> extends AztecMap<K, V> {
+  /**
+   * Gets all the values at the given key.
+   * @param key - The key to get the values from
+   */
+  getValues(key: K): IterableIterator<V>;
+
+  /**
+   * Deletes a specific value at the given key.
+   * @param key - The key to delete the value at
+   * @param val - The value to delete
+   */
+  deleteValue(key: K, val: V): Promise<void>;
+}

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

@@ -0,0 +1,20 @@
+/**
+ * Represents a singleton value in the database.
+ */
+export interface AztecSingleton<T> {
+  /**
+   * Gets the value.
+   */
+  get(): T | undefined;
+
+  /**
+   * Sets the value.
+   * @param val - The new value
+   */
+  set(val: T): Promise<boolean>;
+
+  /**
+   * Deletes the value.
+   */
+  delete(): Promise<boolean>;
+}

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

@@ -0,0 +1,40 @@
+import { AztecArray } from './array.js';
+import { AztecMap, AztecMultiMap } from './map.js';
+import { AztecSingleton } from './singleton.js';
+
+/** A key-value store */
+export interface AztecKVStore {
+  /**
+   * Creates a new map.
+   * @param name - The name of the map
+   * @returns The map
+   */
+  createMap<K extends string | number, V>(name: string): AztecMap<K, V>;
+
+  /**
+   * Creates a new multi-map.
+   * @param name - The name of the multi-map
+   * @returns The multi-map
+   */
+  createMultiMap<K extends string | number, V>(name: string): AztecMultiMap<K, V>;
+
+  /**
+   * Creates a new array.
+   * @param name - The name of the array
+   * @returns The array
+   */
+  createArray<T>(name: string): AztecArray<T>;
+
+  /**
+   * Creates a new singleton.
+   * @param name - The name of the singleton
+   * @returns The singleton
+   */
+  createSingleton<T>(name: string): AztecSingleton<T>;
+
+  /**
+   * 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
+   */
+  transaction<T extends Exclude<any, Promise<any>>>(callback: () => T): Promise<T>;
+}