@openzeppelin/wizard 0.4.1 → 0.4.3

package/src/zip-foundry.ts

@@ -24,10 +24,14 @@ const test = (c: Contract, opts?: GenericOptions) => {
   );
 
   function getImports(c: Contract) {
-    return [
-      'import "forge-std/Test.sol";',
-      `import "../src/${c.name}.sol";`,
+    const result = [
+      'import {Test} from "forge-std/Test.sol";',
     ];
+    if (c.upgradeable) {
+      result.push('import {Upgrades} from "openzeppelin-foundry-upgrades/Upgrades.sol";');
+    }
+    result.push(`import {${c.name}} from "src/${c.name}.sol";`);
+    return result;
   }
 
   function getTestCase(c: Contract) {
@@ -40,10 +44,8 @@ const test = (c: Contract, opts?: GenericOptions) => {
         ],
         [
           'function setUp() public {',
-          getAddressVariables(args),
-          [
-            `instance = new ${getDeploymentCall(c, args)};`,
-          ],
+          getAddressVariables(c, args),
+          getDeploymentCode(c, args),
           '}',
         ],
         getContractSpecificTestFunction(),
@@ -52,11 +54,45 @@ const test = (c: Contract, opts?: GenericOptions) => {
     ];
   }
 
-  function getAddressVariables(args: string[]): Lines[] {
+  function getDeploymentCode(c: Contract, args: string[]): Lines[] {
+    if (c.upgradeable) {
+      if (opts?.upgradeable === 'transparent') {
+        return [
+          `address proxy = Upgrades.deployTransparentProxy(`,
+          [
+            `"${c.name}.sol",`,
+            `initialOwner,`,
+            `abi.encodeCall(${c.name}.initialize, (${args.join(', ')}))`
+          ],
+          ');',
+          `instance = ${c.name}(proxy);`,
+        ];
+      } else {
+        return [
+          `address proxy = Upgrades.deployUUPSProxy(`,
+          [
+            `"${c.name}.sol",`,
+            `abi.encodeCall(${c.name}.initialize, (${args.join(', ')}))`
+          ],
+          ');',
+          `instance = ${c.name}(proxy);`,
+        ];
+      }
+    } else {
+      return [
+        `instance = new ${c.name}(${args.join(', ')});`,
+      ];
+    }
+  }
+
+  function getAddressVariables(c: Contract, args: string[]): Lines[] {
     const vars = [];
-    for (let i = 0; i < args.length; i++) {
-      // use i + 1 as the private key since it must be non-zero
-      vars.push(`address ${args[i]} = vm.addr(${i + 1});`);
+    let i = 1; // private key index starts from 1 since it must be non-zero
+    if (c.upgradeable && opts?.upgradeable === 'transparent' && !args.includes('initialOwner')) {
+      vars.push(`address initialOwner = vm.addr(${i++});`);
+    }
+    for (const arg of args) {
+      vars.push(`address ${arg} = vm.addr(${i++});`);
     }
     return vars;
   }
@@ -67,7 +103,7 @@ const test = (c: Contract, opts?: GenericOptions) => {
         case 'ERC20':
         case 'ERC721':
           return [
-            'function testName() public {',
+            'function testName() public view {',
             [
               `assertEq(instance.name(), "${opts.name}");`
             ],
@@ -76,7 +112,7 @@ const test = (c: Contract, opts?: GenericOptions) => {
 
         case 'ERC1155':
           return [
-            'function testUri() public {',
+            'function testUri() public view {',
             [
               `assertEq(instance.uri(0), "${opts.uri}");`
             ],
@@ -111,11 +147,7 @@ function getAddressArgs(c: Contract): string[] {
   return args;
 }
 
-function getDeploymentCall(c: Contract, args: string[]): string {
-  return `${c.name}(${args.join(', ')})`;
-}
-
-const script = (c: Contract) => {
+const script = (c: Contract, opts?: GenericOptions) => {
   return formatLinesWithSpaces(
     2,
     ...spaceBetween(
@@ -126,18 +158,24 @@ const script = (c: Contract) => {
   );
 
   function getImports(c: Contract) {
-    return [
-      'import "forge-std/Script.sol";',
-      `import "../src/${c.name}.sol";`,
+    const result = [
+      'import {Script} from "forge-std/Script.sol";',
+      'import {console} from "forge-std/console.sol";',
     ];
+    if (c.upgradeable) {
+      result.push('import {Upgrades} from "openzeppelin-foundry-upgrades/Upgrades.sol";');
+    }
+    result.push(`import {${c.name}} from "src/${c.name}.sol";`);
+    return result;
   }
 
   function getScript(c: Contract) {
     const args = getAddressArgs(c);
     const deploymentLines = [
       'vm.startBroadcast();',
-      `${c.name} instance = new ${getDeploymentCall(c, args)};`,
-      'console.log("Contract deployed to %s", address(instance));',
+      ...getAddressVariables(c, args),
+      ...getDeploymentCode(c, args),
+      `console.log("${c.upgradeable ? 'Proxy' : 'Contract'} deployed to %s", address(instance));`,
       'vm.stopBroadcast();',
     ];
     return [
@@ -156,15 +194,59 @@ const script = (c: Contract) => {
     ];
   }
 
-  function addTodoAndCommentOut(lines: string[]) {
+  function getDeploymentCode(c: Contract, args: string[]): Lines[] {
+    if (c.upgradeable) {
+      if (opts?.upgradeable === 'transparent') {
+        return [
+          `address proxy = Upgrades.deployTransparentProxy(`,
+          [
+            `"${c.name}.sol",`,
+            `initialOwner,`,
+            `abi.encodeCall(${c.name}.initialize, (${args.join(', ')}))`
+          ],
+          ');',
+          `${c.name} instance = ${c.name}(proxy);`,
+        ];
+      } else {
+        return [
+          `address proxy = Upgrades.deployUUPSProxy(`,
+          [
+            `"${c.name}.sol",`,
+            `abi.encodeCall(${c.name}.initialize, (${args.join(', ')}))`
+          ],
+          ');',
+          `${c.name} instance = ${c.name}(proxy);`,
+        ];
+      }
+    } else {
+      return [
+        `${c.name} instance = new ${c.name}(${args.join(', ')});`,
+      ];
+    }
+  }
+
+  function getAddressVariables(c: Contract, args: string[]): Lines[] {
+    const vars = [];
+    if (c.upgradeable && opts?.upgradeable === 'transparent' && !args.includes('initialOwner')) {
+      vars.push('address initialOwner = <Set initialOwner address here>;');
+    }
+    for (const arg of args) {
+      vars.push(`address ${arg} = <Set ${arg} address here>;`);
+    }
+    return vars;
+  }
+
+  function addTodoAndCommentOut(lines: Lines[]) {
     return [
-      '// TODO: Set addresses for the contract arguments below, then uncomment the following lines',
-      ...lines.map(l => `// ${l}`),
+      '// TODO: Set addresses for the variables below, then uncomment the following section:',
+      '/*',
+      ...lines,
+      '*/',
     ];
   }
 };
 
-const setupSh = `\
+const setupSh = (c: Contract) => `\
 #!/usr/bin/env bash
 
 # Check if git is installed
@@ -186,23 +268,48 @@ if ! [ -f "foundry.toml" ]
 then
   echo "Initializing Foundry project..."
 
+  # Backup Wizard template readme to avoid it being overwritten
+  mv README.md README-oz.md
+
   # Initialize sample Foundry project
   forge init --force --no-commit --quiet
 
+${c.upgradeable ? `\
+  # Install OpenZeppelin Contracts and Upgrades
+  forge install OpenZeppelin/openzeppelin-contracts-upgradeable@v${contracts.version} --no-commit --quiet
+  forge install OpenZeppelin/openzeppelin-foundry-upgrades --no-commit --quiet\
+` : `\
   # Install OpenZeppelin Contracts
-  forge install OpenZeppelin/openzeppelin-contracts@v${contracts.version} --no-commit --quiet
+  forge install OpenZeppelin/openzeppelin-contracts@v${contracts.version} --no-commit --quiet\
+`}
 
-  # Remove template contracts
+  # Remove unneeded Foundry template files
   rm src/Counter.sol
   rm script/Counter.s.sol
   rm test/Counter.t.sol
+  rm README.md
+
+  # Restore Wizard template readme
+  mv README-oz.md README.md
 
   # Add remappings
   if [ -f "remappings.txt" ]
   then
     echo "" >> remappings.txt
   fi
-  echo "@openzeppelin/contracts/=lib/openzeppelin-contracts/contracts/" >> remappings.txt
+${c.upgradeable ? `\
+  echo "@openzeppelin/contracts/=lib/openzeppelin-contracts-upgradeable/lib/openzeppelin-contracts/contracts/" >> remappings.txt
+  echo "@openzeppelin/contracts-upgradeable/=lib/openzeppelin-contracts-upgradeable/contracts/" >> remappings.txt
+
+  # Add settings in foundry.toml
+  echo "" >> foundry.toml
+  echo "ffi = true" >> foundry.toml
+  echo "ast = true" >> foundry.toml
+  echo "build_info = true" >> foundry.toml
+  echo "extra_output = [\\"storageLayout\\"]" >> foundry.toml\
+` : `\
+  echo "@openzeppelin/contracts/=lib/openzeppelin-contracts/contracts/" >> remappings.txt\
+`}
 
   # Perform initial git commit
   git add .
@@ -232,7 +339,7 @@ bash setup.sh
 ## Testing the contract
 
 \`\`\`
-forge test
+forge test${c.upgradeable ? ' --force' : ''}
 \`\`\`
 
 ## Deploying the contract
@@ -240,7 +347,7 @@ forge test
 You can simulate a deployment by running the script:
 
 \`\`\`
-forge script script/${c.name}.s.sol
+forge script script/${c.name}.s.sol${c.upgradeable ? ' --force' : ''}
 \`\`\`
 
 See [Solidity scripting guide](https://book.getfoundry.sh/tutorials/solidity-scripting) for more information.
@@ -251,8 +358,8 @@ export async function zipFoundry(c: Contract, opts?: GenericOptions) {
 
   zip.file(`src/${c.name}.sol`, printContract(c));
   zip.file(`test/${c.name}.t.sol`, test(c, opts));
-  zip.file(`script/${c.name}.s.sol`, script(c));
-  zip.file('setup.sh', setupSh);
+  zip.file(`script/${c.name}.s.sol`, script(c, opts));
+  zip.file('setup.sh', setupSh(c));
   zip.file('README.md', readme(c));
 
   return zip;