@aztec/world-state 0.56.0 → 0.58.0

package/src/native/native_world_state.ts

@@ -1,16 +1,13 @@
 import {
-  type BatchInsertionResult,
-  type HandleL2BlockAndMessagesResult,
   type IndexedTreeId,
   type L2Block,
-  type MerkleTreeAdminOperations,
   MerkleTreeId,
-  type MerkleTreeLeafType,
-  SiblingPath,
-  type TreeInfo,
+  type MerkleTreeReadOperations,
+  type MerkleTreeWriteOperations,
   TxEffect,
 } from '@aztec/circuit-types';
 import {
+  EthAddress,
   Fr,
   Header,
   MAX_NOTE_HASHES_PER_TX,
@@ -18,279 +15,120 @@ import {
   MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX,
   NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP,
   NullifierLeaf,
-  NullifierLeafPreimage,
+  type NullifierLeafPreimage,
   PartialStateReference,
   PublicDataTreeLeaf,
-  PublicDataTreeLeafPreimage,
   StateReference,
 } from '@aztec/circuits.js';
 import { padArrayEnd } from '@aztec/foundation/collection';
-import { SerialQueue } from '@aztec/foundation/queue';
-import { serializeToBuffer } from '@aztec/foundation/serialize';
-import type { IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
+import { createDebugLogger } from '@aztec/foundation/log';
 
-import bindings from 'bindings';
-import { Decoder, Encoder, addExtension } from 'msgpackr';
-import { isAnyArrayBuffer } from 'util/types';
+import assert from 'assert/strict';
+import { mkdir, mkdtemp, readFile, rm, writeFile } from 'fs/promises';
+import { tmpdir } from 'os';
+import { join } from 'path';
 
-import { type MerkleTreeDb, type TreeSnapshots } from '../world-state-db/merkle_tree_db.js';
-import { MerkleTreeAdminOperationsFacade } from '../world-state-db/merkle_tree_operations_facade.js';
+import { type MerkleTreeAdminDatabase as MerkleTreeDatabase } from '../world-state-db/merkle_tree_db.js';
+import { MerkleTreesFacade, MerkleTreesForkFacade, serializeLeaf } from './merkle_trees_facade.js';
 import {
-  MessageHeader,
-  type SerializedIndexedLeaf,
-  type SerializedLeafValue,
-  TypedMessage,
   WorldStateMessageType,
-  type WorldStateRequest,
-  type WorldStateResponse,
+  type WorldStateStatus,
   blockStateReference,
   treeStateReferenceToSnapshot,
   worldStateRevision,
 } from './message.js';
+import { NativeWorldState } from './native_world_state_instance.js';
 
-// small extension to pack an NodeJS Fr instance to a representation that the C++ code can understand
-// this only works for writes. Unpacking from C++ can't create Fr instances because the data is passed
-// as raw, untagged, buffers. On the NodeJS side we don't know what the buffer represents
-// Adding a tag would be a solution, but it would have to be done on both sides and it's unclear where else
-// C++ fr instances are sent/received/stored.
-addExtension({
-  Class: Fr,
-  write: fr => fr.toBuffer(),
-});
-
-export interface NativeInstance {
-  call(msg: Buffer | Uint8Array): Promise<any>;
-}
-
-export class NativeWorldStateService implements MerkleTreeDb {
-  private nextMessageId = 1;
-
-  private encoder = new Encoder({
-    // always encode JS objects as MessagePack maps
-    // this makes it compatible with other MessagePack decoders
-    useRecords: false,
-    int64AsType: 'bigint',
-  });
+const ROLLUP_ADDRESS_FILE = 'rollup_address';
 
-  private decoder = new Decoder({
-    useRecords: false,
-    int64AsType: 'bigint',
-  });
+export class NativeWorldStateService implements MerkleTreeDatabase {
+  protected initialHeader: Header | undefined;
 
-  private queue = new SerialQueue();
+  protected constructor(
+    protected readonly instance: NativeWorldState,
+    protected readonly log = createDebugLogger('aztec:world-state:database'),
+    private readonly cleanup = () => Promise.resolve(),
+  ) {}
 
-  protected constructor(private instance: NativeInstance) {
-    this.queue.start();
-  }
-
-  static async create(
+  static async new(
+    rollupAddress: EthAddress,
     dataDir: string,
-    libraryName = 'world_state_napi',
-    className = 'WorldState',
+    log = createDebugLogger('aztec:world-state:database'),
+    cleanup = () => Promise.resolve(),
   ): Promise<NativeWorldStateService> {
-    const library = bindings(libraryName);
-    const instance = new library[className](dataDir);
-    const worldState = new NativeWorldStateService(instance);
-    await worldState.init();
-    return worldState;
-  }
+    const rollupAddressFile = join(dataDir, ROLLUP_ADDRESS_FILE);
+    const currentRollupStr = await readFile(rollupAddressFile, 'utf8').catch(() => undefined);
+    const currentRollupAddress = currentRollupStr ? EthAddress.fromString(currentRollupStr.trim()) : undefined;
 
-  protected async init() {
-    const archive = await this.getTreeInfo(MerkleTreeId.ARCHIVE, false);
-    if (archive.size === 0n) {
-      // TODO (alexg) move this to the native module
-      // const header = await this.buildInitialHeader(true);
-      // await this.appendLeaves(MerkleTreeId.ARCHIVE, [header.hash()]);
-      // await this.commit();
+    if (currentRollupAddress && !rollupAddress.equals(currentRollupAddress)) {
+      log.warn('Rollup address changed, deleting database');
+      await rm(dataDir, { recursive: true, force: true });
     }
-  }
-
-  public asLatest(): MerkleTreeAdminOperations {
-    return new MerkleTreeAdminOperationsFacade(this, true);
-  }
 
-  // async buildInitialHeader(ic: boolean = false): Promise<Header> {
-  //   const state = await this.getStateReference(ic);
-  //   return new Header(
-  //     AppendOnlyTreeSnapshot.zero(),
-  //     ContentCommitment.empty(),
-  //     state,
-  //     GlobalVariables.empty(),
-  //     Fr.ZERO,
-  //   );
-  // }
+    await mkdir(dataDir, { recursive: true });
+    await writeFile(rollupAddressFile, rollupAddress.toString(), 'utf8');
 
-  public getInitialHeader(): Header {
-    // TODO (alexg) implement this
-    return Header.empty();
+    const instance = new NativeWorldState(dataDir);
+    const worldState = new this(instance, log, cleanup);
+    await worldState.init();
+    return worldState;
   }
 
-  async appendLeaves<ID extends MerkleTreeId>(treeId: ID, leaves: MerkleTreeLeafType<ID>[]): Promise<void> {
-    await this.call(WorldStateMessageType.APPEND_LEAVES, {
-      leaves: leaves.map(leaf => leaf as any),
-      treeId,
-    });
-  }
+  static async tmp(rollupAddress = EthAddress.ZERO, cleanupTmpDir = true): Promise<NativeWorldStateService> {
+    const log = createDebugLogger('aztec:world-state:database');
+    const dataDir = await mkdtemp(join(tmpdir(), 'aztec-world-state-'));
+    log.debug(`Created temporary world state database: ${dataDir}`);
 
-  async batchInsert<TreeHeight extends number, SubtreeSiblingPathHeight extends number, ID extends IndexedTreeId>(
-    treeId: ID,
-    rawLeaves: Buffer[],
-    subtreeHeight: number,
-  ): Promise<BatchInsertionResult<TreeHeight, SubtreeSiblingPathHeight>> {
-    const leaves = rawLeaves.map((leaf: Buffer) => hydrateLeaf(treeId, leaf)).map(serializeLeaf);
-    const resp = await this.call(WorldStateMessageType.BATCH_INSERT, { leaves, treeId, subtreeDepth: subtreeHeight });
-
-    return {
-      newSubtreeSiblingPath: new SiblingPath<SubtreeSiblingPathHeight>(
-        resp.subtree_path.length as any,
-        resp.subtree_path,
-      ),
-      sortedNewLeaves: resp.sorted_leaves
-        .map(([leaf]) => leaf)
-        .map(deserializeLeafValue)
-        .map(serializeToBuffer),
-      sortedNewLeavesIndexes: resp.sorted_leaves.map(([, index]) => index),
-      lowLeavesWitnessData: resp.low_leaf_witness_data.map(data => ({
-        index: BigInt(data.index),
-        leafPreimage: deserializeIndexedLeaf(data.leaf),
-        siblingPath: new SiblingPath<TreeHeight>(data.path.length as any, data.path),
-      })),
+    // pass a cleanup callback because process.on('beforeExit', cleanup) does not work under Jest
+    const cleanup = async () => {
+      if (cleanupTmpDir) {
+        await rm(dataDir, { recursive: true, force: true });
+        log.debug(`Deleted temporary world state database: ${dataDir}`);
+      } else {
+        log.debug(`Leaving temporary world state database: ${dataDir}`);
+      }
     };
-  }
-
-  async commit(): Promise<void> {
-    await this.call(WorldStateMessageType.COMMIT, void 0);
-  }
-
-  findLeafIndex(
-    treeId: MerkleTreeId,
-    value: MerkleTreeLeafType<MerkleTreeId>,
-    includeUncommitted: boolean,
-  ): Promise<bigint | undefined> {
-    return this.findLeafIndexAfter(treeId, value, 0n, includeUncommitted);
-  }
-
-  async findLeafIndexAfter(
-    treeId: MerkleTreeId,
-    leaf: MerkleTreeLeafType<MerkleTreeId>,
-    startIndex: bigint,
-    includeUncommitted: boolean,
-  ): Promise<bigint | undefined> {
-    const index = await this.call(WorldStateMessageType.FIND_LEAF_INDEX, {
-      leaf: serializeLeaf(hydrateLeaf(treeId, leaf)),
-      revision: worldStateRevision(includeUncommitted),
-      treeId,
-      startIndex,
-    });
-
-    if (typeof index === 'number' || typeof index === 'bigint') {
-      return BigInt(index);
-    } else {
-      return undefined;
-    }
-  }
 
-  async getLeafPreimage(
-    treeId: IndexedTreeId,
-    leafIndex: bigint,
-    args: boolean,
-  ): Promise<IndexedTreeLeafPreimage | undefined> {
-    const resp = await this.call(WorldStateMessageType.GET_LEAF_PREIMAGE, {
-      leafIndex,
-      revision: worldStateRevision(args),
-      treeId,
-    });
-
-    return resp ? deserializeIndexedLeaf(resp) : undefined;
+    return this.new(rollupAddress, dataDir, log, cleanup);
   }
 
-  async getLeafValue(
-    treeId: MerkleTreeId,
-    leafIndex: bigint,
-    includeUncommitted: boolean,
-  ): Promise<MerkleTreeLeafType<MerkleTreeId> | undefined> {
-    const resp = await this.call(WorldStateMessageType.GET_LEAF_VALUE, {
-      leafIndex,
-      revision: worldStateRevision(includeUncommitted),
-      treeId,
-    });
-
-    if (!resp) {
-      return undefined;
-    }
+  protected async init() {
+    this.initialHeader = await this.buildInitialHeader();
+    const committed = this.getCommitted();
 
-    const leaf = deserializeLeafValue(resp);
-    if (leaf instanceof Fr) {
-      return leaf;
-    } else {
-      return leaf.toBuffer();
+    // validate the initial state
+    const archive = await committed.getTreeInfo(MerkleTreeId.ARCHIVE);
+    if (archive.size === 0n) {
+      throw new Error("Archive tree can't be empty");
     }
-  }
 
-  async getPreviousValueIndex(
-    treeId: IndexedTreeId,
-    value: bigint,
-    includeUncommitted: boolean,
-  ): Promise<{ index: bigint; alreadyPresent: boolean } | undefined> {
-    const resp = await this.call(WorldStateMessageType.FIND_LOW_LEAF, {
-      key: new Fr(value),
-      revision: worldStateRevision(includeUncommitted),
-      treeId,
-    });
-    return {
-      alreadyPresent: resp.alreadyPresent,
-      index: BigInt(resp.index),
-    };
+    // the initial header _must_ be the first element in the archive tree
+    // if this assertion fails, check that the hashing done in Header in yarn-project matches the initial header hash done in world_state.cpp
+    const initialHeaderIndex = await committed.findLeafIndex(MerkleTreeId.ARCHIVE, this.initialHeader.hash());
+    assert.strictEqual(initialHeaderIndex, 0n, 'Invalid initial archive state');
   }
 
-  async getSiblingPath(
-    treeId: MerkleTreeId,
-    leafIndex: bigint,
-    includeUncommitted: boolean,
-  ): Promise<SiblingPath<number>> {
-    const siblingPath = await this.call(WorldStateMessageType.GET_SIBLING_PATH, {
-      leafIndex,
-      revision: worldStateRevision(includeUncommitted),
-      treeId,
-    });
-
-    return new SiblingPath(siblingPath.length, siblingPath);
+  public getCommitted(): MerkleTreeReadOperations {
+    return new MerkleTreesFacade(this.instance, this.initialHeader!, worldStateRevision(false, 0, 0));
   }
 
-  getSnapshot(_block: number): Promise<TreeSnapshots> {
-    return Promise.reject(new Error('getSnapshot not implemented'));
+  public getSnapshot(blockNumber: number): MerkleTreeReadOperations {
+    return new MerkleTreesFacade(this.instance, this.initialHeader!, worldStateRevision(false, 0, blockNumber));
   }
 
-  async getStateReference(includeUncommitted: boolean): Promise<StateReference> {
-    const resp = await this.call(WorldStateMessageType.GET_STATE_REFERENCE, {
-      revision: worldStateRevision(includeUncommitted),
+  public async fork(blockNumber?: number): Promise<MerkleTreeWriteOperations> {
+    const resp = await this.instance.call(WorldStateMessageType.CREATE_FORK, {
+      latest: blockNumber === undefined,
+      blockNumber: blockNumber ?? 0,
     });
-
-    return new StateReference(
-      treeStateReferenceToSnapshot(resp.state[MerkleTreeId.L1_TO_L2_MESSAGE_TREE]),
-      new PartialStateReference(
-        treeStateReferenceToSnapshot(resp.state[MerkleTreeId.NOTE_HASH_TREE]),
-        treeStateReferenceToSnapshot(resp.state[MerkleTreeId.NULLIFIER_TREE]),
-        treeStateReferenceToSnapshot(resp.state[MerkleTreeId.PUBLIC_DATA_TREE]),
-      ),
-    );
+    return new MerkleTreesForkFacade(this.instance, this.initialHeader!, worldStateRevision(true, resp.forkId, 0));
   }
 
-  async getTreeInfo(treeId: MerkleTreeId, includeUncommitted: boolean): Promise<TreeInfo> {
-    const resp = await this.call(WorldStateMessageType.GET_TREE_INFO, {
-      treeId: treeId,
-      revision: worldStateRevision(includeUncommitted),
-    });
-
-    return {
-      depth: resp.depth,
-      root: resp.root,
-      size: BigInt(resp.size),
-      treeId,
-    };
+  public getInitialHeader(): Header {
+    return this.initialHeader!;
   }
 
-  async handleL2BlockAndMessages(l2Block: L2Block, l1ToL2Messages: Fr[]): Promise<HandleL2BlockAndMessagesResult> {
+  public async handleL2BlockAndMessages(l2Block: L2Block, l1ToL2Messages: Fr[]): Promise<WorldStateStatus> {
     // We have to pad both the tx effects and the values within tx effects because that's how the trees are built
     // by circuits.
     const paddedTxEffects = padArrayEnd(
@@ -321,142 +159,83 @@ export class NativeWorldStateService implements MerkleTreeDb {
       batchesOfPaddedPublicDataWrites.push(batch);
     }
 
-    return await this.call(WorldStateMessageType.SYNC_BLOCK, {
-      blockHash: l2Block.hash(),
+    const response = await this.instance.call(WorldStateMessageType.SYNC_BLOCK, {
+      blockNumber: l2Block.number,
+      blockHeaderHash: l2Block.header.hash(),
       paddedL1ToL2Messages: paddedL1ToL2Messages.map(serializeLeaf),
       paddedNoteHashes: paddedNoteHashes.map(serializeLeaf),
       paddedNullifiers: paddedNullifiers.map(serializeLeaf),
       batchesOfPaddedPublicDataWrites: batchesOfPaddedPublicDataWrites.map(batch => batch.map(serializeLeaf)),
       blockStateRef: blockStateReference(l2Block.header.state),
     });
+    return response.status;
   }
 
-  async rollback(): Promise<void> {
-    await this.call(WorldStateMessageType.ROLLBACK, void 0);
+  public async close(): Promise<void> {
+    await this.instance.close();
+    await this.cleanup();
   }
 
-  async updateArchive(header: Header, _includeUncommitted: boolean): Promise<void> {
-    await this.call(WorldStateMessageType.UPDATE_ARCHIVE, {
-      blockHash: header.hash().toBuffer(),
-      blockStateRef: blockStateReference(header.state),
-    });
+  private async buildInitialHeader(): Promise<Header> {
+    const state = await this.getInitialStateReference();
+    return Header.empty({ state });
   }
 
-  updateLeaf<ID extends IndexedTreeId>(
-    _treeId: ID,
-    _leaf: NullifierLeafPreimage | Buffer,
-    _index: bigint,
-  ): Promise<void> {
-    return Promise.reject(new Error('Method not implemented'));
+  /**
+   * Advances the finalised block number to be the number provided
+   * @param toBlockNumber The block number that is now the tip of the finalised chain
+   * @returns The new WorldStateStatus
+   */
+  public async setFinalised(toBlockNumber: bigint) {
+    return await this.instance.call(WorldStateMessageType.FINALISE_BLOCKS, {
+      toBlockNumber,
+    });
   }
 
-  private call<T extends WorldStateMessageType>(
-    messageType: T,
-    body: WorldStateRequest[T],
-  ): Promise<WorldStateResponse[T]> {
-    return this.queue.put(async () => {
-      const request = new TypedMessage(messageType, new MessageHeader({ messageId: this.nextMessageId++ }), body);
-
-      const encodedRequest = this.encoder.encode(request);
-      const encodedResponse = await this.instance.call(encodedRequest);
-
-      const buf = Buffer.isBuffer(encodedResponse)
-        ? encodedResponse
-        : isAnyArrayBuffer(encodedResponse)
-        ? Buffer.from(encodedResponse)
-        : encodedResponse;
-
-      if (!Buffer.isBuffer(buf)) {
-        throw new TypeError(
-          'Invalid encoded response: expected Buffer or ArrayBuffer, got ' +
-            (encodedResponse === null ? 'null' : typeof encodedResponse),
-        );
-      }
-
-      const decodedResponse = this.decoder.unpack(buf);
-      if (!TypedMessage.isTypedMessageLike(decodedResponse)) {
-        throw new TypeError(
-          'Invalid response: expected TypedMessageLike, got ' +
-            (decodedResponse === null ? 'null' : typeof decodedResponse),
-        );
-      }
-
-      const response = TypedMessage.fromMessagePack<T, WorldStateResponse[T]>(decodedResponse);
-
-      if (response.header.requestId !== request.header.messageId) {
-        throw new Error(
-          'Response ID does not match request: ' + response.header.requestId + ' != ' + request.header.messageId,
-        );
-      }
-
-      if (response.msgType !== messageType) {
-        throw new Error('Invalid response message type: ' + response.msgType + ' != ' + messageType);
-      }
-
-      return response.value;
+  /**
+   * Removes all historical snapshots up to but not including the given block number
+   * @param toBlockNumber The block number of the new oldest historical block
+   * @returns The new WorldStateStatus
+   */
+  public async removeHistoricalBlocks(toBlockNumber: bigint) {
+    return await this.instance.call(WorldStateMessageType.REMOVE_HISTORICAL_BLOCKS, {
+      toBlockNumber,
     });
   }
 
-  public async stop(): Promise<void> {
-    await this.queue.end();
+  /**
+   * Removes all pending blocks down to but not including the given block number
+   * @param toBlockNumber The block number of the new tip of the pending chain,
+   * @returns The new WorldStateStatus
+   */
+  public async unwindBlocks(toBlockNumber: bigint) {
+    return await this.instance.call(WorldStateMessageType.UNWIND_BLOCKS, {
+      toBlockNumber,
+    });
   }
 
-  public delete(): Promise<void> {
-    return Promise.reject(new Error('Method not implemented'));
+  public async getStatus() {
+    return await this.instance.call(WorldStateMessageType.GET_STATUS, void 0);
   }
 
-  public fork(): Promise<MerkleTreeDb> {
+  updateLeaf<ID extends IndexedTreeId>(
+    _treeId: ID,
+    _leaf: NullifierLeafPreimage | Buffer,
+    _index: bigint,
+  ): Promise<void> {
     return Promise.reject(new Error('Method not implemented'));
   }
-}
-
-function hydrateLeaf<ID extends MerkleTreeId>(treeId: ID, leaf: Fr | Buffer) {
-  if (leaf instanceof Fr) {
-    return leaf;
-  } else if (treeId === MerkleTreeId.NULLIFIER_TREE) {
-    return NullifierLeaf.fromBuffer(leaf);
-  } else if (treeId === MerkleTreeId.PUBLIC_DATA_TREE) {
-    return PublicDataTreeLeaf.fromBuffer(leaf);
-  } else {
-    return Fr.fromBuffer(leaf);
-  }
-}
-
-function serializeLeaf(leaf: Fr | NullifierLeaf | PublicDataTreeLeaf): SerializedLeafValue {
-  if (leaf instanceof Fr) {
-    return leaf.toBuffer();
-  } else if (leaf instanceof NullifierLeaf) {
-    return { value: leaf.nullifier.toBuffer() };
-  } else {
-    return { value: leaf.value.toBuffer(), slot: leaf.slot.toBuffer() };
-  }
-}
 
-function deserializeLeafValue(leaf: SerializedLeafValue): Fr | NullifierLeaf | PublicDataTreeLeaf {
-  if (Buffer.isBuffer(leaf)) {
-    return Fr.fromBuffer(leaf);
-  } else if ('slot' in leaf) {
-    return new PublicDataTreeLeaf(Fr.fromBuffer(leaf.slot), Fr.fromBuffer(leaf.value));
-  } else {
-    return new NullifierLeaf(Fr.fromBuffer(leaf.value));
-  }
-}
+  private async getInitialStateReference(): Promise<StateReference> {
+    const resp = await this.instance.call(WorldStateMessageType.GET_INITIAL_STATE_REFERENCE, void 0);
 
-function deserializeIndexedLeaf(leaf: SerializedIndexedLeaf): IndexedTreeLeafPreimage {
-  if ('slot' in leaf.value) {
-    return new PublicDataTreeLeafPreimage(
-      Fr.fromBuffer(leaf.value.slot),
-      Fr.fromBuffer(leaf.value.value),
-      Fr.fromBuffer(leaf.nextValue),
-      BigInt(leaf.nextIndex),
-    );
-  } else if ('value' in leaf.value) {
-    return new NullifierLeafPreimage(
-      Fr.fromBuffer(leaf.value.value),
-      Fr.fromBuffer(leaf.nextValue),
-      BigInt(leaf.nextIndex),
+    return new StateReference(
+      treeStateReferenceToSnapshot(resp.state[MerkleTreeId.L1_TO_L2_MESSAGE_TREE]),
+      new PartialStateReference(
+        treeStateReferenceToSnapshot(resp.state[MerkleTreeId.NOTE_HASH_TREE]),
+        treeStateReferenceToSnapshot(resp.state[MerkleTreeId.NULLIFIER_TREE]),
+        treeStateReferenceToSnapshot(resp.state[MerkleTreeId.PUBLIC_DATA_TREE]),
+      ),
     );
-  } else {
-    throw new Error('Invalid leaf type');
   }
 }