@aztec/txe 0.0.0-test.1 → 0.0.1-commit.023c3e5

package/src/index.ts

@@ -1,25 +1,27 @@
 import { SchnorrAccountContractArtifact } from '@aztec/accounts/schnorr';
+import { type NoirCompiledContract, loadContractArtifact } from '@aztec/aztec.js/abi';
+import { AztecAddress } from '@aztec/aztec.js/addresses';
 import {
-  AztecAddress,
-  type ContractArtifact,
   type ContractInstanceWithAddress,
-  Fr,
-  PublicKeys,
-  deriveKeys,
-  getContractInstanceFromDeployParams,
-  loadContractArtifact,
-} from '@aztec/aztec.js';
+  getContractInstanceFromInstantiationParams,
+} from '@aztec/aztec.js/contracts';
+import { Fr } from '@aztec/aztec.js/fields';
+import { PublicKeys, deriveKeys } from '@aztec/aztec.js/keys';
 import { createSafeJsonRpcServer } from '@aztec/foundation/json-rpc/server';
 import type { Logger } from '@aztec/foundation/log';
 import { type ProtocolContract, protocolContractNames } from '@aztec/protocol-contracts';
 import { BundledProtocolContractsProvider } from '@aztec/protocol-contracts/providers/bundle';
-import type { ApiSchemaFor, ZodFor } from '@aztec/stdlib/schemas';
+import { computeArtifactHash } from '@aztec/stdlib/contract';
+import type { ApiSchemaFor } from '@aztec/stdlib/schemas';
+import { zodFor } from '@aztec/stdlib/schemas';
 
+import { createHash } from 'crypto';
+import { createReadStream } from 'fs';
 import { readFile, readdir } from 'fs/promises';
-import { join } from 'path';
+import { join, parse } from 'path';
 import { z } from 'zod';
 
-import { TXEService } from './txe_service/txe_service.js';
+import { type TXEOracleFunctionName, TXESession } from './txe_session.js';
 import {
   type ForeignCallArgs,
   ForeignCallArgsSchema,
@@ -29,94 +31,125 @@ import {
   type ForeignCallSingle,
   fromArray,
   fromSingle,
-  toForeignCallResult,
   toSingle,
 } from './util/encoding.js';
+import type { ContractArtifactWithHash } from './util/txe_contract_store.js';
 
-const TXESessions = new Map<number, TXEService>();
+const sessions = new Map<number, TXESession>();
 
-const TXEArtifactsCache = new Map<string, { artifact: ContractArtifact; instance: ContractInstanceWithAddress }>();
-
-type MethodNames<T> = {
-  [K in keyof T]: T[K] extends (...args: any[]) => any ? K : never;
-}[keyof T];
+/*
+ * TXE typically has to load the same contract artifacts over and over again for multiple tests,
+ * so we cache them here to avoid both loading them from disk repeatedly and computing their artifact hashes
+ */
+const TXEArtifactsCache = new Map<
+  string,
+  { artifact: ContractArtifactWithHash; instance: ContractInstanceWithAddress }
+>();
 
 type TXEForeignCallInput = {
   session_id: number;
-  function: MethodNames<TXEService> | 'reset';
+  function: TXEOracleFunctionName;
   root_path: string;
   package_name: string;
   inputs: ForeignCallArgs;
 };
 
-const TXEForeignCallInputSchema = z.object({
-  // eslint-disable-next-line camelcase
-  session_id: z.number().int().nonnegative(),
-  function: z.string() as ZodFor<MethodNames<TXEService> | 'reset'>,
-  // eslint-disable-next-line camelcase
-  root_path: z.string(),
-  // eslint-disable-next-line camelcase
-  package_name: z.string(),
-  inputs: ForeignCallArgsSchema,
-}) satisfies ZodFor<TXEForeignCallInput>;
+const TXEForeignCallInputSchema = zodFor<TXEForeignCallInput>()(
+  z.object({
+    // eslint-disable-next-line camelcase
+    session_id: z.number().int().nonnegative(),
+    function: z.string() as z.ZodType<TXEOracleFunctionName>,
+    // eslint-disable-next-line camelcase
+    root_path: z.string(),
+    // eslint-disable-next-line camelcase
+    package_name: z.string(),
+    inputs: ForeignCallArgsSchema,
+  }),
+);
 
 class TXEDispatcher {
   private protocolContracts!: ProtocolContract[];
 
   constructor(private logger: Logger) {}
 
+  private fastHashFile(path: string) {
+    return new Promise(resolve => {
+      const fd = createReadStream(path);
+      const hash = createHash('sha1');
+      hash.setEncoding('hex');
+
+      fd.on('end', function () {
+        hash.end();
+        resolve(hash.read());
+      });
+
+      fd.pipe(hash);
+    });
+  }
+
   async #processDeployInputs({ inputs, root_path: rootPath, package_name: packageName }: TXEForeignCallInput) {
-    const [pathStr, contractName, initializer] = inputs.slice(0, 3).map(input =>
+    const [contractPath, initializer] = inputs.slice(0, 2).map(input =>
       fromArray(input as ForeignCallArray)
         .map(char => String.fromCharCode(char.toNumber()))
         .join(''),
     );
 
-    const decodedArgs = fromArray(inputs[4] as ForeignCallArray);
-    const secret = fromSingle(inputs[5] as ForeignCallSingle);
+    const decodedArgs = fromArray(inputs[3] as ForeignCallArray);
+    const secret = fromSingle(inputs[4] as ForeignCallSingle);
     const publicKeys = secret.equals(Fr.ZERO) ? PublicKeys.default() : (await deriveKeys(secret)).publicKeys;
     const publicKeysHash = await publicKeys.hash();
 
-    const cacheKey = `${pathStr}-${contractName}-${initializer}-${decodedArgs
+    let artifactPath = '';
+    const { dir: contractDirectory, base: contractFilename } = parse(contractPath);
+    if (contractDirectory) {
+      if (contractDirectory.includes('@')) {
+        // We're deploying a contract that belongs in a workspace
+        // env.deploy("../path/to/workspace/root@packageName/contractName")
+        const [workspace, pkg] = contractDirectory.split('@');
+        const targetPath = join(rootPath, workspace, '/target');
+        this.logger.debug(`Looking for compiled artifact in workspace ${targetPath}`);
+        artifactPath = join(targetPath, `${pkg}-${contractFilename}.json`);
+      } else {
+        // We're deploying a standalone external contract
+        // env.deploy("../path/to/contract/root/contractName")
+        const targetPath = join(rootPath, contractDirectory, '/target');
+        this.logger.debug(`Looking for compiled artifact in ${targetPath}`);
+        [artifactPath] = (await readdir(targetPath)).filter(file => file.endsWith(`-${contractFilename}.json`));
+      }
+    } else {
+      // We're deploying a local contract
+      // env.deploy("contractName")
+      artifactPath = join(rootPath, './target', `${packageName}-${contractFilename}.json`);
+    }
+
+    const fileHash = await this.fastHashFile(artifactPath);
+
+    const cacheKey = `${contractDirectory ?? ''}-${contractFilename}-${initializer}-${decodedArgs
       .map(arg => arg.toString())
-      .join('-')}-${publicKeysHash.toString()}`;
+      .join('-')}-${publicKeysHash}-${fileHash}`;
 
-    let artifact;
     let instance;
+    let artifact: ContractArtifactWithHash;
 
     if (TXEArtifactsCache.has(cacheKey)) {
       this.logger.debug(`Using cached artifact for ${cacheKey}`);
       ({ artifact, instance } = TXEArtifactsCache.get(cacheKey)!);
     } else {
-      let artifactPath = '';
-      // We're deploying the contract under test
-      // env.deploy_self("contractName")
-      if (!pathStr) {
-        artifactPath = join(rootPath, './target', `${packageName}-${contractName}.json`);
-      } else {
-        // We're deploying a contract that belongs in a workspace
-        // env.deploy("../path/to/workspace/root@packageName", "contractName")
-        if (pathStr.includes('@')) {
-          const [workspace, pkg] = pathStr.split('@');
-          const targetPath = join(rootPath, workspace, './target');
-          this.logger.debug(`Looking for compiled artifact in workspace ${targetPath}`);
-          artifactPath = join(targetPath, `${pkg}-${contractName}.json`);
-        } else {
-          // We're deploying a standalone contract
-          // env.deploy("../path/to/contract/root", "contractName")
-          const targetPath = join(rootPath, pathStr, './target');
-          this.logger.debug(`Looking for compiled artifact in ${targetPath}`);
-          [artifactPath] = (await readdir(targetPath)).filter(file => file.endsWith(`-${contractName}.json`));
-        }
-      }
       this.logger.debug(`Loading compiled artifact ${artifactPath}`);
-      artifact = loadContractArtifact(JSON.parse(await readFile(artifactPath, 'utf-8')));
+      const artifactJSON = JSON.parse(await readFile(artifactPath, 'utf-8')) as NoirCompiledContract;
+      const artifactWithoutHash = loadContractArtifact(artifactJSON);
+      artifact = {
+        ...artifactWithoutHash,
+        // Artifact hash is *very* expensive to compute, so we do it here once
+        // and the TXE contract data provider can cache it
+        artifactHash: await computeArtifactHash(artifactWithoutHash),
+      };
       this.logger.debug(
         `Deploy ${
           artifact.name
         } with initializer ${initializer}(${decodedArgs}) and public keys hash ${publicKeysHash.toString()}`,
       );
-      instance = await getContractInstanceFromDeployParams(artifact, {
+      instance = await getContractInstanceFromInstantiationParams(artifact, {
         constructorArgs: decodedArgs,
         skipArgsDecoding: true,
         salt: Fr.ONE,
@@ -127,7 +160,7 @@ class TXEDispatcher {
       TXEArtifactsCache.set(cacheKey, { artifact, instance });
     }
 
-    inputs.splice(0, 2, artifact, instance, toSingle(secret));
+    inputs.splice(0, 1, artifact, instance, toSingle(secret));
   }
 
   async #processAddAccountInputs({ inputs }: TXEForeignCallInput) {
@@ -135,7 +168,7 @@ class TXEDispatcher {
 
     const cacheKey = `SchnorrAccountContract-${secret}`;
 
-    let artifact;
+    let artifact: ContractArtifactWithHash;
     let instance;
 
     if (TXEArtifactsCache.has(cacheKey)) {
@@ -144,8 +177,13 @@ class TXEDispatcher {
     } else {
       const keys = await deriveKeys(secret);
       const args = [keys.publicKeys.masterIncomingViewingPublicKey.x, keys.publicKeys.masterIncomingViewingPublicKey.y];
-      artifact = SchnorrAccountContractArtifact;
-      instance = await getContractInstanceFromDeployParams(artifact, {
+      artifact = {
+        ...SchnorrAccountContractArtifact,
+        // Artifact hash is *very* expensive to compute, so we do it here once
+        // and the TXE contract data provider can cache it
+        artifactHash: await computeArtifactHash(SchnorrAccountContractArtifact),
+      };
+      instance = await getContractInstanceFromInstantiationParams(artifact, {
         constructorArgs: args,
         skipArgsDecoding: true,
         salt: Fr.ONE,
@@ -164,35 +202,28 @@ class TXEDispatcher {
     const { session_id: sessionId, function: functionName, inputs } = callData;
     this.logger.debug(`Calling ${functionName} on session ${sessionId}`);
 
-    if (!TXESessions.has(sessionId) && functionName != 'reset') {
+    if (!sessions.has(sessionId)) {
       this.logger.debug(`Creating new session ${sessionId}`);
       if (!this.protocolContracts) {
         this.protocolContracts = await Promise.all(
           protocolContractNames.map(name => new BundledProtocolContractsProvider().getProtocolContractArtifact(name)),
         );
       }
-      TXESessions.set(sessionId, await TXEService.init(this.logger, this.protocolContracts));
+      sessions.set(sessionId, await TXESession.init(this.protocolContracts));
     }
 
     switch (functionName) {
-      case 'reset': {
-        TXESessions.delete(sessionId) &&
-          this.logger.debug(`Called reset on session ${sessionId}, yeeting it out of existence`);
-        return toForeignCallResult([]);
-      }
-      case 'deploy': {
+      case 'txeDeploy': {
         await this.#processDeployInputs(callData);
         break;
       }
-      case 'addAccount': {
+      case 'txeAddAccount': {
         await this.#processAddAccountInputs(callData);
         break;
       }
     }
 
-    const txeService = TXESessions.get(sessionId);
-    const response = await (txeService as any)[functionName](...inputs);
-    return response;
+    return await sessions.get(sessionId)!.processFunction(functionName, inputs);
   }
 }
 

package/src/oracle/interfaces.ts

@@ -0,0 +1,86 @@
+import type { ContractArtifact } from '@aztec/aztec.js/abi';
+import { CompleteAddress } from '@aztec/aztec.js/addresses';
+import type { ContractInstanceWithAddress } from '@aztec/aztec.js/contracts';
+import { TxHash } from '@aztec/aztec.js/tx';
+import { BlockNumber } from '@aztec/foundation/branded-types';
+import type { Fr } from '@aztec/foundation/curves/bn254';
+import type { EventSelector, FunctionSelector } from '@aztec/stdlib/abi';
+import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+import type { UInt64 } from '@aztec/stdlib/types';
+
+// These interfaces complement the ones defined in PXE, and combined with those contain the full list of oracles used by
+// aztec-nr. In particular, these include the ones needed to run Brillig code associated to #[external("public")] functions that has
+// not been transpiled (e.g. in the context of a Noir test) as well as the ones associated with managing the state of
+// such a Noir test (deploying contracts, manipulating block time, making calls, etc) - the so called 'top level test
+// context'.
+
+/**
+ * Oracle methods associated with the execution of an Aztec #[external("public")] function.
+ *
+ * Note that real contracts have their Brillig calls to these be transpiled into opcodes, the oracles are only executed
+ * as such when running the original Brillig code, e.g. when invoking functions that interact with a PublicContext
+ * directly in a Noir test.
+ */
+export interface IAvmExecutionOracle {
+  isAvm: true;
+
+  avmOpcodeAddress(): Promise<AztecAddress>;
+  avmOpcodeSender(): Promise<AztecAddress>;
+  avmOpcodeBlockNumber(): Promise<BlockNumber>;
+  avmOpcodeTimestamp(): Promise<bigint>;
+  avmOpcodeIsStaticCall(): Promise<boolean>;
+  avmOpcodeChainId(): Promise<Fr>;
+  avmOpcodeVersion(): Promise<Fr>;
+  avmOpcodeEmitNullifier(nullifier: Fr): Promise<void>;
+  avmOpcodeEmitNoteHash(noteHash: Fr): Promise<void>;
+  avmOpcodeNullifierExists(siloedNullifier: Fr): Promise<boolean>;
+  avmOpcodeStorageWrite(slot: Fr, value: Fr): Promise<void>;
+  avmOpcodeStorageRead(slot: Fr, contractAddress: AztecAddress): Promise<Fr>;
+}
+
+/**
+ * Oracle methods associated with the execution of an Aztec Noir test.
+ */
+export interface ITxeExecutionOracle {
+  isTxe: true;
+
+  txeGetDefaultAddress(): AztecAddress;
+  txeGetNextBlockNumber(): Promise<BlockNumber>;
+  txeGetNextBlockTimestamp(): Promise<UInt64>;
+  txeAdvanceBlocksBy(blocks: number): Promise<void>;
+  txeAdvanceTimestampBy(duration: UInt64): void;
+  txeDeploy(artifact: ContractArtifact, instance: ContractInstanceWithAddress, foreignSecret: Fr): Promise<void>;
+  txeCreateAccount(secret: Fr): Promise<CompleteAddress>;
+  txeAddAccount(
+    artifact: ContractArtifact,
+    instance: ContractInstanceWithAddress,
+    secret: Fr,
+  ): Promise<CompleteAddress>;
+  txeAddAuthWitness(address: AztecAddress, messageHash: Fr): Promise<void>;
+  txeGetLastBlockTimestamp(): Promise<bigint>;
+  txeGetLastTxEffects(): Promise<{
+    txHash: TxHash;
+    noteHashes: Fr[];
+    nullifiers: Fr[];
+  }>;
+  txeGetPrivateEvents(selector: EventSelector, contractAddress: AztecAddress, scope: AztecAddress): Promise<Fr[][]>;
+  txePrivateCallNewFlow(
+    from: AztecAddress,
+    targetContractAddress: AztecAddress,
+    functionSelector: FunctionSelector,
+    args: Fr[],
+    argsHash: Fr,
+    isStaticCall: boolean,
+  ): Promise<Fr[]>;
+  txeSimulateUtilityFunction(
+    targetContractAddress: AztecAddress,
+    functionSelector: FunctionSelector,
+    args: Fr[],
+  ): Promise<Fr[]>;
+  txePublicCallNewFlow(
+    from: AztecAddress,
+    targetContractAddress: AztecAddress,
+    calldata: Fr[],
+    isStaticCall: boolean,
+  ): Promise<Fr[]>;
+}

package/src/oracle/txe_oracle_public_context.ts

@@ -0,0 +1,158 @@
+import { BlockNumber } from '@aztec/foundation/branded-types';
+import { Fr } from '@aztec/foundation/curves/bn254';
+import { type Logger, createLogger } from '@aztec/foundation/log';
+import { PublicDataWrite } from '@aztec/stdlib/avm';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+import type { L2Block } from '@aztec/stdlib/block';
+import { computePublicDataTreeLeafSlot, siloNoteHash, siloNullifier } from '@aztec/stdlib/hash';
+import {
+  MerkleTreeId,
+  type MerkleTreeWriteOperations,
+  PublicDataTreeLeaf,
+  PublicDataTreeLeafPreimage,
+} from '@aztec/stdlib/trees';
+import { GlobalVariables, TxEffect, TxHash } from '@aztec/stdlib/tx';
+
+import { insertTxEffectIntoWorldTrees, makeTXEBlock } from '../utils/block_creation.js';
+import type { IAvmExecutionOracle } from './interfaces.js';
+
+export class TXEOraclePublicContext implements IAvmExecutionOracle {
+  isAvm = true as const;
+
+  private logger: Logger;
+  private transientUniqueNoteHashes: Fr[] = [];
+  private transientSiloedNullifiers: Fr[] = [];
+  private publicDataWrites: PublicDataWrite[] = [];
+
+  constructor(
+    private contractAddress: AztecAddress,
+    private forkedWorldTrees: MerkleTreeWriteOperations,
+    private txRequestHash: Fr,
+    private globalVariables: GlobalVariables,
+  ) {
+    this.logger = createLogger('txe:public_context');
+
+    this.logger.debug('Entering Public Context', {
+      contractAddress,
+      blockNumber: globalVariables.blockNumber,
+      timestamp: globalVariables.timestamp,
+    });
+  }
+
+  avmOpcodeAddress(): Promise<AztecAddress> {
+    return Promise.resolve(this.contractAddress);
+  }
+
+  avmOpcodeSender(): Promise<AztecAddress> {
+    return Promise.resolve(AztecAddress.ZERO); // todo: change?
+  }
+
+  avmOpcodeBlockNumber(): Promise<BlockNumber> {
+    return Promise.resolve(this.globalVariables.blockNumber);
+  }
+
+  avmOpcodeTimestamp(): Promise<bigint> {
+    return Promise.resolve(this.globalVariables.timestamp);
+  }
+
+  avmOpcodeIsStaticCall(): Promise<boolean> {
+    return Promise.resolve(false);
+  }
+
+  avmOpcodeChainId(): Promise<Fr> {
+    return Promise.resolve(this.globalVariables.chainId);
+  }
+
+  avmOpcodeVersion(): Promise<Fr> {
+    return Promise.resolve(this.globalVariables.version);
+  }
+
+  async avmOpcodeEmitNullifier(nullifier: Fr) {
+    const siloedNullifier = await siloNullifier(this.contractAddress, nullifier);
+    this.transientSiloedNullifiers.push(siloedNullifier);
+  }
+
+  async avmOpcodeEmitNoteHash(noteHash: Fr) {
+    const siloedNoteHash = await siloNoteHash(this.contractAddress, noteHash);
+    // TODO: make the note hash unique - they are only siloed right now
+    this.transientUniqueNoteHashes.push(siloedNoteHash);
+  }
+
+  async avmOpcodeNullifierExists(siloedNullifier: Fr): Promise<boolean> {
+    const treeIndex = (
+      await this.forkedWorldTrees.findLeafIndices(MerkleTreeId.NULLIFIER_TREE, [siloedNullifier.toBuffer()])
+    )[0];
+    const transientIndex = this.transientSiloedNullifiers.find(n => n.equals(siloedNullifier));
+
+    return treeIndex !== undefined || transientIndex !== undefined;
+  }
+
+  async avmOpcodeStorageWrite(slot: Fr, value: Fr) {
+    this.logger.debug('AVM storage write', { slot, value });
+
+    const dataWrite = new PublicDataWrite(await computePublicDataTreeLeafSlot(this.contractAddress, slot), value);
+
+    this.publicDataWrites.push(dataWrite);
+
+    await this.forkedWorldTrees.sequentialInsert(MerkleTreeId.PUBLIC_DATA_TREE, [
+      new PublicDataTreeLeaf(dataWrite.leafSlot, dataWrite.value).toBuffer(),
+    ]);
+  }
+
+  async avmOpcodeStorageRead(slot: Fr, contractAddress: AztecAddress): Promise<Fr> {
+    const leafSlot = await computePublicDataTreeLeafSlot(contractAddress, slot);
+
+    const lowLeafResult = await this.forkedWorldTrees.getPreviousValueIndex(
+      MerkleTreeId.PUBLIC_DATA_TREE,
+      leafSlot.toBigInt(),
+    );
+
+    const value =
+      !lowLeafResult || !lowLeafResult.alreadyPresent
+        ? Fr.ZERO
+        : (
+            (await this.forkedWorldTrees.getLeafPreimage(
+              MerkleTreeId.PUBLIC_DATA_TREE,
+              lowLeafResult.index,
+            )) as PublicDataTreeLeafPreimage
+          ).leaf.value;
+
+    this.logger.debug('AVM storage read', { slot, contractAddress, value });
+
+    return value;
+  }
+
+  async close(): Promise<L2Block> {
+    this.logger.debug('Exiting Public Context, building block with collected side effects', {
+      blockNumber: this.globalVariables.blockNumber,
+    });
+
+    const txEffect = this.makeTxEffect();
+    await insertTxEffectIntoWorldTrees(txEffect, this.forkedWorldTrees);
+
+    const block = await makeTXEBlock(this.forkedWorldTrees, this.globalVariables, [txEffect]);
+
+    await this.forkedWorldTrees.close();
+
+    this.logger.debug('Exited PublicContext with built block', {
+      blockNumber: block.number,
+      txEffects: block.body.txEffects,
+    });
+
+    return block;
+  }
+
+  private makeTxEffect(): TxEffect {
+    const txEffect = TxEffect.empty();
+
+    txEffect.noteHashes = this.transientUniqueNoteHashes;
+    txEffect.nullifiers = [this.txRequestHash, ...this.transientSiloedNullifiers];
+
+    txEffect.publicDataWrites = this.publicDataWrites;
+    // TODO: support public logs
+
+    txEffect.txHash = new TxHash(new Fr(this.globalVariables.blockNumber));
+
+    return txEffect;
+  }
+}