refactor(txe): normalize deploy and addAccount to oracle registry (#23536)

Author: nchamo

yarn-project/txe/src/index.ts

@@ -1,24 +1,12 @@
-import { type NoirCompiledContract, loadContractArtifact } from '@aztec/aztec.js/abi';
-import { AztecAddress } from '@aztec/aztec.js/addresses';
-import {
-  type ContractInstanceWithAddress,
-  getContractInstanceFromInstantiationParams,
-} from '@aztec/aztec.js/contracts';
 import { Fr } from '@aztec/aztec.js/fields';
-import { PublicKeys, deriveKeys } from '@aztec/aztec.js/keys';
 import type { Logger } from '@aztec/foundation/log';
 import { cloneEphemeralStoreFrom } from '@aztec/kv-store/lmdb-v2';
 import type { ProtocolContractName } from '@aztec/protocol-contracts';
 import { ContractStore } from '@aztec/pxe/client/lazy';
-import { computeArtifactHash } from '@aztec/stdlib/contract';
-import type { ContractArtifactWithHash } 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, parse } from 'path';
+import { join } from 'path';
 import { z } from 'zod';
 
 // Side-effect import: registers the msgpackr Fr extension for the bundled `Fr` class. Must
@@ -28,64 +16,17 @@ import { type TXEOracleFunctionName, TXESession } from './txe_session.js';
 import {
   type ForeignCallArgs,
   ForeignCallArgsSchema,
-  type ForeignCallArray,
   type ForeignCallResult,
   ForeignCallResultSchema,
-  type ForeignCallSingle,
-  addressFromSingle,
-  fromArray,
-  fromSingle,
-  toSingle,
 } from './utils/encoding.js';
+import { TXEArtifactResolver } from './utils/txe_artifact_resolver.js';
 
 // Protocol contracts TXE registers in its contract store. Only AuthRegistry is needed for the
 // current test suites; add a contract here if a lookup against a `0x000…00X` address fails.
 export const TXE_REQUIRED_PROTOCOL_CONTRACTS: ProtocolContractName[] = [];
 
 const sessions = new Map<number, TXESession>();
 
-/**
- * Cache + in-flight map pair. Lookup hits the cache, then awaits an in-flight `compute()` if one
- * exists, otherwise starts one and stores it. Guarantees `compute()` runs at most once per `key`
- * across concurrent callers, which matters because `computeArtifactHash` is expensive.
- */
-class AsyncCache<K, V> {
-  private readonly cache = new Map<K, V>();
-  private readonly inFlight = new Map<K, Promise<V>>();
-
-  getOrCompute(key: K, compute: () => Promise<V>): Promise<V> {
-    const cached = this.cache.get(key);
-    if (cached !== undefined) {
-      return Promise.resolve(cached);
-    }
-    let pending = this.inFlight.get(key);
-    if (!pending) {
-      pending = (async () => {
-        try {
-          const value = await compute();
-          this.cache.set(key, value);
-          return value;
-        } finally {
-          this.inFlight.delete(key);
-        }
-      })();
-      this.inFlight.set(key, pending);
-    }
-    return pending;
-  }
-}
-
-// Full deploys (artifact + computed instance), keyed by the full deploy context (contract +
-// constructor args + publicKeys + salt + deployer). Hits on repeated identical deploys.
-const TXEDeploymentsCache = new AsyncCache<
-  string,
-  { artifact: ContractArtifactWithHash; instance: ContractInstanceWithAddress }
->();
-
-// Loaded + hashed contract artifact, keyed by compiled-bytecode hash. Hits across deploys of the
-// same contract when constructor args / salt / deployer differ.
-const TXEArtifactsCache = new AsyncCache<string, ContractArtifactWithHash>();
-
 export type TXEForeignCallInput = {
   session_id: number;
   function: TXEOracleFunctionName;
@@ -119,15 +60,16 @@ export interface TXEDispatcherOptions {
    */
   contractStoreSourceDir: string;
   /**
-   * Class id (hex) of the SchnorrAccount artifact pre-registered in the shared LMDB. When set,
-   * `#processAddAccountInputs` looks the artifact up from the cloned store instead of
-   * recomputing it via `getSchnorrAccountContractArtifact()` + `computeArtifactHash()`.
+   * Class id (hex) of the SchnorrAccount artifact pre-registered in the shared LMDB. The
+   * {@link TXEArtifactResolver} looks the artifact up from the cloned store via this class id
+   * instead of recomputing it via `getSchnorrAccountContractArtifact()` + `computeArtifactHash()`.
    */
   schnorrClassId: string;
 }
 
 export class TXEDispatcher {
   private contractStore!: ContractStore;
+  private artifactResolver!: TXEArtifactResolver;
   private readonly contractStoreSourceDir: string;
   private readonly schnorrClassId: Fr;
 
@@ -156,146 +98,25 @@ export class TXEDispatcher {
       2,
     );
     this.contractStore = new ContractStore(kvStore);
+    this.artifactResolver = new TXEArtifactResolver(this.contractStore, this.schnorrClassId);
     this.logger.debug('Cloned shared protocol-contracts store', { totalMs: Date.now() - t0 });
   }
 
-  private fastHashFile(path: string): Promise<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() as string);
-      });
-
-      fd.pipe(hash);
-    });
-  }
-
-  async #processDeployInputs({ inputs, root_path: rootPath, package_name: packageName }: TXEForeignCallInput) {
-    const [contractPath, initializer] = inputs.slice(0, 2).map(input =>
-      fromArray(input as ForeignCallArray)
-        .map(char => String.fromCharCode(char.toNumber()))
-        .join(''),
-    );
-
-    const decodedArgs = fromArray(inputs[3] as ForeignCallArray);
-    const secret = fromSingle(inputs[4] as ForeignCallSingle);
-    const salt = fromSingle(inputs[5] as ForeignCallSingle);
-    const deployer = addressFromSingle(inputs[6] as ForeignCallSingle);
-    const publicKeys = secret.equals(Fr.ZERO) ? PublicKeys.default() : (await deriveKeys(secret)).publicKeys;
-    const publicKeysHash = await publicKeys.hash();
-
-    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}-${salt}-${deployer}-${fileHash}`;
-
-    const { artifact, instance } = await TXEDeploymentsCache.getOrCompute(cacheKey, async () => {
-      this.logger.debug(`Loading compiled artifact ${artifactPath}`);
-      // Inner cache: artifact load + hash depends only on the compiled bytecode (`fileHash`), so
-      // subsequent deploys of the same contract — regardless of constructor args / deployer /
-      // salt — reuse the same `ContractArtifactWithHash`.
-      const computedArtifact = await TXEArtifactsCache.getOrCompute(fileHash, async () => {
-        const artifactJSON = JSON.parse(await readFile(artifactPath, 'utf-8')) as NoirCompiledContract;
-        const artifactWithoutHash = loadContractArtifact(artifactJSON);
-        return { ...artifactWithoutHash, artifactHash: await computeArtifactHash(artifactWithoutHash) };
-      });
-      this.logger.debug(
-        `Deploy ${computedArtifact.name} with initializer ${initializer}(${decodedArgs}) and public keys hash ${publicKeysHash.toString()}`,
-      );
-      const computedInstance = await getContractInstanceFromInstantiationParams(computedArtifact, {
-        constructorArgs: decodedArgs,
-        skipArgsDecoding: true,
-        salt,
-        publicKeys,
-        constructorArtifact: initializer ? initializer : undefined,
-        deployer,
-      });
-      return { artifact: computedArtifact, instance: computedInstance };
-    });
-
-    inputs.splice(0, 1, artifact, instance, toSingle(secret));
-  }
-
-  async #processAddAccountInputs({ inputs }: TXEForeignCallInput) {
-    const secret = fromSingle(inputs[0] as ForeignCallSingle);
-
-    const cacheKey = `SchnorrAccountContract-${secret}`;
-
-    const { artifact, instance } = await TXEDeploymentsCache.getOrCompute(cacheKey, async () => {
-      const [artifactFromStore, classWithPreimage] = await Promise.all([
-        this.contractStore.getContractArtifact(this.schnorrClassId),
-        this.contractStore.getContractClassWithPreimage(this.schnorrClassId),
-      ]);
-      if (!artifactFromStore || !classWithPreimage) {
-        throw new Error(
-          `SchnorrAccount not found in shared contract store at class id ${this.schnorrClassId.toString()}`,
-        );
-      }
-      const computedArtifact = { ...artifactFromStore, artifactHash: classWithPreimage.artifactHash };
-      const keys = await deriveKeys(secret);
-      const args = [keys.publicKeys.ivpkM.x, keys.publicKeys.ivpkM.y];
-      const computedInstance = await getContractInstanceFromInstantiationParams(computedArtifact, {
-        constructorArgs: args,
-        skipArgsDecoding: true,
-        salt: Fr.ONE,
-        publicKeys: keys.publicKeys,
-        constructorArtifact: 'constructor',
-        deployer: AztecAddress.ZERO,
-      });
-      return { artifact: computedArtifact, instance: computedInstance };
-    });
-
-    inputs.splice(0, 0, artifact, instance);
-  }
-
   // eslint-disable-next-line camelcase
   async resolve_foreign_call(callData: TXEForeignCallInput): Promise<ForeignCallResult> {
-    const { session_id: sessionId, function: functionName, inputs } = callData;
+    const {
+      session_id: sessionId,
+      function: functionName,
+      inputs,
+      root_path: rootPath,
+      package_name: packageName,
+    } = callData;
     this.logger.debug(`Calling ${functionName} on session ${sessionId}`);
 
     if (!sessions.has(sessionId)) {
       this.logger.debug(`Creating new session ${sessionId}`);
       await this.warmUp();
-      sessions.set(sessionId, await TXESession.init(this.contractStore));
-    }
-
-    switch (functionName) {
-      case 'aztec_txe_deploy': {
-        await this.#processDeployInputs(callData);
-        break;
-      }
-      case 'aztec_txe_addAccount': {
-        await this.#processAddAccountInputs(callData);
-        break;
-      }
+      sessions.set(sessionId, await TXESession.init(this.contractStore, this.artifactResolver, rootPath, packageName));
     }
 
     return await sessions.get(sessionId)!.processFunction(functionName, inputs);

yarn-project/txe/src/oracle/interfaces.ts

@@ -1,6 +1,4 @@
-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';
@@ -60,9 +58,16 @@ export interface ITxeExecutionOracle {
   getNextBlockTimestamp(): Promise<UInt64>;
   advanceBlocksBy(blocks: number): Promise<void>;
   advanceTimestampBy(duration: UInt64): void;
-  deploy(artifact: ContractArtifact, instance: ContractInstanceWithAddress, foreignSecret: Fr): Promise<void>;
+  deploy(
+    contractPath: string,
+    initializer: string,
+    args: Fr[],
+    secret: Fr,
+    salt: Fr,
+    deployer: AztecAddress,
+  ): Promise<Fr[]>;
   createAccount(secret: Fr): Promise<CompleteAddress>;
-  addAccount(artifact: ContractArtifact, instance: ContractInstanceWithAddress, secret: Fr): Promise<CompleteAddress>;
+  addAccount(secret: Fr): Promise<CompleteAddress>;
   addAuthWitness(address: AztecAddress, messageHash: Fr): Promise<void>;
   getLastBlockTimestamp(): Promise<bigint>;
   getLastTxEffects(): Promise<{

yarn-project/txe/src/oracle/txe_oracle_registry.ts

@@ -22,6 +22,7 @@ import {
   ORACLE_REGISTRY,
   type OracleRegistryEntry,
   type ParamTypes,
+  STR,
   type TypeMapping,
   U32,
   makeEntry,
@@ -188,11 +189,29 @@ export const TXE_ORACLE_REGISTRY = {
     params: [{ name: 'duration', type: BIGINT }],
   }),
 
+  aztec_txe_deploy: makeEntry({
+    params: [
+      { name: 'contractPath', type: STR },
+      { name: 'initializer', type: STR },
+      { name: 'argsLength', type: U32 },
+      { name: 'args', type: ARRAY(FIELD) },
+      { name: 'secret', type: FIELD },
+      { name: 'salt', type: FIELD },
+      { name: 'deployer', type: AZTEC_ADDRESS },
+    ],
+    returnType: ARRAY(FIELD),
+  }),
+
   aztec_txe_createAccount: makeEntry({
     params: [{ name: 'secret', type: FIELD }],
     returnType: COMPLETE_ADDRESS,
   }),
 
+  aztec_txe_addAccount: makeEntry({
+    params: [{ name: 'secret', type: FIELD }],
+    returnType: COMPLETE_ADDRESS,
+  }),
+
   aztec_txe_addAuthWitness: makeEntry({
     params: [
       { name: 'address', type: AZTEC_ADDRESS },