refactor Account upgradeability

Author: Amxx

packages/core/solidity/src/account.test.ts.md

@@ -2,13 +2,11 @@ import { ContractBuilder } from './contract';
 import type { Contract } from './contract';
 import { defineFunctions } from './utils/define-functions';
 import { printContract } from './print';
-import { makeUpgradeable } from './helpers';
 import { defaults as commonDefaults, withCommonDefaults, type CommonOptions } from './common-options';
+import { upgradeableName } from './options';
 import { setInfo } from './set-info';
 import {
-  addLockingConstructorAllowReachable,
   addSigner,
-  signerArgs,
   signerFunctions,
   signers,
   type SignerOptions,
@@ -75,6 +73,7 @@ export function buildAccount(opts: AccountOptions): Contract {
     c.addImportOnly({
       name: 'PackedUserOperation',
       path: '@openzeppelin/contracts/interfaces/draft-IERC4337.sol',
+      transpiled: false, // PackedUserOperation doesn't start with "I" so its not recognized as an "interface object"
     });
   }
 
@@ -86,16 +85,16 @@ function addParents(c: ContractBuilder, opts: AccountOptions): void {
   c.addParent({
     name: 'Account',
     path: `@openzeppelin/contracts/account/Account.sol`,
+    transpiled: false,
   });
-  c.addOverride({ name: 'Account' }, functions._validateUserOp);
+  c.addOverride({ name: 'Account', transpiled: false }, functions._validateUserOp);
 
   if (opts.signatureValidation === 'ERC7739') addEIP712(c, opts);
 
   // Extensions
   addSignatureValidation(c, opts);
   addERC7579Modules(c, opts);
   addSigner(c, opts.signer ?? false, opts.upgradeable ?? false);
-  addSignerInitializer(c, opts);
   addMultisigFunctions(c, opts);
   addBatchedExecution(c, opts);
   addERC721Holder(c, opts);
@@ -108,6 +107,7 @@ function addSignatureValidation(c: ContractBuilder, opts: AccountOptions) {
       c.addParent({
         name: 'ERC7739',
         path: '@openzeppelin/contracts/utils/cryptography/signers/draft-ERC7739.sol',
+        transpiled: false,
       });
       break;
     case 'ERC1271':
@@ -133,6 +133,7 @@ function addERC721Holder(c: ContractBuilder, opts: AccountOptions): void {
   c.addParent({
     name: 'ERC721Holder',
     path: '@openzeppelin/contracts/token/ERC721/utils/ERC721Holder.sol',
+    transpiled: false,
   });
 }
 
@@ -141,6 +142,7 @@ function addERC1155Holder(c: ContractBuilder, opts: AccountOptions): void {
   c.addParent({
     name: 'ERC1155Holder',
     path: '@openzeppelin/contracts/token/ERC1155/utils/ERC1155Holder.sol',
+    transpiled: false,
   });
 }
 
@@ -150,8 +152,9 @@ function addBatchedExecution(c: ContractBuilder, opts: AccountOptions): void {
   c.addParent({
     name: 'ERC7821',
     path: '@openzeppelin/contracts/account/extensions/draft-ERC7821.sol',
+    transpiled: false,
   });
-  c.addOverride({ name: 'ERC7821' }, functions._erc7821AuthorizedExecutor);
+  c.addOverride({ name: 'ERC7821', transpiled: false }, functions._erc7821AuthorizedExecutor);
   c.setFunctionBody(
     ['return caller == address(entryPoint()) || super._erc7821AuthorizedExecutor(caller, mode, executionData);'],
     functions._erc7821AuthorizedExecutor,
@@ -162,59 +165,38 @@ function addERC7579Modules(c: ContractBuilder, opts: AccountOptions): void {
   if (!opts.ERC7579Modules) return;
 
   // Base AccountERC7579 account (upgradeable or not)
-  const name = makeUpgradeable('AccountERC7579', opts.upgradeable);
+  const name = 'AccountERC7579';
 
   c.addParent({
-    name: makeUpgradeable(opts.ERC7579Modules, opts.upgradeable),
-    path: makeUpgradeable(
-      `@openzeppelin/contracts/account/extensions/draft-${opts.ERC7579Modules}.sol`,
-      opts.upgradeable,
-    ),
+    name: opts.ERC7579Modules,
+    path: `@openzeppelin/contracts/account/extensions/draft-${opts.ERC7579Modules}.sol`,
   });
   if (opts.ERC7579Modules !== 'AccountERC7579') {
     c.addImportOnly({
-      name: makeUpgradeable('AccountERC7579', opts.upgradeable),
-      path: makeUpgradeable('@openzeppelin/contracts/account/extensions/draft-AccountERC7579.sol', opts.upgradeable),
+      name: 'AccountERC7579',
+      path: '@openzeppelin/contracts/account/extensions/draft-AccountERC7579.sol',
     });
   }
 
   // Accounts that use ERC7579 without a signer must be constructed with at least one module (executor of validation)
   if (!opts.signer) {
-    const args = [
-      { type: 'uint256', name: 'moduleTypeId' },
-      { type: 'address', name: 'module' },
-      { type: 'bytes calldata', name: 'initData' },
-    ];
-    const body = [
-      'require(moduleTypeId == MODULE_TYPE_VALIDATOR || moduleTypeId == MODULE_TYPE_EXECUTOR);',
-      '_installModule(moduleTypeId, module, initData);',
-    ];
-    if (opts.upgradeable) {
-      c.addParent({
-        name: 'Initializable',
-        path: '@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol',
-      });
-      addLockingConstructorAllowReachable(c);
-
-      const fn = { name: 'initialize', kind: 'public' as const, args };
-      c.addModifier('initializer', fn);
-      c.setFunctionBody(body, fn);
-    } else {
-      for (const arg of args) c.addConstructorArgument(arg);
-      for (const line of body) c.addConstructorCode(line);
-    }
+    c.addConstructorArgument({ type: 'uint256', name: 'moduleTypeId' });
+    c.addConstructorArgument({ type: 'address', name: 'module' });
+    c.addConstructorArgument({ type: 'bytes calldata', name: 'initData' });
+    c.addConstructorCode('require(moduleTypeId == MODULE_TYPE_VALIDATOR || moduleTypeId == MODULE_TYPE_EXECUTOR);');
+    c.addConstructorCode('_installModule(moduleTypeId, module, initData);');
   }
 
   // isValidSignature override
   c.addOverride({ name }, functions.isValidSignature);
   if (opts.signatureValidation === 'ERC7739') {
-    c.addOverride({ name: 'ERC7739' }, functions.isValidSignature);
+    c.addOverride({ name: 'ERC7739', transpiled: false }, functions.isValidSignature);
     c.setFunctionBody(
       [
         '// ERC-7739 can return the ERC-1271 magic value, 0xffffffff (invalid) or 0x77390001 (detection).',
         '// If the returned value is 0xffffffff, fallback to ERC-7579 validation.',
         'bytes4 erc7739magic = ERC7739.isValidSignature(hash, signature);',
-        `return erc7739magic == bytes4(0xffffffff) ? ${name}.isValidSignature(hash, signature) : erc7739magic;`,
+        `return erc7739magic == bytes4(0xffffffff) ? ${opts.upgradeable ? upgradeableName(name) : name}.isValidSignature(hash, signature) : erc7739magic;`,
       ],
       functions.isValidSignature,
     );
@@ -224,51 +206,6 @@ function addERC7579Modules(c: ContractBuilder, opts: AccountOptions): void {
   c.addOverride({ name }, functions._validateUserOp);
 }
 
-function addSignerInitializer(c: ContractBuilder, opts: AccountOptions): void {
-  if (opts.upgradeable) {
-    if (!opts.signer) {
-      addLockingConstructorAllowReachable(c);
-    }
-    return; // Initializer added in signer.ts
-  }
-  if (!opts.signer || opts.signer === 'ERC7702') return; // No initialization required
-
-  c.addParent({
-    name: 'Initializable',
-    path: opts.upgradeable
-      ? '@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol'
-      : '@openzeppelin/contracts/proxy/utils/Initializable.sol',
-  });
-
-  addLockingConstructorAllowReachable(c, [
-    '// Accounts are typically deployed and initialized as clones during their first user op,',
-    '// therefore, initializers are disabled for the implementation contract',
-  ]);
-
-  const fn = { name: 'initialize', kind: 'public' as const, args: signerArgs[opts.signer] };
-  c.addModifier('initializer', fn);
-
-  switch (opts.signer) {
-    case 'Multisig':
-      c.addFunctionCode(`_addSigners(${signerArgs[opts.signer][0]!.name});`, fn);
-      c.addFunctionCode(`_setThreshold(${signerArgs[opts.signer][1]!.name});`, fn);
-      break;
-    case 'MultisigWeighted':
-      c.addFunctionCode(`_addSigners(${signerArgs[opts.signer][0]!.name});`, fn);
-      c.addFunctionCode(
-        `_setSignerWeights(${signerArgs[opts.signer][0]!.name}, ${signerArgs[opts.signer][1]!.name});`,
-        fn,
-      );
-      c.addFunctionCode(`_setThreshold(${signerArgs[opts.signer][2]!.name});`, fn);
-      break;
-    case 'ECDSA':
-    case 'P256':
-    case 'RSA':
-      c.addFunctionCode(`_setSigner(${signerArgs[opts.signer].map(({ name }) => name).join(', ')});`, fn);
-      break;
-  }
-}
-
 function addMultisigFunctions(c: ContractBuilder, opts: AccountOptions): void {
   switch (opts.signer) {
     case 'MultisigWeighted':
@@ -296,6 +233,7 @@ function addEIP712(c: ContractBuilder, opts: AccountOptions): void {
     {
       name: 'EIP712',
       path: '@openzeppelin/contracts/utils/cryptography/EIP712.sol',
+      transpiled: false, // do not use the upgradeable variant for in Accounts
     },
     [opts.name, '1'],
   );
@@ -309,21 +247,22 @@ function overrideRawSignatureValidation(c: ContractBuilder, opts: AccountOptions
   // to provide a custom validation logic
   if (!opts.signer && !opts.ERC7579Modules) {
     // Custom validation logic
-    c.addOverride({ name: 'Account' }, signerFunctions._rawSignatureValidation);
+    c.addOverride({ name: 'Account', transpiled: false }, signerFunctions._rawSignatureValidation);
     c.setFunctionBody(['// Custom validation logic', 'return false;'], signerFunctions._rawSignatureValidation);
   }
 
   // Disambiguate between Signer and AccountERC7579
   if (opts.signer && opts.ERC7579Modules) {
-    const accountName = makeUpgradeable('AccountERC7579', opts.upgradeable);
-    const signerName = makeUpgradeable(`Signer${opts.signer}`, opts.upgradeable);
+    const accountName = opts.upgradeable ? upgradeableName('AccountERC7579') : 'AccountERC7579';
+    const signerName = opts.upgradeable ? upgradeableName(`Signer${opts.signer}`) : `Signer${opts.signer}`;
 
     c.addImportOnly({
       name: 'AbstractSigner',
       path: '@openzeppelin/contracts/utils/cryptography/signers/AbstractSigner.sol',
+      transpiled: false,
     });
-    c.addOverride({ name: 'AbstractSigner' }, signerFunctions._rawSignatureValidation);
-    c.addOverride({ name: accountName }, signerFunctions._rawSignatureValidation);
+    c.addOverride({ name: 'AbstractSigner', transpiled: false }, signerFunctions._rawSignatureValidation);
+    c.addOverride({ name: 'AccountERC7579' }, signerFunctions._rawSignatureValidation);
     c.setFunctionComments(
       [
         `// IMPORTANT: Make sure ${signerName} is most derived than ${accountName}`,
@@ -337,8 +276,8 @@ function overrideRawSignatureValidation(c: ContractBuilder, opts: AccountOptions
     // Base override for `_rawSignatureValidation` given MultiSignerERC7913Weighted is MultiSignerERC7913
     if (opts.signer === 'MultisigWeighted') {
       c.addImportOnly({
-        name: makeUpgradeable(signers.Multisig.name, opts.upgradeable),
-        path: makeUpgradeable(signers.Multisig.path, opts.upgradeable),
+        name: signers.Multisig.name,
+        path: signers.Multisig.path,
       });
     }
   }

packages/core/solidity/src/account.test.ts.snap

@@ -9,11 +9,8 @@ export interface Contract {
   functions: ContractFunction[];
   constructorCode: string[];
   constructorArgs: FunctionArgument[];
-  constructorComments: string[];
   variables: string[];
-  shouldAutoTranspileImports: boolean;
-  shouldInstallContractsUpgradeable: boolean;
-  shouldUseUpgradesPluginsForProxyDeployment: boolean;
+  upgradeable: boolean;
 }
 
 export type Value = string | number | { lit: string } | { note: string; value: Value };
@@ -78,17 +75,13 @@ export interface NatspecTag {
 export class ContractBuilder implements Contract {
   readonly name: string;
   license: string = 'MIT';
-
-  shouldAutoTranspileImports: boolean = false;
-  shouldInstallContractsUpgradeable: boolean = false;
-  shouldUseUpgradesPluginsForProxyDeployment: boolean = false;
+  upgradeable = false;
 
   readonly using: Using[] = [];
   readonly natspecTags: NatspecTag[] = [];
 
   readonly constructorArgs: FunctionArgument[] = [];
   readonly constructorCode: string[] = [];
-  readonly constructorComments: string[] = [];
   readonly variableSet: Set<string> = new Set();
 
   private parentMap: Map<string, Parent> = new Map<string, Parent>();
@@ -187,15 +180,6 @@ export class ContractBuilder implements Contract {
     this.constructorCode.push(code);
   }
 
-  addConstructorComment(comment: string) {
-    if (this.shouldAutoTranspileImports) {
-      throw new Error(
-        'Constructor comments are not supported when `shouldAutoTranspileImports` is true, since constructor will be transformed into an initializer',
-      );
-    }
-    this.constructorComments.push(comment);
-  }
-
   addFunctionCode(code: string, baseFn: BaseFunction, mutability?: FunctionMutability) {
     const fn = this.addFunction(baseFn);
     if (fn.final) {

packages/core/solidity/src/account.ts

@@ -3,15 +3,15 @@ import path from 'path';
 import type { Contract, ReferencedContract, ImportContract } from './contract';
 import { inferTranspiled } from './infer-transpiled';
 
-const upgradeableName = (n: string) => {
+export function upgradeableName(n: string) {
   if (n === 'Initializable') {
     return n;
   } else {
     return n.replace(/(Upgradeable)?(?=\.|$)/, 'Upgradeable');
   }
 };
 
-const upgradeableImport = (p: ImportContract): ImportContract => {
+export function upgradeableImport(p: ImportContract): ImportContract {
   const { dir, ext, name } = path.parse(p.path);
   // Use path.posix to get forward slashes
   return {
@@ -30,19 +30,19 @@ export interface Options {
 }
 
 export interface Helpers extends Required<Options> {
-  shouldUseInitializers: boolean;
+  upgradeable: boolean;
   transformName: (name: ReferencedContract) => string;
+  transformImport: (name: ImportContract) => ImportContract;
 }
 
 export function withHelpers(contract: Contract, opts: Options = {}): Helpers {
-  const shouldAutoTranspileImports = contract.shouldAutoTranspileImports;
-  const transformName = (n: ReferencedContract) =>
-    shouldAutoTranspileImports && inferTranspiled(n) ? upgradeableName(n.name) : n.name;
+  const contractUpgradeable = contract.upgradeable;
   return {
-    shouldUseInitializers: shouldAutoTranspileImports,
-    transformName,
-    transformImport: p1 => {
-      const p2 = shouldAutoTranspileImports && inferTranspiled(p1) ? upgradeableImport(p1) : p1;
+    upgradeable: contractUpgradeable,
+    transformName: (n: ReferencedContract) =>
+      contractUpgradeable && inferTranspiled(n) ? upgradeableName(n.name) : n.name,
+    transformImport: (p1: ImportContract) => {
+      const p2 = contractUpgradeable && inferTranspiled(p1) ? upgradeableImport(p1) : p1;
       return opts.transformImport?.(p2) ?? p2;
     },
   };