Use namespaced storage for variables when upgradeable

.changeset/rude-nails-mate.md

@@ -0,0 +1,5 @@
+---
+'@openzeppelin/wizard': minor
+---
+
+Support for ERC721 with Mintable → Auto Increment IDs and Upgradeability, using \_nextTokenId in namespaced storage for safe upgradeable minting.

package.json

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

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.ts

@@ -7,6 +7,7 @@ export interface Contract {
   natspecTags: NatspecTag[];
   imports: ImportContract[];
   functions: ContractFunction[];
+  structs: ContractStruct[];
   constructorCode: string[];
   constructorArgs: FunctionArgument[];
   variables: string[];
@@ -52,6 +53,12 @@ export interface ContractFunction extends BaseFunction {
   comments: string[];
 }
 
+export interface ContractStruct {
+  name: string;
+  comments: string[];
+  variables: string[];
+}
+
 export type FunctionKind = 'internal' | 'public';
 export type FunctionMutability = (typeof mutabilityRank)[number];
 
@@ -86,6 +93,7 @@ export class ContractBuilder implements Contract {
 
   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,6 +121,10 @@ export class ContractBuilder implements Contract {
     return [...this.functionMap.values()];
   }
 
+  get structs(): ContractStruct[] {
+    return [...this.structMap.values()];
+  }
+
   get variables(): string[] {
     return [...this.variableSet];
   }
@@ -172,6 +184,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);
   }
@@ -219,4 +244,14 @@ export class ContractBuilder implements Contract {
     this.variableSet.add(code);
     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/erc721.test.ts

@@ -133,6 +133,17 @@ testERC721('full upgradeable uups + managed', {
   access: 'managed',
 });
 
+testERC721('full upgradeable uups + managed + incremental', {
+  mintable: true,
+  enumerable: true,
+  pausable: true,
+  burnable: true,
+  incremental: true,
+  votes: true,
+  upgradeable: 'uups',
+  access: 'managed',
+});
+
 testAPIEquivalence('API default');
 
 testAPIEquivalence('API basic', { name: 'CustomToken', symbol: 'CTK' });

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

@@ -777,3 +777,98 @@ Generated by [AVA](https://avajs.dev).
         }␊
     }␊
     `
+
+## full upgradeable uups + managed + incremental
+
+> Snapshot 1
+
+    `// SPDX-License-Identifier: MIT␊
+    // Compatible with OpenZeppelin Contracts ^5.4.0␊
+    pragma solidity ^0.8.27;␊
+    ␊
+    import {AccessManagedUpgradeable} from "@openzeppelin/contracts-upgradeable/access/manager/AccessManagedUpgradeable.sol";␊
+    import {EIP712Upgradeable} from "@openzeppelin/contracts-upgradeable/utils/cryptography/EIP712Upgradeable.sol";␊
+    import {ERC721Upgradeable} from "@openzeppelin/contracts-upgradeable/token/ERC721/ERC721Upgradeable.sol";␊
+    import {ERC721BurnableUpgradeable} from "@openzeppelin/contracts-upgradeable/token/ERC721/extensions/ERC721BurnableUpgradeable.sol";␊
+    import {ERC721EnumerableUpgradeable} from "@openzeppelin/contracts-upgradeable/token/ERC721/extensions/ERC721EnumerableUpgradeable.sol";␊
+    import {ERC721PausableUpgradeable} from "@openzeppelin/contracts-upgradeable/token/ERC721/extensions/ERC721PausableUpgradeable.sol";␊
+    import {ERC721VotesUpgradeable} from "@openzeppelin/contracts-upgradeable/token/ERC721/extensions/ERC721VotesUpgradeable.sol";␊
+    import {Initializable} from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";␊
+    import {UUPSUpgradeable} from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";␊
+    ␊
+    contract MyToken is Initializable, ERC721Upgradeable, ERC721EnumerableUpgradeable, ERC721PausableUpgradeable, AccessManagedUpgradeable, ERC721BurnableUpgradeable, EIP712Upgradeable, ERC721VotesUpgradeable, UUPSUpgradeable {␊
+        /// @custom:storage-location erc7201:myProject.MyToken␊
+        struct Storage {␊
+            uint256 _nextTokenId;␊
+        }␊
+    ␊
+        bytes32 private constant MYTOKEN_STORAGE_LOCATION = 0xfbb7c9e4123fcf4b1aad53c70358f7b1c1d7cf28092f5178b53e55db565e9200;␊
+    ␊
+        /// @custom:oz-upgrades-unsafe-allow constructor␊
+        constructor() {␊
+            _disableInitializers();␊
+        }␊
+    ␊
+        function initialize(address initialAuthority) public initializer {␊
+            __ERC721_init("MyToken", "MTK");␊
+            __ERC721Enumerable_init();␊
+            __ERC721Pausable_init();␊
+            __AccessManaged_init(initialAuthority);␊
+            __ERC721Burnable_init();␊
+            __EIP712_init("MyToken", "1");␊
+            __ERC721Votes_init();␊
+            __UUPSUpgradeable_init();␊
+        }␊
+    ␊
+        function pause() public restricted {␊
+            _pause();␊
+        }␊
+    ␊
+        function unpause() public restricted {␊
+            _unpause();␊
+        }␊
+    ␊
+        function safeMint(address to) public restricted returns (uint256) {␊
+            Storage storage $ = _getStorage();␊
+            uint256 tokenId = $._nextTokenId++;␊
+            _safeMint(to, tokenId);␊
+            return tokenId;␊
+        }␊
+    ␊
+        function _getStorage() internal pure returns (Storage storage $) {␊
+            assembly { $.slot :=  MYTOKEN_STORAGE_LOCATION}␊
+        }␊
+    ␊
+        function _authorizeUpgrade(address newImplementation)␊
+            internal␊
+            override␊
+            restricted␊
+        {}␊
+    ␊
+        // The following functions are overrides required by Solidity.␊
+    ␊
+        function _update(address to, uint256 tokenId, address auth)␊
+            internal␊
+            override(ERC721Upgradeable, ERC721EnumerableUpgradeable, ERC721PausableUpgradeable, ERC721VotesUpgradeable)␊
+            returns (address)␊
+        {␊
+            return super._update(to, tokenId, auth);␊
+        }␊
+    ␊
+        function _increaseBalance(address account, uint128 value)␊
+            internal␊
+            override(ERC721Upgradeable, ERC721EnumerableUpgradeable, ERC721VotesUpgradeable)␊
+        {␊
+            super._increaseBalance(account, value);␊
+        }␊
+    ␊
+        function supportsInterface(bytes4 interfaceId)␊
+            public␊
+            view␊
+            override(ERC721Upgradeable, ERC721EnumerableUpgradeable)␊
+            returns (bool)␊
+        {␊
+            return super.supportsInterface(interfaceId);␊
+        }␊
+    }␊
+    `

packages/core/solidity/src/erc721.ts

@@ -1,4 +1,4 @@
-import type { BaseFunction, Contract } from './contract';
+import type { BaseFunction, Contract, ContractStruct } from './contract';
 import { ContractBuilder } from './contract';
 import type { Access } from './set-access-control';
 import { setAccessControl, requireAccessControl } from './set-access-control';
@@ -12,6 +12,7 @@ import { setInfo } from './set-info';
 import { printContract } from './print';
 import type { ClockMode } from './set-clock-mode';
 import { clockModeDefault, setClockMode } from './set-clock-mode';
+import { computeNamespacedStorageSlot, getNamespaceId } from './utils/namespaced-storage-generator';
 
 export interface ERC721Options extends CommonOptions {
   name: string;
@@ -99,7 +100,7 @@ export function buildERC721(opts: ERC721Options): Contract {
   }
 
   if (allOpts.mintable) {
-    addMintable(c, access, allOpts.incremental, allOpts.uriStorage);
+    addMintable(c, access, allOpts.incremental, allOpts.uriStorage, allOpts.upgradeable === 'uups');
   }
 
   if (allOpts.votes) {
@@ -174,14 +175,28 @@ function addBurnable(c: ContractBuilder) {
   });
 }
 
-function addMintable(c: ContractBuilder, access: Access, incremental = false, uriStorage = false) {
+function addMintable(c: ContractBuilder, access: Access, incremental = false, uriStorage = false, uups = false) {
   const fn = getMintFunction(incremental, uriStorage);
   requireAccessControl(c, fn, access, 'MINTER', 'minter');
-
   if (incremental) {
-    c.addVariable('uint256 private _nextTokenId;');
-    c.addFunctionCode('uint256 tokenId = _nextTokenId++;', fn);
-    c.addFunctionCode('_safeMint(to, tokenId);', fn);
+    if (!uups) {
+      c.addVariable('uint256 private _nextTokenId;');
+      c.addFunctionCode('uint256 tokenId = _nextTokenId++;', fn);
+      c.addFunctionCode('_safeMint(to, tokenId);', fn);
+    } else {
+      const storageFn = getStorageFunction();
+      const storageStruct = getStorageStruct(c.name);
+      const namespacedStorageName = `${c.name.toUpperCase()}_STORAGE_LOCATION`;
+      c.addStructVariable(storageStruct, 'uint256 _nextTokenId;');
+      c.addVariable(
+        `bytes32 private constant ${namespacedStorageName} = ${computeNamespacedStorageSlot(getNamespaceId(c.name))};`,
+      );
+      c.addFunctionCode(`assembly { $.slot :=  ${namespacedStorageName}}`, storageFn);
+
+      c.addFunctionCode('Storage storage $ = _getStorage();', fn);
+      c.addFunctionCode('uint256 tokenId = $._nextTokenId++;', fn);
+      c.addFunctionCode('_safeMint(to, tokenId);', fn);
+    }
   } else {
     c.addFunctionCode('_safeMint(to, tokenId);', fn);
   }
@@ -264,3 +279,24 @@ function getMintFunction(incremental: boolean, uriStorage: boolean): BaseFunctio
 
   return fn;
 }
+
+function getStorageFunction(): BaseFunction {
+  const fn: BaseFunction = {
+    name: '_getStorage',
+    kind: 'internal' as const,
+    mutability: 'pure',
+    args: [],
+    returns: ['Storage storage $'],
+  };
+
+  return fn;
+}
+
+function getStorageStruct(name: string) {
+  const struct: ContractStruct = {
+    name: 'Storage',
+    comments: [`/// @custom:storage-location erc7201:${getNamespaceId(name)}`],
+    variables: [],
+  };
+  return struct;
+}

packages/core/solidity/src/print.ts

@@ -6,6 +6,7 @@ import type {
   Value,
   NatspecTag,
   ImportContract,
+  ContractStruct,
 } from './contract';
 import type { Options, Helpers } from './options';
 import { withHelpers } from './options';
@@ -23,10 +24,9 @@ import { getCommunityContractsGitCommit } from './utils/community-contracts-git-
 export function printContract(contract: Contract, opts?: Options): string {
   const helpers = withHelpers(contract, opts);
 
+  const structs = contract.structs.map(_struct => printStruct(_struct));
   const fns = mapValues(sortedFunctions(contract), fns => fns.map(fn => printFunction(fn, helpers)));
-
   const hasOverrides = fns.override.some(l => l.length > 0);
-
   return formatLines(
     ...spaceBetween(
       [
@@ -42,6 +42,7 @@ export function printContract(contract: Contract, opts?: Options): string {
         [`contract ${contract.name}`, ...printInheritance(contract, helpers), '{'].join(' '),
 
         spaceBetween(
+          ...structs,
           contract.variables,
           printConstructor(contract, helpers),
           ...fns.code,
@@ -258,6 +259,20 @@ function printFunction2(
   return fn;
 }
 
+function printStruct(_struct: ContractStruct): Lines[] {
+  const [comments, kindedName, code] = [_struct.comments, _struct.name, _struct.variables];
+  const struct: Lines[] = [...comments];
+
+  const braces = code.length > 0 ? '{' : '{}';
+  struct.push([`struct ${kindedName}`, braces].join(' '));
+
+  if (code.length > 0) {
+    struct.push(code, '}');
+  }
+
+  return struct;
+}
+
 function printArgument(arg: FunctionArgument, { transformName }: Helpers): string {
   let type: string;
   if (typeof arg.type === 'string') {

packages/core/solidity/src/utils/namespaced-storage-generator.test.ts

@@ -0,0 +1,23 @@
+import test from 'ava';
+import { computeNamespacedStorageSlot } from './namespaced-storage-generator';
+
+test('namespaced storage slot generation', t => {
+  const cases = [
+    {
+      input: 'myProject.MyToken',
+      expected: '0xfbb7c9e4123fcf4b1aad53c70358f7b1c1d7cf28092f5178b53e55db565e9200',
+    },
+    {
+      input: 'myProject.token',
+      expected: '0x86796099e489af07082cc4e6965fe431aadf035a7b4d4b46f81d8dfb81822d00',
+    },
+    {
+      input: 'myProject.token123456',
+      expected: '0x824a9aeab482b3e91ee3e454c74509cca55ad57e0185a36d070359384be52800',
+    },
+  ];
+
+  for (const { input, expected } of cases) {
+    t.is(computeNamespacedStorageSlot(input), expected);
+  }
+});

packages/core/solidity/src/utils/namespaced-storage-generator.ts

@@ -0,0 +1,23 @@
+
+import { keccak256 } from 'ethereum-cryptography/keccak';
+import { hexToBytes, toHex, utf8ToBytes } from 'ethereum-cryptography/utils';
+
+export function getNamespaceId(name: string, namespace: string = 'myProject') {
+  return `${namespace}.${name}`;
+}
+
+/**
+ * Returns the ERC-7201 storage location for a given namespace id
+ */
+export function computeNamespacedStorageSlot(id: string): string {
+  const innerHash = keccak256(utf8ToBytes(id));
+  const minusOne = BigInt('0x' + toHex(innerHash)) - 1n;
+  const minusOneBytes = hexToBytes(minusOne.toString(16).padStart(64, '0'));
+
+  const outerHash = keccak256(minusOneBytes);
+
+  const mask = BigInt('0xff');
+  const masked = BigInt('0x' + toHex(outerHash)) & ~mask;
+
+  return '0x' + masked.toString(16).padStart(64, '0');
+}