|
1 | 1 | import { FunctionSelector } from '@aztec/stdlib/abi'; |
2 | | -import { AztecAddress } from '@aztec/stdlib/aztec-address'; |
| 2 | +import type { AztecAddress } from '@aztec/stdlib/aztec-address'; |
3 | 3 | import type { ContractOverrides } from '@aztec/stdlib/tx'; |
4 | 4 |
|
5 | 5 | import type { ContractStore } from '../storage/contract_store/contract_store.js'; |
6 | 6 |
|
7 | 7 | /* |
8 | | - * Proxy generator for a ContractStore that allows overriding contract instances and artifacts, so |
9 | | - * the contract function simulator can execute different bytecode on certain addresses. An example use case |
10 | | - * would be overriding your own account contract so that valid signatures don't have to be provided while simulating. |
| 8 | + * Proxy generator for a ContractStore that allows overriding contract instances at given addresses, |
| 9 | + * so the contract function simulator can execute different bytecode on certain addresses. An example |
| 10 | + * use case would be overriding your own account contract so that valid signatures don't have to be |
| 11 | + * provided while simulating. |
| 12 | + * |
| 13 | + * Function artifact lookups for an overridden address are routed via the override-instance's |
| 14 | + * `currentContractClassId` rather than the underlying ContractStore's address→class mapping — |
| 15 | + * `ContractStore.getFunctionArtifact` calls `this.getContractInstance` internally, which the proxy |
| 16 | + * cannot intercept (the binding sees the raw target, not the proxy). The target class must be |
| 17 | + * registered ahead of time via `pxe.registerContractClass(...)`. |
11 | 18 | */ |
12 | 19 | export class ProxiedContractStoreFactory { |
13 | 20 | static create(contractStore: ContractStore, overrides?: ContractOverrides) { |
14 | 21 | if (!overrides) { |
15 | 22 | return contractStore; |
16 | 23 | } |
17 | 24 |
|
| 25 | + const findFunctionInArtifact = async ( |
| 26 | + artifact: { name: string; functions: { name: string; parameters: any[] }[] }, |
| 27 | + selector: FunctionSelector, |
| 28 | + contractAddress: AztecAddress, |
| 29 | + ) => { |
| 30 | + for (const fn of artifact.functions) { |
| 31 | + const fnSelector = await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters); |
| 32 | + if (fnSelector.equals(selector)) { |
| 33 | + return { ...fn, contractName: artifact.name } as any; |
| 34 | + } |
| 35 | + } |
| 36 | + throw new Error( |
| 37 | + `Function with selector ${selector} not found in stub artifact for overridden contract at ${contractAddress}.`, |
| 38 | + ); |
| 39 | + }; |
| 40 | + |
18 | 41 | return new Proxy(contractStore, { |
19 | 42 | get(target, prop: keyof ContractStore) { |
20 | | - switch (prop) { |
21 | | - case 'getContractInstance': { |
22 | | - return async (address: AztecAddress) => { |
23 | | - if (overrides[address.toString()]) { |
24 | | - const { instance } = overrides[address.toString()]!; |
25 | | - instance.address = address; |
26 | | - const realInstance = await target.getContractInstance(address); |
27 | | - if (!realInstance) { |
28 | | - throw new Error(`Contract instance not found for address: ${address}`); |
29 | | - } |
30 | | - instance.currentContractClassId = realInstance.currentContractClassId; |
31 | | - instance.originalContractClassId = realInstance.originalContractClassId; |
32 | | - instance.initializationHash = realInstance.initializationHash; |
33 | | - return instance; |
34 | | - } else { |
35 | | - return target.getContractInstance(address); |
36 | | - } |
37 | | - }; |
38 | | - } |
39 | | - case 'getFunctionArtifact': { |
40 | | - return async (contractAddress: AztecAddress, selector: FunctionSelector) => { |
41 | | - if (overrides[contractAddress.toString()]) { |
42 | | - const { artifact } = overrides[contractAddress.toString()]!; |
43 | | - const functions = artifact.functions; |
44 | | - for (let i = 0; i < functions.length; i++) { |
45 | | - const fn = functions[i]; |
46 | | - const fnSelector = await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters); |
47 | | - if (fnSelector.equals(selector)) { |
48 | | - return fn; |
49 | | - } |
50 | | - } |
51 | | - throw new Error( |
52 | | - `Function with selector ${selector} not found in stub artifact for overridden contract at ${contractAddress}. The stub does not implement this function.`, |
53 | | - ); |
54 | | - } else { |
55 | | - return target.getFunctionArtifact(contractAddress, selector); |
56 | | - } |
57 | | - }; |
58 | | - } |
59 | | - case 'getFunctionArtifactWithDebugMetadata': { |
60 | | - return async (contractAddress: AztecAddress, selector: FunctionSelector) => { |
61 | | - if (overrides[contractAddress.toString()]) { |
62 | | - const { artifact } = overrides[contractAddress.toString()]!; |
63 | | - const functions = artifact.functions; |
64 | | - for (let i = 0; i < functions.length; i++) { |
65 | | - const fn = functions[i]; |
66 | | - const fnSelector = await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters); |
67 | | - if (fnSelector.equals(selector)) { |
68 | | - return fn; |
69 | | - } |
70 | | - } |
71 | | - throw new Error( |
72 | | - `Function with selector ${selector} not found in stub artifact for overridden contract at ${contractAddress}. The stub does not implement this function.`, |
73 | | - ); |
74 | | - } else { |
75 | | - return target.getFunctionArtifactWithDebugMetadata(contractAddress, selector); |
76 | | - } |
77 | | - }; |
78 | | - } |
79 | | - default: { |
80 | | - const value = Reflect.get(target, prop); |
81 | | - if (typeof value === 'function') { |
82 | | - return value.bind(target); |
| 43 | + if (prop === 'getContractInstance') { |
| 44 | + return (address: AztecAddress) => { |
| 45 | + const override = overrides[address.toString()]; |
| 46 | + return override ? Promise.resolve(override.instance) : target.getContractInstance(address); |
| 47 | + }; |
| 48 | + } |
| 49 | + if (prop === 'getFunctionArtifact' || prop === 'getFunctionArtifactWithDebugMetadata') { |
| 50 | + return async (contractAddress: AztecAddress, selector: FunctionSelector) => { |
| 51 | + const override = overrides[contractAddress.toString()]; |
| 52 | + if (!override) { |
| 53 | + return (target[prop] as Function).call(target, contractAddress, selector); |
| 54 | + } |
| 55 | + const artifact = await target.getContractArtifact(override.instance.currentContractClassId); |
| 56 | + if (!artifact) { |
| 57 | + throw new Error( |
| 58 | + `No artifact registered for override class ${override.instance.currentContractClassId} ` + |
| 59 | + `at ${contractAddress}. Register it via pxe.registerContractClass(...) before simulating.`, |
| 60 | + ); |
83 | 61 | } |
84 | | - return value; |
85 | | - } |
| 62 | + return findFunctionInArtifact(artifact, selector, contractAddress); |
| 63 | + }; |
86 | 64 | } |
| 65 | + const value = Reflect.get(target, prop); |
| 66 | + return typeof value === 'function' ? value.bind(target) : value; |
87 | 67 | }, |
88 | 68 | }) satisfies ContractStore; |
89 | 69 | } |
|
0 commit comments