Use namespaced storage for variables when upgradeable (#704)

Co-authored-by: Eric Lau <[email protected]>
Co-authored-by: Ernesto García <[email protected]>

Author: 3 people

.changeset/rude-nails-mate.md

@@ -0,0 +1,7 @@
+---
+'@openzeppelin/wizard': minor
+---
+
+**Breaking changes**: Use namespaced storage instead of state variables when upgradeability is enabled.
+  - For ERC-20, use namespaced storage for `tokenBridge` when cross-chain bridging is set to `'custom'` and upgradeability is enabled.
+  - For ERC-721, use namespaced storage for `_nextTokenId` when mintable, auto increment IDs, and upgradeability are enabled.

package.json

@@ -28,6 +28,9 @@
       "**/typescript"
     ]
   },
+  "resolutions": {
+    "ethereum-cryptography": "^3.2.0"
+  },
   "devDependencies": {
     "@eslint/js": "^9.21.0",
     "concurrently": "^9.1.2",

packages/common/src/ai/descriptions/solidity.ts

@@ -18,6 +18,8 @@ export const solidityCommonDescriptions = {
     'The type of access control to provision. Ownable is a simple mechanism with a single account authorized for all privileged actions. Roles is a flexible mechanism with a separate role for each privileged action. A role can have many authorized accounts. Managed enables a central contract to define a policy that allows certain callers to access certain functions.',
   upgradeable:
     'Whether the smart contract is upgradeable. Transparent uses more complex proxy with higher overhead, requires less changes in your contract. Can also be used with beacons. UUPS uses simpler proxy with less overhead, requires including extra code in your contract. Allows flexibility for authorizing upgrades.',
+  namespacePrefix:
+    'The prefix for ERC-7201 namespace identifiers. It should be derived from the project name or a unique naming convention specific to the project. Used only if the contract includes storage variables and upgradeability is enabled. Default is "myProject".',
 };
 
 export const solidityERC20Descriptions = {

packages/core/solidity/package.json

@@ -21,6 +21,9 @@
     "test:watch": "ava --watch",
     "update-env": "rm ./src/environments/hardhat/package-lock.json && npm install --package-lock-only --prefix ./src/environments/hardhat && rm ./src/environments/hardhat/upgradeable/package-lock.json && npm install --package-lock-only --prefix ./src/environments/hardhat/upgradeable"
   },
+  "dependencies": {
+    "ethereum-cryptography": "^3.2.0"
+  },
   "devDependencies": {
     "@openzeppelin/community-contracts": "git+https://github.com/OpenZeppelin/openzeppelin-community-contracts.git#b0ddd27",
     "@openzeppelin/contracts": "^5.4.0",

packages/core/solidity/src/contract.test.ts

@@ -135,17 +135,38 @@ test('contract with overridden function with code', t => {
 
 test('contract with one variable', t => {
   const Foo = new ContractBuilder('Foo');
-  Foo.addVariable('uint value = 42;');
+  Foo.addStateVariable('uint value = 42;', false);
   t.snapshot(printContract(Foo));
 });
 
 test('contract with two variables', t => {
   const Foo = new ContractBuilder('Foo');
-  Foo.addVariable('uint value = 42;');
-  Foo.addVariable('string name = "john";');
+  Foo.addStateVariable('uint value = 42;', false);
+  Foo.addStateVariable('string name = "john";', false);
   t.snapshot(printContract(Foo));
 });
 
+test('contract with immutable', t => {
+  const Foo = new ContractBuilder('Foo');
+  Foo.addConstantOrImmutableOrErrorDefinition('uint immutable value = 42;');
+  t.snapshot(printContract(Foo));
+});
+
+test('contract with constant and comment', t => {
+  const Foo = new ContractBuilder('Foo');
+  Foo.addConstantOrImmutableOrErrorDefinition('uint constant value = 42;', ['// This is a comment']);
+  t.snapshot(printContract(Foo));
+});
+
+test('contract with variable and upgradeable - error', t => {
+  const Foo = new ContractBuilder('Foo');
+  const error = t.throws(() => Foo.addStateVariable('uint value = 42;', true));
+  t.is(
+    (error as Error).message,
+    'State variables should not be used when upgradeable is true. Set namespaced storage instead.',
+  );
+});
+
 test('name with special characters', t => {
   const Foo = new ContractBuilder('foo bar baz');
   t.snapshot(printContract(Foo));

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

@@ -262,6 +262,33 @@ Generated by [AVA](https://avajs.dev).
     }␊
     `
 
+## contract with immutable
+
+> Snapshot 1
+
+    `// SPDX-License-Identifier: MIT␊
+    // Compatible with OpenZeppelin Contracts ^5.4.0␊
+    pragma solidity ^0.8.27;␊
+    ␊
+    contract Foo {␊
+        uint immutable value = 42;␊
+    }␊
+    `
+
+## contract with constant and comment
+
+> Snapshot 1
+
+    `// SPDX-License-Identifier: MIT␊
+    // Compatible with OpenZeppelin Contracts ^5.4.0␊
+    pragma solidity ^0.8.27;␊
+    ␊
+    contract Foo {␊
+        // This is a comment␊
+        uint constant value = 42;␊
+    }␊
+    `
+
 ## name with special characters
 
 > Snapshot 1

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

@@ -7,9 +7,10 @@ export interface Contract {
   natspecTags: NatspecTag[];
   imports: ImportContract[];
   functions: ContractFunction[];
+  structs: ContractStruct[];
   constructorCode: string[];
   constructorArgs: FunctionArgument[];
-  variables: string[];
+  variableOrErrorDefinitions: VariableOrErrorDefinition[];
   upgradeable: boolean;
 }
 
@@ -52,7 +53,13 @@ export interface ContractFunction extends BaseFunction {
   comments: string[];
 }
 
-export type FunctionKind = 'internal' | 'public';
+export interface ContractStruct {
+  name: string;
+  comments: string[];
+  variables: string[];
+}
+
+export type FunctionKind = 'internal' | 'public' | 'private';
 export type FunctionMutability = (typeof mutabilityRank)[number];
 
 // Order is important
@@ -72,6 +79,11 @@ export interface NatspecTag {
   value: string;
 }
 
+export interface VariableOrErrorDefinition {
+  code: string;
+  comments?: string[];
+}
+
 export class ContractBuilder implements Contract {
   readonly name: string;
   license: string = 'MIT';
@@ -82,10 +94,11 @@ export class ContractBuilder implements Contract {
 
   readonly constructorArgs: FunctionArgument[] = [];
   readonly constructorCode: string[] = [];
-  readonly variableSet: Set<string> = new Set();
 
+  readonly variableOrErrorMap: Map<string, VariableOrErrorDefinition> = new Map<string, VariableOrErrorDefinition>();
   private parentMap: Map<string, Parent> = new Map<string, Parent>();
   private functionMap: Map<string, ContractFunction> = new Map();
+  private structMap: Map<string, ContractStruct> = new Map();
 
   constructor(name: string) {
     this.name = toIdentifier(name, true);
@@ -113,8 +126,12 @@ export class ContractBuilder implements Contract {
     return [...this.functionMap.values()];
   }
 
-  get variables(): string[] {
-    return [...this.variableSet];
+  get structs(): ContractStruct[] {
+    return [...this.structMap.values()];
+  }
+
+  get variableOrErrorDefinitions(): VariableOrErrorDefinition[] {
+    return [...this.variableOrErrorMap.values()];
   }
 
   addParent(contract: ImportContract, params: Value[] = []): boolean {
@@ -172,6 +189,19 @@ export class ContractBuilder implements Contract {
     }
   }
 
+  private addStruct(_struct: ContractStruct): ContractStruct {
+    const got = this.structMap.get(_struct.name);
+    if (got !== undefined) {
+      return got;
+    } else {
+      const struct: ContractStruct = {
+        ..._struct,
+      };
+      this.structMap.set(_struct.name, struct);
+      return struct;
+    }
+  }
+
   addConstructorArgument(arg: FunctionArgument) {
     this.constructorArgs.push(arg);
   }
@@ -212,11 +242,36 @@ export class ContractBuilder implements Contract {
   }
 
   /**
-   * Note: The type in the variable is not currently transpiled, even if it refers to a contract
+   * Note: The type in the code is not currently transpiled, even if it refers to a contract
+   */
+  addStateVariable(code: string, upgradeable: boolean): boolean {
+    if (upgradeable) {
+      throw new Error('State variables should not be used when upgradeable is true. Set namespaced storage instead.');
+    } else {
+      return this._addVariableOrErrorDefinition({ code });
+    }
+  }
+
+  /**
+   * Note: The type in the code is not currently transpiled, even if it refers to a contract
    */
-  addVariable(code: string): boolean {
-    const present = this.variableSet.has(code);
-    this.variableSet.add(code);
+  addConstantOrImmutableOrErrorDefinition(code: string, comments?: string[]): boolean {
+    return this._addVariableOrErrorDefinition({ code, comments });
+  }
+
+  private _addVariableOrErrorDefinition(variableOrErrorDefinition: VariableOrErrorDefinition): boolean {
+    const present = this.variableOrErrorMap.has(variableOrErrorDefinition.code);
+    this.variableOrErrorMap.set(variableOrErrorDefinition.code, variableOrErrorDefinition);
+    return !present;
+  }
+
+  addStructVariable(baseStruct: ContractStruct, code: string): boolean {
+    let struct = this.structMap.get(baseStruct.name);
+    if (!struct) {
+      struct = this.addStruct(baseStruct);
+    }
+    const present = struct.variables.includes(code);
+    if (!present) struct.variables.push(code);
     return !present;
   }
 }

packages/core/solidity/src/contract.ts

@@ -22,16 +22,14 @@ export interface ERC1155Options extends CommonOptions {
 }
 
 export const defaults: Required<ERC1155Options> = {
+  ...commonDefaults,
   name: 'MyToken',
   uri: '',
   burnable: false,
   pausable: false,
   mintable: false,
   supply: false,
   updatableUri: true,
-  access: commonDefaults.access,
-  upgradeable: commonDefaults.upgradeable,
-  info: commonDefaults.info,
 } as const;
 
 function withDefaults(opts: ERC1155Options): Required<ERC1155Options> {

packages/core/solidity/src/erc1155.ts

@@ -863,7 +863,14 @@ Generated by [AVA](https://avajs.dev).
     import {Initializable} from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";␊
     ␊
     contract MyToken is Initializable, ERC20Upgradeable, ERC20BridgeableUpgradeable, ERC20PermitUpgradeable {␊
-        address public tokenBridge;␊
+        /// @custom:storage-location erc7201:myProject.MyToken␊
+        struct MyTokenStorage {␊
+            address tokenBridge;␊
+        }␊
+    ␊
+        // keccak256(abi.encode(uint256(keccak256("myProject.MyToken")) - 1)) & ~bytes32(uint256(0xff))␊
+        bytes32 private constant MYTOKEN_STORAGE_LOCATION = 0xfbb7c9e4123fcf4b1aad53c70358f7b1c1d7cf28092f5178b53e55db565e9200;␊
+    ␊
         error Unauthorized();␊
     ␊
         /// @custom:oz-upgrades-unsafe-allow constructor␊
@@ -877,11 +884,17 @@ Generated by [AVA](https://avajs.dev).
             __ERC20Permit_init("MyToken");␊
     ␊
             require(tokenBridge_ != address(0), "Invalid tokenBridge_ address");␊
-            tokenBridge = tokenBridge_;␊
+            MyTokenStorage storage $ = _getMyTokenStorage();␊
+            $.tokenBridge = tokenBridge_;␊
         }␊
     ␊
         function _checkTokenBridge(address caller) internal view override {␊
-            if (caller != tokenBridge) revert Unauthorized();␊
+            MyTokenStorage storage $ = _getMyTokenStorage();␊
+            if (caller != $.tokenBridge) revert Unauthorized();␊
+        }␊
+    ␊
+        function _getMyTokenStorage() private pure returns (MyTokenStorage storage $) {␊
+            assembly { $.slot := MYTOKEN_STORAGE_LOCATION }␊
         }␊
     }␊
     `