Skip to content

Commit

Permalink
feat: Allow registering contract classes in PXE
Browse files Browse the repository at this point in the history
  • Loading branch information
spalladino committed Mar 18, 2024
1 parent 82f8cf5 commit b5b3839
Show file tree
Hide file tree
Showing 21 changed files with 165 additions and 119 deletions.
10 changes: 4 additions & 6 deletions yarn-project/aztec.js/src/account_manager/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -104,12 +104,10 @@ export class AccountManager {
*/
public async register(opts: WaitOpts = DefaultWaitOpts): Promise<AccountWalletWithPrivateKey> {
await this.#register();
await this.pxe.addContracts([
{
artifact: this.accountContract.getContractArtifact(),
instance: this.getInstance(),
},
]);
await this.pxe.registerContract({
artifact: this.accountContract.getContractArtifact(),
instance: this.getInstance(),
});

await waitForAccountSynch(this.pxe, this.getCompleteAddress(), opts);
return this.getWallet();
Expand Down
3 changes: 1 addition & 2 deletions yarn-project/aztec.js/src/contract/contract_base.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
import { DeployedContract } from '@aztec/circuit-types';
import { computePartialAddress } from '@aztec/circuits.js';
import { ContractArtifact, FunctionArtifact, FunctionSelector } from '@aztec/foundation/abi';
import { ContractInstanceWithAddress } from '@aztec/types/contracts';
Expand All @@ -20,7 +19,7 @@ export type ContractMethod = ((...args: any[]) => ContractFunctionInteraction) &
/**
* Abstract implementation of a contract extended by the Contract class and generated contract types.
*/
export class ContractBase implements DeployedContract {
export class ContractBase {
/**
* An object containing contract methods mapped to their respective names.
*/
Expand Down
2 changes: 1 addition & 1 deletion yarn-project/aztec.js/src/contract/deploy_method.ts
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,7 @@ export class DeployMethod<TContract extends ContractBase = Contract> extends Bas
}
this.txRequest = await this.wallet.createTxExecutionRequest(await this.request(options));
// TODO: Should we add the contracts to the DB here, or once the tx has been sent or mined?
await this.pxe.addContracts([{ artifact: this.artifact, instance: this.instance! }]);
await this.pxe.registerContract({ artifact: this.artifact, instance: this.instance! });
}
return this.txRequest;
}
Expand Down
2 changes: 1 addition & 1 deletion yarn-project/aztec.js/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@ export {
AztecNode,
Body,
CompleteAddress,
DeployedContract,
ContractWithArtifact,
ExtendedNote,
FunctionCall,
GrumpkinPrivateKey,
Expand Down
10 changes: 7 additions & 3 deletions yarn-project/aztec.js/src/wallet/base_wallet.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import {
AuthWitness,
DeployedContract,
ContractWithArtifact,
ExtendedNote,
FunctionCall,
GetUnencryptedLogsResponse,
Expand All @@ -16,6 +16,7 @@ import {
TxReceipt,
} from '@aztec/circuit-types';
import { AztecAddress, CompleteAddress, Fr, GrumpkinPrivateKey, PartialAddress } from '@aztec/circuits.js';
import { ContractArtifact } from '@aztec/foundation/abi';
import { ContractClassWithId, ContractInstanceWithAddress } from '@aztec/types/contracts';
import { NodeInfo } from '@aztec/types/interfaces';

Expand Down Expand Up @@ -75,8 +76,11 @@ export abstract class BaseWallet implements Wallet {
getRecipient(address: AztecAddress): Promise<CompleteAddress | undefined> {
return this.pxe.getRecipient(address);
}
addContracts(contracts: DeployedContract[]): Promise<void> {
return this.pxe.addContracts(contracts);
registerContract(contract: ContractWithArtifact): Promise<void> {
return this.pxe.registerContract(contract);
}
registerContractClass(artifact: ContractArtifact): Promise<void> {
return this.pxe.registerContractClass(artifact);
}
getContracts(): Promise<AztecAddress[]> {
return this.pxe.getContracts();
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
import { ContractArtifact } from '@aztec/foundation/abi';
import { Fr } from '@aztec/foundation/fields';
import { ContractInstanceWithAddress } from '@aztec/types/contracts';

/**
* Container for contract instance and artifact or class id.
*/
export type ContractWithArtifact = (
| {
/** The artifact of the contract. */
artifact: ContractArtifact;
}
| {
/** The class id of the contract artifact. */
contractClassId: Fr;
}
) & {
/** The contract instance. */
instance: ContractInstanceWithAddress;
};
17 changes: 0 additions & 17 deletions yarn-project/circuit-types/src/interfaces/deployed-contract.ts

This file was deleted.

2 changes: 1 addition & 1 deletion yarn-project/circuit-types/src/interfaces/index.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
export * from './aztec-node.js';
export * from './pxe.js';
export * from './deployed-contract.js';
export * from './contract-with-artifact.js';
export * from './sync-status.js';
export * from './configs.js';
export * from './nullifier_tree.js';
Expand Down
14 changes: 11 additions & 3 deletions yarn-project/circuit-types/src/interfaces/pxe.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import { AztecAddress, CompleteAddress, Fr, GrumpkinPrivateKey, PartialAddress } from '@aztec/circuits.js';
import { ContractArtifact } from '@aztec/foundation/abi';
import { ContractClassWithId, ContractInstanceWithAddress } from '@aztec/types/contracts';
import { NodeInfo } from '@aztec/types/interfaces';

Expand All @@ -10,7 +11,7 @@ import { NoteFilter } from '../notes/note_filter.js';
import { Tx, TxHash, TxReceipt } from '../tx/index.js';
import { TxEffect } from '../tx_effect.js';
import { TxExecutionRequest } from '../tx_execution_request.js';
import { DeployedContract } from './deployed-contract.js';
import { ContractWithArtifact } from './contract-with-artifact.js';
import { SyncStatus } from './sync-status.js';

// docs:start:pxe-interface
Expand Down Expand Up @@ -98,15 +99,22 @@ export interface PXE {
*/
getRecipient(address: AztecAddress): Promise<CompleteAddress | undefined>;

/**
* Registers a contract class in the PXE without registering any associated contract instance with it.
*
* @param artifact - The build artifact for the contract class.
*/
registerContractClass(artifact: ContractArtifact): Promise<void>;

/**
* Adds deployed contracts to the PXE Service. Deployed contract information is used to access the
* contract code when simulating local transactions. This is automatically called by aztec.js when
* deploying a contract. Dapps that wish to interact with contracts already deployed should register
* these contracts in their users' PXE Service through this method.
*
* @param contracts - An array of DeployedContract objects containing contract ABI, address, and portal contract.
* @param contract - An object containing contract artifact and instance.
*/
addContracts(contracts: DeployedContract[]): Promise<void>;
registerContract(contract: ContractWithArtifact): Promise<void>;

/**
* Retrieves the addresses of contracts added to this PXE Service.
Expand Down
10 changes: 5 additions & 5 deletions yarn-project/circuit-types/src/mocks.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
import {
AztecAddress,
CallRequest,
Fr,
MAX_PUBLIC_CALL_STACK_LENGTH_PER_TX,
MAX_REVERTIBLE_PUBLIC_CALL_STACK_LENGTH_PER_TX,
Proof,
Expand All @@ -10,10 +9,11 @@ import { ContractArtifact } from '@aztec/foundation/abi';
import { makeTuple } from '@aztec/foundation/array';
import { times } from '@aztec/foundation/collection';
import { randomBytes } from '@aztec/foundation/crypto';
import { Fr } from '@aztec/foundation/fields';
import { Tuple } from '@aztec/foundation/serialize';
import { ContractInstanceWithAddress, SerializableContractInstance } from '@aztec/types/contracts';

import { DeployedContract } from './interfaces/index.js';
import { ContractWithArtifact } from './interfaces/index.js';
import { FunctionL2Logs, Note, TxL2Logs } from './logs/index.js';
import { makePrivateKernelTailCircuitPublicInputs, makePublicCallRequest } from './mocks_to_purge.js';
import { ExtendedNote } from './notes/index.js';
Expand Down Expand Up @@ -59,10 +59,10 @@ export const randomContractArtifact = (): ContractArtifact => ({
fileMap: {},
});

export const randomContractInstanceWithAddress = (): ContractInstanceWithAddress =>
SerializableContractInstance.random().withAddress(AztecAddress.random());
export const randomContractInstanceWithAddress = (opts: { contractClassId?: Fr } = {}): ContractInstanceWithAddress =>
SerializableContractInstance.random(opts).withAddress(AztecAddress.random());

export const randomDeployedContract = (): DeployedContract => ({
export const randomDeployedContract = (): ContractWithArtifact => ({
artifact: randomContractArtifact(),
instance: randomContractInstanceWithAddress(),
});
Expand Down
2 changes: 1 addition & 1 deletion yarn-project/cli/src/cmds/add_contract.ts
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,6 @@ export async function addContract(

const client = await createCompatibleClient(rpcUrl, debugLogger);

await client.addContracts([{ artifact, instance }]);
await client.registerContract({ artifact, instance });
log(`\nContract added to PXE at ${address.toString()} with class ${instance.contractClassId.toString()}\n`);
}
2 changes: 1 addition & 1 deletion yarn-project/end-to-end/src/benchmarks/utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ export async function waitNewPXESynced(
startingBlock: number = INITIAL_L2_BLOCK_NUM,
): Promise<PXEService> {
const pxe = await createPXEService(node, { l2BlockPollingIntervalMS: 100, l2StartingBlock: startingBlock });
await pxe.addContracts([contract]);
await pxe.registerContract(contract);
await retryUntil(() => pxe.isGlobalStateSynchronized(), 'pxe-global-sync');
return pxe;
}
Expand Down
50 changes: 20 additions & 30 deletions yarn-project/end-to-end/src/e2e_2_pxes.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -130,12 +130,10 @@ describe('e2e_2_pxes', () => {
await pxeB.registerRecipient(userA);

// Add token to PXE B (PXE A already has it because it was deployed through it)
await pxeB.addContracts([
{
artifact: TokenContract.artifact,
instance: tokenInstance,
},
]);
await pxeB.registerContract({
artifact: TokenContract.artifact,
instance: tokenInstance,
});

// Check initial balances and logs are as expected
await expectTokenBalance(walletA, tokenAddress, userA.address, initialBalance);
Expand Down Expand Up @@ -188,12 +186,10 @@ describe('e2e_2_pxes', () => {
await awaitServerSynchronized(pxeA);

// Add Child to PXE B
await pxeB.addContracts([
{
artifact: ChildContract.artifact,
instance: childCompleteAddress,
},
]);
await pxeB.registerContract({
artifact: ChildContract.artifact,
instance: childCompleteAddress,
});

const newValueToSet = new Fr(256n);

Expand Down Expand Up @@ -222,12 +218,10 @@ describe('e2e_2_pxes', () => {
await pxeB.registerRecipient(userA);

// Add token to PXE B (PXE A already has it because it was deployed through it)
await pxeB.addContracts([
{
artifact: TokenContract.artifact,
instance: tokenInstance,
},
]);
await pxeB.registerContract({
artifact: TokenContract.artifact,
instance: tokenInstance,
});

// Mint tokens to user B
await mintTokens(contractWithWalletA, userB.address, userBBalance, pxeA);
Expand Down Expand Up @@ -287,12 +281,10 @@ describe('e2e_2_pxes', () => {
await contractWithWalletA.methods.transfer(userA.address, userB.address, transferAmount1, 0).send().wait();

// now add the contract and check balances
await pxeB.addContracts([
{
artifact: TokenContract.artifact,
instance: tokenInstance,
},
]);
await pxeB.registerContract({
artifact: TokenContract.artifact,
instance: tokenInstance,
});
await expectTokenBalance(walletA, tokenAddress, userA.address, initialBalance - transferAmount1);
await expectTokenBalance(walletB, tokenAddress, userB.address, transferAmount1);
});
Expand Down Expand Up @@ -347,12 +339,10 @@ describe('e2e_2_pxes', () => {
// PXE-B had previously deferred the notes from A -> Shared, and Shared -> B
// PXE-B adds the contract
// PXE-B reprocesses the deferred notes, and sees the nullifier for A -> Shared
await pxeB.addContracts([
{
artifact: TokenContract.artifact,
instance: tokenInstance,
},
]);
await pxeB.registerContract({
artifact: TokenContract.artifact,
instance: tokenInstance,
});
await expectTokenBalance(walletB, tokenAddress, userB.address, transferAmount2);
await expect(sharedWalletOnB.isAccountStateSynchronized(sharedAccountAddress.address)).resolves.toBe(true);
await expectTokenBalance(
Expand Down
4 changes: 2 additions & 2 deletions yarn-project/end-to-end/src/e2e_deploy_contract.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -454,7 +454,7 @@ describe('e2e_deploy_contract', () => {

testDeployingAnInstance('from a contract', async instance => {
// Register the instance to be deployed in the pxe
await wallet.addContracts([{ artifact, instance }]);
await wallet.registerContract({ artifact, instance });
// Set up the contract that calls the deployer (which happens to be the TestContract) and call it
const deployer = await TestContract.deploy(wallet).send().deployed();
await deployer.methods.deploy_contract(instance.address).send().wait();
Expand Down Expand Up @@ -594,6 +594,6 @@ async function registerContract<T extends ContractBase>(
portalAddress,
deployer,
});
await wallet.addContracts([{ artifact: contractArtifact.artifact, instance }]);
await wallet.registerContract({ artifact: contractArtifact.artifact, instance });
return contractArtifact.at(instance.address, wallet);
}
40 changes: 16 additions & 24 deletions yarn-project/end-to-end/src/e2e_persistence.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -202,12 +202,10 @@ describe('Aztec persistence', () => {
});

it("pxe does not have owner's private notes", async () => {
await context.pxe.addContracts([
{
artifact: TokenContract.artifact,
instance: contractInstance,
},
]);
await context.pxe.registerContract({
artifact: TokenContract.artifact,
instance: contractInstance,
});
await context.pxe.registerRecipient(ownerAddress);

const wallet = await getUnsafeSchnorrAccount(context.pxe, Fq.random(), Fr.ZERO).waitSetup();
Expand All @@ -216,12 +214,10 @@ describe('Aztec persistence', () => {
});

it('has access to public storage', async () => {
await context.pxe.addContracts([
{
artifact: TokenContract.artifact,
instance: contractInstance,
},
]);
await context.pxe.registerContract({
artifact: TokenContract.artifact,
instance: contractInstance,
});

const wallet = await getUnsafeSchnorrAccount(context.pxe, Fq.random(), Fr.ZERO).waitSetup();
const contract = await TokenContract.at(contractAddress, wallet);
Expand All @@ -230,12 +226,10 @@ describe('Aztec persistence', () => {
});

it('pxe restores notes after registering the owner', async () => {
await context.pxe.addContracts([
{
artifact: TokenContract.artifact,
instance: contractInstance,
},
]);
await context.pxe.registerContract({
artifact: TokenContract.artifact,
instance: contractInstance,
});

const ownerAccount = getUnsafeSchnorrAccount(context.pxe, ownerPrivateKey, ownerSalt);
await ownerAccount.register();
Expand Down Expand Up @@ -263,12 +257,10 @@ describe('Aztec persistence', () => {
beforeAll(async () => {
const temporaryContext = await setup(0, { deployL1ContractsValues }, {});

await temporaryContext.pxe.addContracts([
{
artifact: TokenContract.artifact,
instance: contractInstance,
},
]);
await temporaryContext.pxe.registerContract({
artifact: TokenContract.artifact,
instance: contractInstance,
});

const ownerAccount = getUnsafeSchnorrAccount(temporaryContext.pxe, ownerPrivateKey, ownerSalt);
await ownerAccount.register();
Expand Down
3 changes: 2 additions & 1 deletion yarn-project/foundation/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,8 @@
"./committable": "./dest/committable/index.js",
"./noir": "./dest/noir/index.js",
"./testing": "./dest/testing/index.js",
"./array": "./dest/array/index.js"
"./array": "./dest/array/index.js",
"./validation": "./dest/validation/index.js"
},
"scripts": {
"build": "yarn clean && tsc -b",
Expand Down
Loading

0 comments on commit b5b3839

Please sign in to comment.