@aztec/pxe 0.0.0-test.1 → 0.0.1-commit.1142ef1

package/src/storage/{address_data_provider/address_data_provider.ts → address_store/address_store.ts}

@@ -3,9 +3,7 @@ import type { AztecAsyncArray, AztecAsyncKVStore, AztecAsyncMap } from '@aztec/k
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
 import { CompleteAddress } from '@aztec/stdlib/contract';
 
-import type { DataProvider } from '../data_provider.js';
-
-export class AddressDataProvider implements DataProvider {
+export class AddressStore {
   #store: AztecAsyncKVStore;
   #completeAddresses: AztecAsyncArray<Buffer>;
   #completeAddressIndex: AztecAsyncMap<string, number>;
@@ -64,8 +62,4 @@ export class AddressDataProvider implements DataProvider {
       (await toArray(this.#completeAddresses.valuesAsync())).map(v => CompleteAddress.fromBuffer(v)),
     );
   }
-
-  async getSize(): Promise<number> {
-    return (await this.#completeAddresses.lengthAsync()) * CompleteAddress.SIZE_IN_BYTES;
-  }
 }

package/src/storage/address_store/index.ts

@@ -0,0 +1 @@
+export { AddressStore } from './address_store.js';

package/src/storage/{sync_data_provider/sync_data_provider.ts → anchor_block_store/anchor_block_store.ts}

@@ -1,9 +1,7 @@
 import type { AztecAsyncKVStore, AztecAsyncSingleton } from '@aztec/kv-store';
 import { BlockHeader } from '@aztec/stdlib/tx';
 
-import type { DataProvider } from '../data_provider.js';
-
-export class SyncDataProvider implements DataProvider {
+export class AnchorBlockStore {
   #store: AztecAsyncKVStore;
   #synchronizedHeader: AztecAsyncSingleton<Buffer>;
 
@@ -16,25 +14,12 @@ export class SyncDataProvider implements DataProvider {
     await this.#synchronizedHeader.set(header.toBuffer());
   }
 
-  async getBlockNumber(): Promise<number | undefined> {
-    const headerBuffer = await this.#synchronizedHeader.getAsync();
-    if (!headerBuffer) {
-      return undefined;
-    }
-
-    return Number(BlockHeader.fromBuffer(headerBuffer).globalVariables.blockNumber.toBigInt());
-  }
-
   async getBlockHeader(): Promise<BlockHeader> {
     const headerBuffer = await this.#synchronizedHeader.getAsync();
     if (!headerBuffer) {
-      throw new Error(`Header not set`);
+      throw new Error(`Trying to get block header with a not-yet-synchronized PXE - this should never happen`);
     }
 
     return BlockHeader.fromBuffer(headerBuffer);
   }
-
-  async getSize(): Promise<number> {
-    return (await this.#synchronizedHeader.getAsync())?.length ?? 0;
-  }
 }

package/src/storage/anchor_block_store/index.ts

@@ -0,0 +1 @@
+export { AnchorBlockStore } from './anchor_block_store.js';

package/src/storage/capsule_store/capsule_store.ts

@@ -0,0 +1,179 @@
+import { Fr } from '@aztec/foundation/curves/bn254';
+import { type Logger, createLogger } from '@aztec/foundation/log';
+import type { AztecAsyncKVStore, AztecAsyncMap } from '@aztec/kv-store';
+import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+
+export class CapsuleStore {
+  #store: AztecAsyncKVStore;
+
+  // Arbitrary data stored by contracts. Key is computed as `${contractAddress}:${key}`
+  #capsules: AztecAsyncMap<string, Buffer>;
+
+  logger: Logger;
+
+  constructor(store: AztecAsyncKVStore) {
+    this.#store = store;
+
+    this.#capsules = this.#store.openMap('capsules');
+
+    this.logger = createLogger('pxe:capsule-data-provider');
+  }
+
+  /**
+   * Stores arbitrary information in a per-contract non-volatile database, which can later be retrieved with `loadCapsule`.
+   * * If data was already stored at this slot, it is overwritten.
+   * @param contractAddress - The contract address to scope the data under.
+   * @param slot - The slot in the database in which to store the value. Slots need not be contiguous.
+   * @param capsule - An array of field elements representing the capsule.
+   * @remarks A capsule is a "blob" of data that is passed to the contract through an oracle. It works similarly
+   * to public contract storage in that it's indexed by the contract address and storage slot but instead of the global
+   * network state it's backed by local PXE db.
+   */
+  async storeCapsule(contractAddress: AztecAddress, slot: Fr, capsule: Fr[]): Promise<void> {
+    await this.#capsules.set(dbSlotToKey(contractAddress, slot), Buffer.concat(capsule.map(value => value.toBuffer())));
+  }
+
+  /**
+   * Returns data previously stored via `storeCapsule` in the per-contract non-volatile database.
+   * @param contractAddress - The contract address under which the data is scoped.
+   * @param slot - The slot in the database to read.
+   * @returns The stored data or `null` if no data is stored under the slot.
+   */
+  async loadCapsule(contractAddress: AztecAddress, slot: Fr): Promise<Fr[] | null> {
+    const dataBuffer = await this.#capsules.getAsync(dbSlotToKey(contractAddress, slot));
+    if (!dataBuffer) {
+      this.logger.trace(`Data not found for contract ${contractAddress.toString()} and slot ${slot.toString()}`);
+      return null;
+    }
+    const capsule: Fr[] = [];
+    for (let i = 0; i < dataBuffer.length; i += Fr.SIZE_IN_BYTES) {
+      capsule.push(Fr.fromBuffer(dataBuffer.subarray(i, i + Fr.SIZE_IN_BYTES)));
+    }
+    return capsule;
+  }
+
+  /**
+   * Deletes data in the per-contract non-volatile database. Does nothing if no data was present.
+   * @param contractAddress - The contract address under which the data is scoped.
+   * @param slot - The slot in the database to delete.
+   */
+  async deleteCapsule(contractAddress: AztecAddress, slot: Fr): Promise<void> {
+    await this.#capsules.delete(dbSlotToKey(contractAddress, slot));
+  }
+
+  /**
+   * Copies a number of contiguous entries in the per-contract non-volatile database. This allows for efficient data
+   * structures by avoiding repeated calls to `loadCapsule` and `storeCapsule`.
+   * Supports overlapping source and destination regions (which will result in the overlapped source values being
+   * overwritten). All copied slots must exist in the database (i.e. have been stored and not deleted)
+   *
+   * @param contractAddress - The contract address under which the data is scoped.
+   * @param srcSlot - The first slot to copy from.
+   * @param dstSlot - The first slot to copy to.
+   * @param numEntries - The number of entries to copy.
+   */
+  copyCapsule(contractAddress: AztecAddress, srcSlot: Fr, dstSlot: Fr, numEntries: number): Promise<void> {
+    return this.#store.transactionAsync(async () => {
+      // In order to support overlapping source and destination regions, we need to check the relative positions of source
+      // and destination. If destination is ahead of source, then by the time we overwrite source elements using forward
+      // indexes we'll have already read those. On the contrary, if source is ahead of destination we need to use backward
+      // indexes to avoid reading elements that've been overwritten.
+
+      const indexes = Array.from(Array(numEntries).keys());
+      if (srcSlot.lt(dstSlot)) {
+        indexes.reverse();
+      }
+
+      for (const i of indexes) {
+        const currentSrcSlot = dbSlotToKey(contractAddress, srcSlot.add(new Fr(i)));
+        const currentDstSlot = dbSlotToKey(contractAddress, dstSlot.add(new Fr(i)));
+
+        const toCopy = await this.#capsules.getAsync(currentSrcSlot);
+        if (!toCopy) {
+          throw new Error(`Attempted to copy empty slot ${currentSrcSlot} for contract ${contractAddress.toString()}`);
+        }
+
+        await this.#capsules.set(currentDstSlot, toCopy);
+      }
+    });
+  }
+
+  /**
+   * Appends multiple capsules to a capsule array stored at the base slot.
+   * The array length is stored at the base slot, and elements are stored in consecutive slots after it.
+   * All operations are performed in a single transaction.
+   * @param contractAddress - The contract address that owns the capsule array
+   * @param baseSlot - The slot where the array length is stored
+   * @param content - Array of capsule data to append
+   */
+  appendToCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr, content: Fr[][]): Promise<void> {
+    return this.#store.transactionAsync(async () => {
+      // Load current length, defaulting to 0 if not found
+      const lengthData = await this.loadCapsule(contractAddress, baseSlot);
+      const currentLength = lengthData ? lengthData[0].toNumber() : 0;
+
+      // Store each capsule at consecutive slots after baseSlot + 1 + currentLength
+      for (let i = 0; i < content.length; i++) {
+        const nextSlot = arraySlot(baseSlot, currentLength + i);
+        await this.storeCapsule(contractAddress, nextSlot, content[i]);
+      }
+
+      // Update length to include all new capsules
+      const newLength = currentLength + content.length;
+      await this.storeCapsule(contractAddress, baseSlot, [new Fr(newLength)]);
+    });
+  }
+
+  readCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr): Promise<Fr[][]> {
+    return this.#store.transactionAsync(async () => {
+      // Load length, defaulting to 0 if not found
+      const maybeLength = await this.loadCapsule(contractAddress, baseSlot);
+      const length = maybeLength ? maybeLength[0].toBigInt() : 0n;
+
+      const values: Fr[][] = [];
+
+      // Read each capsule at consecutive slots after baseSlot
+      for (let i = 0; i < length; i++) {
+        const currentValue = await this.loadCapsule(contractAddress, arraySlot(baseSlot, i));
+        if (currentValue == undefined) {
+          throw new Error(
+            `Expected non-empty value at capsule array in base slot ${baseSlot} at index ${i} for contract ${contractAddress}`,
+          );
+        }
+
+        values.push(currentValue);
+      }
+
+      return values;
+    });
+  }
+
+  setCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr, content: Fr[][]) {
+    return this.#store.transactionAsync(async () => {
+      // Load current length, defaulting to 0 if not found
+      const maybeLength = await this.loadCapsule(contractAddress, baseSlot);
+      const originalLength = maybeLength ? maybeLength[0].toNumber() : 0;
+
+      // Set the new length
+      await this.storeCapsule(contractAddress, baseSlot, [new Fr(content.length)]);
+
+      // Store the new content, possibly overwriting existing values
+      for (let i = 0; i < content.length; i++) {
+        await this.storeCapsule(contractAddress, arraySlot(baseSlot, i), content[i]);
+      }
+
+      // Clear any stragglers
+      for (let i = content.length; i < originalLength; i++) {
+        await this.deleteCapsule(contractAddress, arraySlot(baseSlot, i));
+      }
+    });
+  }
+}
+
+function dbSlotToKey(contractAddress: AztecAddress, slot: Fr): string {
+  return `${contractAddress.toString()}:${slot.toString()}`;
+}
+
+function arraySlot(baseSlot: Fr, index: number) {
+  return baseSlot.add(new Fr(1)).add(new Fr(index));
+}

package/src/storage/capsule_store/index.ts

@@ -0,0 +1 @@
+export { CapsuleStore } from './capsule_store.js';

package/src/storage/contract_store/contract_store.ts

@@ -0,0 +1,339 @@
+import type { FUNCTION_TREE_HEIGHT } from '@aztec/constants';
+import type { Fr } from '@aztec/foundation/curves/bn254';
+import { toArray } from '@aztec/foundation/iterable';
+import type { MembershipWitness } from '@aztec/foundation/trees';
+import type { AztecAsyncKVStore, AztecAsyncMap } from '@aztec/kv-store';
+import { isProtocolContract } from '@aztec/protocol-contracts';
+import {
+  type ContractArtifact,
+  type FunctionAbi,
+  type FunctionArtifact,
+  type FunctionArtifactWithContractName,
+  FunctionCall,
+  type FunctionDebugMetadata,
+  FunctionSelector,
+  FunctionType,
+  contractArtifactFromBuffer,
+  contractArtifactToBuffer,
+  encodeArguments,
+  getFunctionDebugMetadata,
+} from '@aztec/stdlib/abi';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+import {
+  type ContractClass,
+  type ContractInstanceWithAddress,
+  SerializableContractInstance,
+  getContractClassFromArtifact,
+} from '@aztec/stdlib/contract';
+
+import { PrivateFunctionsTree } from './private_functions_tree.js';
+
+/**
+ * ContractStore serves as a data manager and retriever for Aztec.nr contracts.
+ * It provides methods to obtain contract addresses, function ABI, bytecode, and membership witnesses
+ * from a given contract address and function selector. The class maintains a cache of ContractTree instances
+ * to efficiently serve the requested data. It interacts with the ContractDatabase and AztecNode to fetch
+ * the required information and facilitate cryptographic proof generation.
+ */
+export class ContractStore {
+  /** Map from contract class id to private function tree. */
+  // TODO: Update it to be LRU cache so that it doesn't keep all the data all the time.
+  #privateFunctionTrees: Map<string, PrivateFunctionsTree> = new Map();
+
+  /** Map from contract address to contract class id */
+  #contractClassIdMap: Map<string, Fr> = new Map();
+
+  #contractArtifacts: AztecAsyncMap<string, Buffer>;
+  #contractInstances: AztecAsyncMap<string, Buffer>;
+
+  constructor(store: AztecAsyncKVStore) {
+    this.#contractArtifacts = store.openMap('contract_artifacts');
+    this.#contractInstances = store.openMap('contracts_instances');
+  }
+
+  // Setters
+
+  public async addContractArtifact(id: Fr, contract: ContractArtifact): Promise<void> {
+    const privateFunctions = contract.functions.filter(
+      functionArtifact => functionArtifact.functionType === FunctionType.PRIVATE,
+    );
+
+    const privateSelectors = await Promise.all(
+      privateFunctions.map(async privateFunctionArtifact =>
+        (
+          await FunctionSelector.fromNameAndParameters(privateFunctionArtifact.name, privateFunctionArtifact.parameters)
+        ).toString(),
+      ),
+    );
+
+    if (privateSelectors.length !== new Set(privateSelectors).size) {
+      throw new Error('Repeated function selectors of private functions');
+    }
+
+    await this.#contractArtifacts.set(id.toString(), contractArtifactToBuffer(contract));
+  }
+
+  async addContractInstance(contract: ContractInstanceWithAddress): Promise<void> {
+    this.#contractClassIdMap.set(contract.address.toString(), contract.currentContractClassId);
+
+    await this.#contractInstances.set(
+      contract.address.toString(),
+      new SerializableContractInstance(contract).toBuffer(),
+    );
+  }
+
+  // Private getters
+
+  async #getContractClassId(contractAddress: AztecAddress): Promise<Fr | undefined> {
+    if (!this.#contractClassIdMap.has(contractAddress.toString())) {
+      const instance = await this.getContractInstance(contractAddress);
+      if (!instance) {
+        return;
+      }
+      this.#contractClassIdMap.set(contractAddress.toString(), instance.currentContractClassId);
+    }
+    return this.#contractClassIdMap.get(contractAddress.toString());
+  }
+
+  /**
+   * Retrieve or create a ContractTree instance based on the provided class id.
+   * If an existing tree with the same class id is found in the cache, it will be returned.
+   * Otherwise, a new ContractTree instance will be created using the contract data from the database
+   * and added to the cache before returning.
+   *
+   * @param classId - The class id of the contract for which the ContractTree is required.
+   * @returns A ContractTree instance associated with the specified contract address.
+   * @throws An Error if the contract is not found in the ContractDatabase.
+   */
+  async #getPrivateFunctionTreeForClassId(classId: Fr): Promise<PrivateFunctionsTree | undefined> {
+    if (!this.#privateFunctionTrees.has(classId.toString())) {
+      const artifact = await this.getContractArtifact(classId);
+      if (!artifact) {
+        return;
+      }
+      const tree = await PrivateFunctionsTree.create(artifact);
+      this.#privateFunctionTrees.set(classId.toString(), tree);
+    }
+    return this.#privateFunctionTrees.get(classId.toString())!;
+  }
+
+  async #getContractArtifactByAddress(contractAddress: AztecAddress): Promise<ContractArtifact | undefined> {
+    const contractClassId = await this.#getContractClassId(contractAddress);
+    return contractClassId && this.getContractArtifact(contractClassId);
+  }
+
+  // Public getters
+
+  async getContractsAddresses(): Promise<AztecAddress[]> {
+    const keys = await toArray(this.#contractInstances.keysAsync());
+    return keys.map(AztecAddress.fromString);
+  }
+
+  /** Returns a contract instance for a given address. Throws if not found. */
+  public async getContractInstance(contractAddress: AztecAddress): Promise<ContractInstanceWithAddress | undefined> {
+    const contract = await this.#contractInstances.getAsync(contractAddress.toString());
+    return contract && SerializableContractInstance.fromBuffer(contract).withAddress(contractAddress);
+  }
+
+  public async getContractArtifact(contractClassId: Fr): Promise<ContractArtifact | undefined> {
+    const contract = await this.#contractArtifacts.getAsync(contractClassId.toString());
+    // TODO(@spalladino): AztecAsyncMap lies and returns Uint8Arrays instead of Buffers, hence the extra Buffer.from.
+    return contract && contractArtifactFromBuffer(Buffer.from(contract));
+  }
+
+  /** Returns a contract class for a given class id. Throws if not found. */
+  public async getContractClass(contractClassId: Fr): Promise<ContractClass | undefined> {
+    const artifact = await this.getContractArtifact(contractClassId);
+    return artifact && getContractClassFromArtifact(artifact);
+  }
+
+  public async getContract(
+    address: AztecAddress,
+  ): Promise<(ContractInstanceWithAddress & ContractArtifact) | undefined> {
+    const instance = await this.getContractInstance(address);
+    if (!instance) {
+      return;
+    }
+    const artifact = await this.getContractArtifact(instance.currentContractClassId);
+    if (!artifact) {
+      return;
+    }
+    return { ...instance, ...artifact };
+  }
+
+  /**
+   * Retrieves the artifact of a specified function within a given contract.
+   * The function is identified by its selector, which is a unique code generated from the function's signature.
+   * Throws an error if the contract address or function selector are invalid or not found.
+   *
+   * @param contractAddress - The AztecAddress representing the contract containing the function.
+   * @param selector - The function selector.
+   * @returns The corresponding function's artifact as an object.
+   */
+  public async getFunctionArtifact(
+    contractAddress: AztecAddress,
+    selector: FunctionSelector,
+  ): Promise<FunctionArtifactWithContractName | undefined> {
+    const artifact = await this.#getContractArtifactByAddress(contractAddress);
+    const fnArtifact = artifact && (await this.#findFunctionArtifactBySelector(artifact, selector));
+    return fnArtifact && { ...fnArtifact, contractName: artifact.name };
+  }
+
+  public async getFunctionArtifactWithDebugMetadata(
+    contractAddress: AztecAddress,
+    selector: FunctionSelector,
+  ): Promise<FunctionArtifactWithContractName> {
+    const artifact = await this.getFunctionArtifact(contractAddress, selector);
+    if (!artifact) {
+      throw new Error(`Function artifact not found for contract ${contractAddress} and selector ${selector}.`);
+    }
+    const debug = await this.getFunctionDebugMetadata(contractAddress, selector);
+    return {
+      ...artifact,
+      debug,
+    };
+  }
+
+  public async getPublicFunctionArtifact(
+    contractAddress: AztecAddress,
+  ): Promise<FunctionArtifactWithContractName | undefined> {
+    const artifact = await this.#getContractArtifactByAddress(contractAddress);
+    const fnArtifact = artifact && artifact.functions.find(fn => fn.functionType === FunctionType.PUBLIC);
+    return fnArtifact && { ...fnArtifact, contractName: artifact.name };
+  }
+
+  public async getFunctionAbi(
+    contractAddress: AztecAddress,
+    selector: FunctionSelector,
+  ): Promise<FunctionAbi | undefined> {
+    const artifact = await this.#getContractArtifactByAddress(contractAddress);
+    return artifact && (await this.#findFunctionAbiBySelector(artifact, selector));
+  }
+
+  /**
+   * Retrieves the debug metadata of a specified function within a given contract.
+   * The function is identified by its selector, which is a unique code generated from the function's signature.
+   * Returns undefined if the debug metadata for the given function is not found.
+   * Throws if the contract has not been added to the database.
+   *
+   * @param contractAddress - The AztecAddress representing the contract containing the function.
+   * @param selector - The function selector.
+   * @returns The corresponding function's artifact as an object.
+   */
+  public async getFunctionDebugMetadata(
+    contractAddress: AztecAddress,
+    selector: FunctionSelector,
+  ): Promise<FunctionDebugMetadata | undefined> {
+    const artifact = await this.#getContractArtifactByAddress(contractAddress);
+    const fnArtifact = artifact && (await this.#findFunctionArtifactBySelector(artifact, selector));
+    return fnArtifact && getFunctionDebugMetadata(artifact, fnArtifact);
+  }
+
+  public async getPublicFunctionDebugMetadata(
+    contractAddress: AztecAddress,
+  ): Promise<FunctionDebugMetadata | undefined> {
+    const artifact = await this.#getContractArtifactByAddress(contractAddress);
+    const fnArtifact = artifact && artifact.functions.find(fn => fn.functionType === FunctionType.PUBLIC);
+    return fnArtifact && getFunctionDebugMetadata(artifact, fnArtifact);
+  }
+
+  /**
+   * Retrieve the function membership witness for the given contract class and function selector.
+   * The function membership witness represents a proof that the function belongs to the specified contract.
+   * Throws an error if the contract address or function selector is unknown.
+   *
+   * @param contractClassId - The id of the class.
+   * @param selector - The function selector.
+   * @returns A promise that resolves with the MembershipWitness instance for the specified contract's function.
+   */
+  public async getFunctionMembershipWitness(
+    contractClassId: Fr,
+    selector: FunctionSelector,
+  ): Promise<MembershipWitness<typeof FUNCTION_TREE_HEIGHT> | undefined> {
+    const tree = await this.#getPrivateFunctionTreeForClassId(contractClassId);
+    return tree?.getFunctionMembershipWitness(selector);
+  }
+
+  public async getDebugContractName(contractAddress: AztecAddress) {
+    const artifact = await this.#getContractArtifactByAddress(contractAddress);
+    return artifact?.name;
+  }
+
+  public async getDebugFunctionName(contractAddress: AztecAddress, selector: FunctionSelector) {
+    const artifact = await this.#getContractArtifactByAddress(contractAddress);
+    const fnArtifact = artifact && (await this.#findFunctionAbiBySelector(artifact, selector));
+    return `${artifact?.name ?? contractAddress}:${fnArtifact?.name ?? selector}`;
+  }
+
+  async #findFunctionArtifactBySelector(
+    artifact: ContractArtifact,
+    selector: FunctionSelector,
+  ): Promise<FunctionArtifact | undefined> {
+    const functions = artifact.functions;
+    for (let i = 0; i < functions.length; i++) {
+      const fn = functions[i];
+      const fnSelector = await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters);
+      if (fnSelector.equals(selector)) {
+        return fn;
+      }
+    }
+  }
+
+  async #findFunctionAbiBySelector(
+    artifact: ContractArtifact,
+    selector: FunctionSelector,
+  ): Promise<FunctionAbi | undefined> {
+    const functions = [...artifact.functions, ...(artifact.nonDispatchPublicFunctions ?? [])];
+    for (let i = 0; i < functions.length; i++) {
+      const fn = functions[i];
+      const fnSelector = await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters);
+      if (fnSelector.equals(selector)) {
+        return fn;
+      }
+    }
+  }
+
+  public async getFunctionCall(functionName: string, args: any[], to: AztecAddress): Promise<FunctionCall> {
+    const contract = await this.getContract(to);
+    if (!contract) {
+      throw new Error(
+        `Unknown contract ${to}: add it to PXE by calling server.addContracts(...).\nSee docs for context: https://docs.aztec.network/developers/resources/debugging/aztecnr-errors#unknown-contract-0x0-add-it-to-pxe-by-calling-serveraddcontracts`,
+      );
+    }
+
+    const functionDao = contract.functions.find(f => f.name === functionName);
+    if (!functionDao) {
+      throw new Error(`Unknown function ${functionName} in contract ${contract.name}.`);
+    }
+
+    return {
+      name: functionDao.name,
+      args: encodeArguments(functionDao, args),
+      selector: await FunctionSelector.fromNameAndParameters(functionDao.name, functionDao.parameters),
+      type: functionDao.functionType,
+      to,
+      hideMsgSender: false,
+      isStatic: functionDao.isStatic,
+      returnTypes: functionDao.returnTypes,
+    };
+  }
+
+  // Synchronize target contract data
+  public async syncPrivateState(
+    contractAddress: AztecAddress,
+    functionToInvokeAfterSync: FunctionSelector | null,
+    utilityExecutor: (privateSyncCall: FunctionCall) => Promise<any>,
+  ) {
+    // Protocol contracts don't have private state to sync
+    if (!isProtocolContract(contractAddress)) {
+      const syncPrivateStateFunctionCall = await this.getFunctionCall('sync_private_state', [], contractAddress);
+      if (functionToInvokeAfterSync && functionToInvokeAfterSync.equals(syncPrivateStateFunctionCall.selector)) {
+        throw new Error(
+          'Forbidden `sync_private_state` invocation. `sync_private_state` can only be invoked by PXE, manual execution can lead to inconsistencies.',
+        );
+      }
+
+      return utilityExecutor(syncPrivateStateFunctionCall);
+    }
+  }
+}

package/src/storage/contract_store/index.ts

@@ -0,0 +1 @@
+export { ContractStore } from './contract_store.js';

package/src/storage/contract_store/private_functions_tree.ts

@@ -0,0 +1,67 @@
+import { FUNCTION_TREE_HEIGHT } from '@aztec/constants';
+import { Fr } from '@aztec/foundation/curves/bn254';
+import { assertLength } from '@aztec/foundation/serialize';
+import { MembershipWitness, type MerkleTree } from '@aztec/foundation/trees';
+import { type ContractArtifact, FunctionSelector, FunctionType } from '@aztec/stdlib/abi';
+import {
+  type PrivateFunction,
+  computePrivateFunctionLeaf,
+  computePrivateFunctionsTree,
+  getContractClassPrivateFunctionFromArtifact,
+} from '@aztec/stdlib/contract';
+
+/**
+ * Represents a Merkle tree of functions for a particular Contract Class.
+ * It manages the construction of the function tree, computes its root, and generates membership witnesses
+ * for constrained functions. This class also enables lookup of specific function artifact using selectors.
+ * It is used in combination with the AztecNode to compute various data for executing private transactions.
+ */
+export class PrivateFunctionsTree {
+  private tree?: MerkleTree;
+
+  private constructor(private readonly privateFunctions: PrivateFunction[]) {}
+
+  static async create(artifact: ContractArtifact) {
+    const privateFunctions = await Promise.all(
+      artifact.functions
+        .filter(fn => fn.functionType === FunctionType.PRIVATE)
+        .map(getContractClassPrivateFunctionFromArtifact),
+    );
+    return new PrivateFunctionsTree(privateFunctions);
+  }
+
+  /**
+   * Retrieve the membership witness of a function within a contract's function tree.
+   * A membership witness represents the position and authentication path of a target function
+   * in the Merkle tree of constrained functions. It is required to prove the existence of the
+   * function within the contract during execution. Throws if fn does not exist or is not private.
+   *
+   * @param selector - The function selector.
+   * @returns A MembershipWitness instance representing the position and authentication path of the function in the function tree.
+   */
+  public async getFunctionMembershipWitness(
+    selector: FunctionSelector,
+  ): Promise<MembershipWitness<typeof FUNCTION_TREE_HEIGHT>> {
+    const fn = this.privateFunctions.find(f => f.selector.equals(selector));
+    if (!fn) {
+      throw new Error(`Private function with selector ${selector.toString()} not found in contract class.`);
+    }
+
+    const leaf = await computePrivateFunctionLeaf(fn);
+    const tree = await this.getTree();
+    const index = tree.getIndex(leaf);
+    const path = tree.getSiblingPath(index);
+    return new MembershipWitness<typeof FUNCTION_TREE_HEIGHT>(
+      FUNCTION_TREE_HEIGHT,
+      BigInt(index),
+      assertLength(path.map(Fr.fromBuffer), FUNCTION_TREE_HEIGHT),
+    );
+  }
+
+  private async getTree() {
+    if (!this.tree) {
+      this.tree = await computePrivateFunctionsTree(this.privateFunctions);
+    }
+    return this.tree;
+  }
+}

package/src/storage/index.ts

@@ -1,10 +1,8 @@
-export * from './address_data_provider/index.js';
-export * from './auth_witness_data_provider/index.js';
-export * from './capsule_data_provider/index.js';
-export * from './contract_data_provider/index.js';
-export * from './note_data_provider/index.js';
-export * from './sync_data_provider/index.js';
-export * from './tagging_data_provider/index.js';
-export * from './data_provider.js';
-
-export const PXE_DATA_SCHEMA_VERSION = 2;
+export * from './address_store/index.js';
+export * from './capsule_store/index.js';
+export * from './contract_store/index.js';
+export * from './note_store/index.js';
+export * from './anchor_block_store/index.js';
+export * from './tagging_store/index.js';
+export * from './metadata.js';
+export * from './private_event_store/private_event_store.js';

package/src/storage/metadata.ts

@@ -0,0 +1 @@
+export const PXE_DATA_SCHEMA_VERSION = 2;