@aztec/pxe 4.0.0-rc.3 → 4.0.0-rc.5

package/dest/storage/contract_store/contract_store.js

@@ -1,8 +1,53 @@
+import { Fr } from '@aztec/foundation/curves/bn254';
 import { toArray } from '@aztec/foundation/iterable';
+import { BufferReader, numToUInt8, serializeToBuffer } from '@aztec/foundation/serialize';
 import { FunctionCall, FunctionSelector, FunctionType, contractArtifactFromBuffer, contractArtifactToBuffer, encodeArguments, getFunctionDebugMetadata } from '@aztec/stdlib/abi';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
 import { SerializableContractInstance, getContractClassFromArtifact } from '@aztec/stdlib/contract';
 import { PrivateFunctionsTree } from './private_functions_tree.js';
+const VERSION = 1;
+/**
+ * All contract class data except the large packedBytecode.
+ * The expensive data from the ContractClass is precomputed and stored in this format to avoid redundant hashing.
+ * Since we have to store the artifacts anyway, the final ContractClass is reconstructed by combining this data
+ * with the packedBytecode obtained from the former. That way we can have quick class lookups without wasted storage.
+ */ export class SerializableContractClassData {
+    version = VERSION;
+    id;
+    artifactHash;
+    privateFunctionsRoot;
+    publicBytecodeCommitment;
+    privateFunctions;
+    constructor(data){
+        this.id = data.id;
+        this.artifactHash = data.artifactHash;
+        this.privateFunctionsRoot = data.privateFunctionsRoot;
+        this.publicBytecodeCommitment = data.publicBytecodeCommitment;
+        this.privateFunctions = data.privateFunctions;
+    }
+    toBuffer() {
+        return serializeToBuffer(numToUInt8(this.version), this.id, this.artifactHash, this.privateFunctionsRoot, this.publicBytecodeCommitment, this.privateFunctions.length, ...this.privateFunctions.map((fn)=>serializeToBuffer(fn.selector, fn.vkHash)));
+    }
+    static fromBuffer(bufferOrReader) {
+        const reader = BufferReader.asReader(bufferOrReader);
+        const version = reader.readUInt8();
+        if (version !== VERSION) {
+            throw new Error(`Unexpected contract class data version ${version}`);
+        }
+        return new SerializableContractClassData({
+            id: reader.readObject(Fr),
+            artifactHash: reader.readObject(Fr),
+            privateFunctionsRoot: reader.readObject(Fr),
+            publicBytecodeCommitment: reader.readObject(Fr),
+            privateFunctions: reader.readVector({
+                fromBuffer: (r)=>({
+                        selector: r.readObject(FunctionSelector),
+                        vkHash: r.readObject(Fr)
+                    })
+            })
+        });
+    }
+}
 /**
  * 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
@@ -12,24 +57,46 @@ import { PrivateFunctionsTree } from './private_functions_tree.js';
  */ 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 = new Map();
-    /** Map from contract address to contract class id */ #contractClassIdMap = new Map();
+    /**
+   * In-memory cache of deserialized ContractArtifact objects, keyed by class id string.
+   * Avoids repeated LMDB reads + JSON.parse + Zod validation on every oracle call.
+   * Artifacts are large but immutable after registration — safe to cache for the lifetime of the store.
+   */ // TODO: Update it to be LRU cache so that it doesn't keep all the data all the time.
+    #contractArtifactCache = new Map();
+    /** Map from contract address to contract class id (avoids KV round-trip on hot path). */ #contractClassIdMap = new Map();
     #store;
     #contractArtifacts;
+    #contractClassData;
     #contractInstances;
     constructor(store){
         this.#store = store;
         this.#contractArtifacts = store.openMap('contract_artifacts');
+        this.#contractClassData = store.openMap('contract_classes');
         this.#contractInstances = store.openMap('contracts_instances');
     }
     // Setters
-    async addContractArtifact(id, contract) {
-        // Validation outside transactionAsync - these are not DB operations
+    /**
+   * Registers a new contract artifact and its corresponding class data.
+   * IMPORTANT: This method does not verify that the provided artifact matches the class data or that the class id matches the artifact.
+   * It is the caller's responsibility to ensure the consistency and correctness of the provided data.
+   * This is done to avoid redundant, expensive contract class computations
+   */ async addContractArtifact(contract, contractClassWithIdAndPreimage) {
+        const contractClass = contractClassWithIdAndPreimage ?? await getContractClassFromArtifact(contract);
+        const key = contractClass.id.toString();
+        if (this.#contractArtifactCache.has(key)) {
+            return contractClass.id;
+        }
         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()));
+        const privateSelectors = await Promise.all(privateFunctions.map(async (fn)=>(await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters)).toString()));
         if (privateSelectors.length !== new Set(privateSelectors).size) {
             throw new Error('Repeated function selectors of private functions');
         }
-        await this.#store.transactionAsync(()=>this.#contractArtifacts.set(id.toString(), contractArtifactToBuffer(contract)));
+        this.#contractArtifactCache.set(key, contract);
+        await this.#store.transactionAsync(async ()=>{
+            await this.#contractArtifacts.set(key, contractArtifactToBuffer(contract));
+            await this.#contractClassData.set(key, new SerializableContractClassData(contractClass).toBuffer());
+        });
+        return contractClass.id;
     }
     async addContractInstance(contract) {
         this.#contractClassIdMap.set(contract.address.toString(), contract.currentContractClassId);
@@ -37,25 +104,17 @@ import { PrivateFunctionsTree } from './private_functions_tree.js';
     }
     // Private getters
     async #getContractClassId(contractAddress) {
-        if (!this.#contractClassIdMap.has(contractAddress.toString())) {
+        const key = contractAddress.toString();
+        if (!this.#contractClassIdMap.has(key)) {
             const instance = await this.getContractInstance(contractAddress);
             if (!instance) {
                 return;
             }
-            this.#contractClassIdMap.set(contractAddress.toString(), instance.currentContractClassId);
+            this.#contractClassIdMap.set(key, instance.currentContractClassId);
         }
-        return this.#contractClassIdMap.get(contractAddress.toString());
+        return this.#contractClassIdMap.get(key);
     }
-    /**
-   * 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) {
+    async #getPrivateFunctionTreeForClassId(classId) {
         if (!this.#privateFunctionTrees.has(classId.toString())) {
             const artifact = await this.getContractArtifact(classId);
             if (!artifact) {
@@ -66,9 +125,9 @@ import { PrivateFunctionsTree } from './private_functions_tree.js';
         }
         return this.#privateFunctionTrees.get(classId.toString());
     }
-    async #getContractArtifactByAddress(contractAddress) {
-        const contractClassId = await this.#getContractClassId(contractAddress);
-        return contractClassId && this.getContractArtifact(contractClassId);
+    async #getArtifactByAddress(contractAddress) {
+        const classId = await this.#getContractClassId(contractAddress);
+        return classId && this.getContractArtifact(classId);
     }
     // Public getters
     getContractsAddresses() {
@@ -77,22 +136,43 @@ import { PrivateFunctionsTree } from './private_functions_tree.js';
             return keys.map(AztecAddress.fromString);
         });
     }
-    /** Returns a contract instance for a given address. Throws if not found. */ getContractInstance(contractAddress) {
+    /** Returns a contract instance for a given address. */ getContractInstance(contractAddress) {
         return this.#store.transactionAsync(async ()=>{
             const contract = await this.#contractInstances.getAsync(contractAddress.toString());
             return contract && SerializableContractInstance.fromBuffer(contract).withAddress(contractAddress);
         });
     }
-    getContractArtifact(contractClassId) {
-        return this.#store.transactionAsync(async ()=>{
-            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 the raw contract artifact for a given class id. */ async getContractArtifact(contractClassId) {
+        const key = contractClassId.toString();
+        const cached = this.#contractArtifactCache.get(key);
+        if (cached) {
+            return cached;
+        }
+        const artifact = await this.#store.transactionAsync(async ()=>{
+            const buf = await this.#contractArtifacts.getAsync(key);
+            return buf && contractArtifactFromBuffer(buf);
         });
+        if (artifact) {
+            this.#contractArtifactCache.set(key, artifact);
+        }
+        return artifact;
     }
-    /** Returns a contract class for a given class id. Throws if not found. */ async getContractClass(contractClassId) {
+    /** Returns a contract class for a given class id. */ async getContractClassWithPreimage(contractClassId) {
+        const key = contractClassId.toString();
+        const buf = await this.#contractClassData.getAsync(key);
+        if (!buf) {
+            return undefined;
+        }
+        const classData = SerializableContractClassData.fromBuffer(buf);
         const artifact = await this.getContractArtifact(contractClassId);
-        return artifact && getContractClassFromArtifact(artifact);
+        if (!artifact) {
+            return undefined;
+        }
+        const packedBytecode = artifact.functions.find((f)=>f.name === 'public_dispatch')?.bytecode ?? Buffer.alloc(0);
+        return {
+            ...classData,
+            packedBytecode
+        };
     }
     async getContract(address) {
         const instance = await this.getContractInstance(address);
@@ -110,17 +190,18 @@ import { PrivateFunctionsTree } from './private_functions_tree.js';
     }
     /**
    * 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.
    */ async getFunctionArtifact(contractAddress, selector) {
-        const artifact = await this.#getContractArtifactByAddress(contractAddress);
-        const fnArtifact = artifact && await this.#findFunctionArtifactBySelector(artifact, selector);
-        return fnArtifact && {
-            ...fnArtifact,
+        const artifact = await this.#getArtifactByAddress(contractAddress);
+        if (!artifact) {
+            return undefined;
+        }
+        const fn = await this.#findFunctionArtifactBySelector(artifact, selector);
+        return fn && {
+            ...fn,
             contractName: artifact.name
         };
     }
@@ -136,40 +217,38 @@ import { PrivateFunctionsTree } from './private_functions_tree.js';
         };
     }
     async getPublicFunctionArtifact(contractAddress) {
-        const artifact = await this.#getContractArtifactByAddress(contractAddress);
-        const fnArtifact = artifact && artifact.functions.find((fn)=>fn.functionType === FunctionType.PUBLIC);
-        return fnArtifact && {
-            ...fnArtifact,
+        const artifact = await this.#getArtifactByAddress(contractAddress);
+        const fn = artifact && artifact.functions.find((f)=>f.functionType === FunctionType.PUBLIC);
+        return fn && {
+            ...fn,
             contractName: artifact.name
         };
     }
     async getFunctionAbi(contractAddress, selector) {
-        const artifact = await this.#getContractArtifactByAddress(contractAddress);
+        const artifact = await this.#getArtifactByAddress(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.
+   * @returns The corresponding function's debug metadata, or undefined.
    */ async getFunctionDebugMetadata(contractAddress, selector) {
-        const artifact = await this.#getContractArtifactByAddress(contractAddress);
-        const fnArtifact = artifact && await this.#findFunctionArtifactBySelector(artifact, selector);
-        return fnArtifact && getFunctionDebugMetadata(artifact, fnArtifact);
+        const artifact = await this.#getArtifactByAddress(contractAddress);
+        if (!artifact) {
+            return undefined;
+        }
+        const fn = await this.#findFunctionArtifactBySelector(artifact, selector);
+        return fn && getFunctionDebugMetadata(artifact, fn);
     }
     async getPublicFunctionDebugMetadata(contractAddress) {
-        const artifact = await this.#getContractArtifactByAddress(contractAddress);
-        const fnArtifact = artifact && artifact.functions.find((fn)=>fn.functionType === FunctionType.PUBLIC);
-        return fnArtifact && getFunctionDebugMetadata(artifact, fnArtifact);
+        const artifact = await this.#getArtifactByAddress(contractAddress);
+        const fn = artifact && artifact.functions.find((f)=>f.functionType === FunctionType.PUBLIC);
+        return fn && getFunctionDebugMetadata(artifact, fn);
     }
     /**
    * 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.
@@ -179,18 +258,16 @@ import { PrivateFunctionsTree } from './private_functions_tree.js';
         return tree?.getFunctionMembershipWitness(selector);
     }
     async getDebugContractName(contractAddress) {
-        const artifact = await this.#getContractArtifactByAddress(contractAddress);
+        const artifact = await this.#getArtifactByAddress(contractAddress);
         return artifact?.name;
     }
     async getDebugFunctionName(contractAddress, selector) {
-        const artifact = await this.#getContractArtifactByAddress(contractAddress);
-        const fnArtifact = artifact && await this.#findFunctionAbiBySelector(artifact, selector);
-        return `${artifact?.name ?? contractAddress}:${fnArtifact?.name ?? selector}`;
+        const artifact = await this.#getArtifactByAddress(contractAddress);
+        const fn = artifact && await this.#findFunctionAbiBySelector(artifact, selector);
+        return `${artifact?.name ?? contractAddress}:${fn?.name ?? selector}`;
     }
     async #findFunctionArtifactBySelector(artifact, selector) {
-        const functions = artifact.functions;
-        for(let i = 0; i < functions.length; i++){
-            const fn = functions[i];
+        for (const fn of artifact.functions){
             const fnSelector = await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters);
             if (fnSelector.equals(selector)) {
                 return fn;
@@ -198,12 +275,10 @@ import { PrivateFunctionsTree } from './private_functions_tree.js';
         }
     }
     async #findFunctionAbiBySelector(artifact, selector) {
-        const functions = [
+        for (const fn of [
             ...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;
@@ -219,10 +294,11 @@ import { PrivateFunctionsTree } from './private_functions_tree.js';
         if (!functionDao) {
             throw new Error(`Unknown function ${functionName} in contract ${contract.name}.`);
         }
+        const selector = await FunctionSelector.fromNameAndParameters(functionDao.name, functionDao.parameters);
         return FunctionCall.from({
             name: functionDao.name,
             to,
-            selector: await FunctionSelector.fromNameAndParameters(functionDao.name, functionDao.parameters),
+            selector,
             type: functionDao.functionType,
             hideMsgSender: false,
             isStatic: functionDao.isStatic,

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@aztec/pxe",
-  "version": "4.0.0-rc.3",
+  "version": "4.0.0-rc.5",
   "type": "module",
   "typedocOptions": {
     "entryPoints": [
@@ -70,19 +70,19 @@
     ]
   },
   "dependencies": {
-    "@aztec/bb-prover": "4.0.0-rc.3",
-    "@aztec/bb.js": "4.0.0-rc.3",
-    "@aztec/builder": "4.0.0-rc.3",
-    "@aztec/constants": "4.0.0-rc.3",
-    "@aztec/ethereum": "4.0.0-rc.3",
-    "@aztec/foundation": "4.0.0-rc.3",
-    "@aztec/key-store": "4.0.0-rc.3",
-    "@aztec/kv-store": "4.0.0-rc.3",
-    "@aztec/noir-protocol-circuits-types": "4.0.0-rc.3",
-    "@aztec/noir-types": "4.0.0-rc.3",
-    "@aztec/protocol-contracts": "4.0.0-rc.3",
-    "@aztec/simulator": "4.0.0-rc.3",
-    "@aztec/stdlib": "4.0.0-rc.3",
+    "@aztec/bb-prover": "4.0.0-rc.5",
+    "@aztec/bb.js": "4.0.0-rc.5",
+    "@aztec/builder": "4.0.0-rc.5",
+    "@aztec/constants": "4.0.0-rc.5",
+    "@aztec/ethereum": "4.0.0-rc.5",
+    "@aztec/foundation": "4.0.0-rc.5",
+    "@aztec/key-store": "4.0.0-rc.5",
+    "@aztec/kv-store": "4.0.0-rc.5",
+    "@aztec/noir-protocol-circuits-types": "4.0.0-rc.5",
+    "@aztec/noir-types": "4.0.0-rc.5",
+    "@aztec/protocol-contracts": "4.0.0-rc.5",
+    "@aztec/simulator": "4.0.0-rc.5",
+    "@aztec/stdlib": "4.0.0-rc.5",
     "koa": "^2.16.1",
     "koa-router": "^13.1.1",
     "lodash.omit": "^4.5.0",
@@ -91,8 +91,8 @@
     "viem": "npm:@aztec/viem@2.38.2"
   },
   "devDependencies": {
-    "@aztec/merkle-tree": "4.0.0-rc.3",
-    "@aztec/noir-test-contracts.js": "4.0.0-rc.3",
+    "@aztec/merkle-tree": "4.0.0-rc.5",
+    "@aztec/noir-test-contracts.js": "4.0.0-rc.5",
     "@jest/globals": "^30.0.0",
     "@types/jest": "^30.0.0",
     "@types/lodash.omit": "^4.5.7",

package/src/contract_function_simulator/contract_function_simulator.ts

@@ -4,8 +4,6 @@ import {
   AVM_SENDL2TOL1MSG_BASE_L2_GAS,
   DA_GAS_PER_FIELD,
   FIXED_AVM_STARTUP_L2_GAS,
-  FIXED_DA_GAS,
-  FIXED_L2_GAS,
   L2_GAS_PER_CONTRACT_CLASS_LOG,
   L2_GAS_PER_L2_TO_L1_MSG,
   L2_GAS_PER_NOTE_HASH,
@@ -19,6 +17,9 @@ import {
   MAX_NULLIFIERS_PER_TX,
   MAX_NULLIFIER_READ_REQUESTS_PER_TX,
   MAX_PRIVATE_LOGS_PER_TX,
+  PRIVATE_TX_L2_GAS_OVERHEAD,
+  PUBLIC_TX_L2_GAS_OVERHEAD,
+  TX_DA_GAS_OVERHEAD,
 } from '@aztec/constants';
 import { arrayNonEmptyLength, padArrayEnd } from '@aztec/foundation/collection';
 import { Fr } from '@aztec/foundation/curves/bn254';
@@ -653,7 +654,12 @@ export async function generateSimulatedProvingResult(
 
   const publicInputs = new PrivateKernelTailCircuitPublicInputs(
     constantData,
-    /*gasUsed=*/ gasUsed.add(Gas.from({ l2Gas: FIXED_L2_GAS, daGas: FIXED_DA_GAS })),
+    /*gasUsed=*/ gasUsed.add(
+      Gas.from({
+        l2Gas: isPrivateOnlyTx ? PRIVATE_TX_L2_GAS_OVERHEAD : PUBLIC_TX_L2_GAS_OVERHEAD,
+        daGas: TX_DA_GAS_OVERHEAD,
+      }),
+    ),
     /*feePayer=*/ AztecAddress.zero(),
     /*expirationTimestamp=*/ 0n,
     hasPublicCalls ? inputsForPublic : undefined,

package/src/private_kernel/private_kernel_oracle.ts

@@ -8,11 +8,7 @@ import { ProtocolContractAddress } from '@aztec/protocol-contracts';
 import type { FunctionSelector } from '@aztec/stdlib/abi';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
 import { BlockHash } from '@aztec/stdlib/block';
-import {
-  type ContractInstanceWithAddress,
-  computeContractClassIdPreimage,
-  computeSaltedInitializationHash,
-} from '@aztec/stdlib/contract';
+import { type ContractInstanceWithAddress, computeSaltedInitializationHash } from '@aztec/stdlib/contract';
 import { DelayedPublicMutableValues, DelayedPublicMutableValuesWithHash } from '@aztec/stdlib/delayed-public-mutable';
 import { computePublicDataTreeLeafSlot } from '@aztec/stdlib/hash';
 import type { AztecNode } from '@aztec/stdlib/interfaces/client';
@@ -49,11 +45,15 @@ export class PrivateKernelOracle {
 
   /** Retrieves the preimage of a contract class id from the contract classes db. */
   public async getContractClassIdPreimage(contractClassId: Fr) {
-    const contractClass = await this.contractStore.getContractClass(contractClassId);
+    const contractClass = await this.contractStore.getContractClassWithPreimage(contractClassId);
     if (!contractClass) {
       throw new Error(`Contract class not found when getting class id preimage. Class id: ${contractClassId}.`);
     }
-    return computeContractClassIdPreimage(contractClass);
+    return {
+      artifactHash: contractClass.artifactHash,
+      privateFunctionsRoot: contractClass.privateFunctionsRoot,
+      publicBytecodeCommitment: contractClass.publicBytecodeCommitment,
+    };
   }
 
   /** Returns a membership witness with the sibling path and leaf index in our private functions tree. */

package/src/pxe.ts

@@ -344,9 +344,8 @@ export class PXE {
   async #registerProtocolContracts() {
     const registered: Record<string, string> = {};
     for (const name of protocolContractNames) {
-      const { address, contractClass, instance, artifact } =
-        await this.protocolContractsProvider.getProtocolContractArtifact(name);
-      await this.contractStore.addContractArtifact(contractClass.id, artifact);
+      const { address, instance, artifact } = await this.protocolContractsProvider.getProtocolContractArtifact(name);
+      await this.contractStore.addContractArtifact(artifact);
       await this.contractStore.addContractInstance(instance);
       registered[name] = address.toString();
     }
@@ -601,8 +600,7 @@ export class PXE {
    * @param artifact - The build artifact for the contract class.
    */
   public async registerContractClass(artifact: ContractArtifact): Promise<void> {
-    const { id: contractClassId } = await getContractClassFromArtifact(artifact);
-    await this.contractStore.addContractArtifact(contractClassId, artifact);
+    const contractClassId = await this.contractStore.addContractArtifact(artifact);
     this.log.info(`Added contract class ${artifact.name} with id ${contractClassId}`);
   }
 
@@ -621,17 +619,17 @@ export class PXE {
     if (artifact) {
       // If the user provides an artifact, validate it against the expected class id and register it
       const contractClass = await getContractClassFromArtifact(artifact);
-      const contractClassId = contractClass.id;
-      if (!contractClassId.equals(instance.currentContractClassId)) {
+      if (!contractClass.id.equals(instance.currentContractClassId)) {
         throw new Error(
-          `Artifact does not match expected class id (computed ${contractClassId} but instance refers to ${instance.currentContractClassId})`,
+          `Artifact does not match expected class id (computed ${contractClass.id} but instance refers to ${instance.currentContractClassId})`,
         );
       }
       const computedAddress = await computeContractAddressFromInstance(instance);
       if (!computedAddress.equals(instance.address)) {
         throw new Error('Added a contract in which the address does not match the contract instance.');
       }
-      await this.contractStore.addContractArtifact(contractClass.id, artifact);
+
+      await this.contractStore.addContractArtifact(artifact, contractClass);
 
       const publicFunctionSignatures = artifact.functions
         .filter(fn => fn.functionType === FunctionType.PUBLIC)
@@ -680,15 +678,16 @@ export class PXE {
         throw new Error('Could not update contract to a class different from the current one.');
       }
 
-      await this.contractStore.addContractArtifact(contractClass.id, artifact);
-
       const publicFunctionSignatures = artifact.functions
         .filter(fn => fn.functionType === FunctionType.PUBLIC)
         .map(fn => decodeFunctionSignature(fn.name, fn.parameters));
       await this.node.registerContractFunctionSignatures(publicFunctionSignatures);
 
       currentInstance.currentContractClassId = contractClass.id;
-      await this.contractStore.addContractInstance(currentInstance);
+      await Promise.all([
+        this.contractStore.addContractArtifact(artifact, contractClass),
+        this.contractStore.addContractInstance(currentInstance),
+      ]);
       this.log.info(`Updated contract ${artifact.name} at ${contractAddress.toString()} to class ${contractClass.id}`);
     });
   }