@openzeppelin/upgrades-core 1.12.0 → 1.13.0

package/src/standalone.test.ts

@@ -1,4 +1,4 @@
-import _test, { TestInterface } from 'ava';
+import _test, { TestFn } from 'ava';
 import { artifacts } from 'hardhat';
 
 import { SolcInput, SolcOutput } from './solc-api';
@@ -9,7 +9,7 @@ interface Context {
   solcOutput: SolcOutput;
 }
 
-const test = _test as TestInterface<Context>;
+const test = _test as TestFn<Context>;
 
 test.before(async t => {
   const buildInfo = await artifacts.getBuildInfo('contracts/test/Standalone.sol:StandaloneV1');

package/src/storage/extract.ts

@@ -1,11 +1,18 @@
 import assert from 'assert';
-import { ContractDefinition, StructDefinition, EnumDefinition, TypeDescriptions } from 'solidity-ast';
+import {
+  ContractDefinition,
+  StructDefinition,
+  EnumDefinition,
+  TypeDescriptions,
+  VariableDeclaration,
+} from 'solidity-ast';
 import { isNodeType, findAll } from 'solidity-ast/utils';
-
 import { StorageLayout, TypeItem } from './layout';
 import { normalizeTypeIdentifier } from '../utils/type-id';
 import { SrcDecoder } from '../src-decoder';
 import { ASTDereferencer } from '../ast-dereferencer';
+import { mapValues } from '../utils/map-values';
+import { pick } from '../utils/pick';
 
 const currentLayoutVersion = '1.1';
 
@@ -17,62 +24,48 @@ export function extractStorageLayout(
   contractDef: ContractDefinition,
   decodeSrc: SrcDecoder,
   deref: ASTDereferencer,
+  storageLayout?: StorageLayout | undefined,
 ): StorageLayout {
-  const layout: StorageLayout = { storage: [], types: {}, layoutVersion: currentLayoutVersion };
-
-  // Note: A UserDefinedTypeName can also refer to a ContractDefinition but we won't care about those.
-  const derefUserDefinedType = deref(['StructDefinition', 'EnumDefinition']);
+  const layout: StorageLayout = { storage: [], types: {}, layoutVersion: currentLayoutVersion, flat: false };
+  if (storageLayout !== undefined) {
+    layout.types = mapValues(storageLayout.types, m => {
+      return {
+        label: m.label,
+        members: m.members?.map(m => (typeof m === 'string' ? m : pick(m, ['label', 'type']))) as TypeItem['members'],
+      };
+    });
 
-  for (const varDecl of contractDef.nodes) {
-    if (isNodeType('VariableDeclaration', varDecl)) {
-      if (!varDecl.constant && varDecl.mutability !== 'immutable') {
-        const type = normalizeTypeIdentifier(typeDescriptions(varDecl).typeIdentifier);
-
-        layout.storage.push({
-          contract: contractDef.name,
-          label: varDecl.name,
-          type,
-          src: decodeSrc(varDecl),
-        });
-
-        assert(varDecl.typeName != null);
-
-        // We will recursively look for all types involved in this variable declaration in order to store their type
-        // information. We iterate over a Map that is indexed by typeIdentifier to ensure we visit each type only once.
-        // Note that there can be recursive types.
-        const typeNames = new Map(
-          [...findTypeNames(varDecl.typeName)].map(n => [typeDescriptions(n).typeIdentifier, n]),
-        );
-
-        for (const typeName of typeNames.values()) {
-          const { typeIdentifier, typeString: label } = typeDescriptions(typeName);
-
-          const type = normalizeTypeIdentifier(typeIdentifier);
-
-          if (type in layout.types) {
-            continue;
-          }
-
-          let members;
-
-          if ('referencedDeclaration' in typeName && !/^t_contract\b/.test(type)) {
-            const typeDef = derefUserDefinedType(typeName.referencedDeclaration);
-            members = getTypeMembers(typeDef);
-            // Recursively look for the types referenced in this definition and add them to the queue.
-            for (const typeName of findTypeNames(typeDef)) {
-              const { typeIdentifier } = typeDescriptions(typeName);
-              if (!typeNames.has(typeIdentifier)) {
-                typeNames.set(typeIdentifier, typeName);
-              }
-            }
-          }
-
-          layout.types[type] = { label, members };
+    for (const storage of storageLayout.storage) {
+      const origin = getOriginContract(contractDef, storage.astId, deref);
+
+      if (origin) {
+        const [varDecl, contract] = origin;
+        // Solc layout doesn't bring members for enums so we get them using the ast method
+        loadLayoutType(varDecl, layout, deref);
+        const { label, offset, slot, type } = storage;
+        const src = decodeSrc(varDecl);
+        layout.storage.push({ label, offset, slot, type, contract, src });
+        layout.flat = true;
+      }
+    }
+  } else {
+    for (const varDecl of contractDef.nodes) {
+      if (isNodeType('VariableDeclaration', varDecl)) {
+        if (!varDecl.constant && varDecl.mutability !== 'immutable') {
+          const type = normalizeTypeIdentifier(typeDescriptions(varDecl).typeIdentifier);
+
+          layout.storage.push({
+            contract: contractDef.name,
+            label: varDecl.name,
+            type,
+            src: decodeSrc(varDecl),
+          });
+
+          loadLayoutType(varDecl, layout, deref);
         }
       }
     }
   }
-
   return layout;
 }
 
@@ -108,3 +101,56 @@ function getTypeMembers(typeDef: StructDefinition | EnumDefinition): TypeItem['m
     return typeDef.members.map(m => m.name);
   }
 }
+
+function getOriginContract(
+  contract: ContractDefinition,
+  astId: number | undefined,
+  deref: ASTDereferencer,
+): undefined | [VariableDeclaration, string] {
+  for (const id of contract.linearizedBaseContracts.reverse()) {
+    const parentContract = deref(['ContractDefinition'], id);
+
+    const varDecl = parentContract.nodes.find(n => n.id == astId);
+    if (varDecl && isNodeType('VariableDeclaration', varDecl)) {
+      return [varDecl, parentContract.name];
+    }
+  }
+}
+
+function loadLayoutType(varDecl: VariableDeclaration, layout: StorageLayout, deref: ASTDereferencer) {
+  // Note: A UserDefinedTypeName can also refer to a ContractDefinition but we won't care about those.
+  const derefUserDefinedType = deref(['StructDefinition', 'EnumDefinition']);
+
+  assert(varDecl.typeName != null);
+
+  // We will recursively look for all types involved in this variable declaration in order to store their type
+  // information. We iterate over a Map that is indexed by typeIdentifier to ensure we visit each type only once.
+  // Note that there can be recursive types.
+  const typeNames = new Map([...findTypeNames(varDecl.typeName)].map(n => [typeDescriptions(n).typeIdentifier, n]));
+
+  for (const typeName of typeNames.values()) {
+    const { typeIdentifier, typeString: label } = typeDescriptions(typeName);
+
+    const type = normalizeTypeIdentifier(typeIdentifier);
+    const hasMembers = /^t_(enum|struct)\b/.test(type);
+    if (type in layout.types && (!hasMembers || layout.types[type].members !== undefined)) {
+      continue;
+    }
+
+    let members;
+
+    if ('referencedDeclaration' in typeName && !/^t_contract\b/.test(type)) {
+      const typeDef = derefUserDefinedType(typeName.referencedDeclaration);
+      members = getTypeMembers(typeDef);
+      // Recursively look for the types referenced in this definition and add them to the queue.
+      for (const typeName of findTypeNames(typeDef)) {
+        const { typeIdentifier } = typeDescriptions(typeName);
+        if (!typeNames.has(typeIdentifier)) {
+          typeNames.set(typeIdentifier, typeName);
+        }
+      }
+    }
+
+    layout.types[type] = { label, members };
+  }
+}

package/src/storage/layout.ts

@@ -9,13 +9,17 @@ export interface StorageLayout {
   layoutVersion?: string;
   storage: StorageItem[];
   types: Record<string, TypeItem>;
+  flat?: boolean;
 }
 
 export interface StorageItem<Type = string> {
+  astId?: number;
   contract: string;
   label: string;
   type: Type;
   src: string;
+  offset?: number;
+  slot?: string;
 }
 
 export interface TypeItem<Type = string> {

package/src/storage/report.test.ts

@@ -1,4 +1,4 @@
-import _test, { TestInterface } from 'ava';
+import _test, { TestFn } from 'ava';
 import { ContractDefinition } from 'solidity-ast';
 import { findAll } from 'solidity-ast/utils';
 import { artifacts } from 'hardhat';
@@ -13,7 +13,7 @@ interface Context {
   extractStorageLayout: (contract: string) => ReturnType<typeof extractStorageLayout>;
 }
 
-const test = _test as TestInterface<Context>;
+const test = _test as TestFn<Context>;
 
 const dummyDecodeSrc = () => 'file.sol:1';
 

package/src/storage.test.ts

@@ -1,4 +1,4 @@
-import _test, { TestInterface } from 'ava';
+import _test, { TestFn } from 'ava';
 import { ContractDefinition } from 'solidity-ast';
 import { findAll } from 'solidity-ast/utils';
 import { artifacts } from 'hardhat';
@@ -14,7 +14,7 @@ interface Context {
   extractStorageLayout: (contract: string) => ReturnType<typeof extractStorageLayout>;
 }
 
-const test = _test as TestInterface<Context>;
+const test = _test as TestFn<Context>;
 
 test.before(async t => {
   const buildInfo = await artifacts.getBuildInfo('contracts/test/Storage.sol:Storage1');
@@ -344,7 +344,7 @@ test('storage upgrade with recursive type', t => {
   const v1 = t.context.extractStorageLayout('StorageUpgrade_Recursive_V1');
   const v2 = t.context.extractStorageLayout('StorageUpgrade_Recursive_V2');
   const e = t.throws(() => getStorageUpgradeErrors(v1, v2));
-  t.true(e.message.includes('Recursion found'));
+  t.true(e?.message.includes('Recursion found'));
 });
 
 test('storage upgrade with contract type', t => {

package/src/usage-error.ts

@@ -51,6 +51,15 @@ export class PrepareUpgradeUnsupportedError extends UpgradesError {
   }
 }
 
+export class ForceImportUnsupportedError extends UpgradesError {
+  constructor(proxyOrBeaconAddress: string) {
+    super(
+      `Contract at address ${proxyOrBeaconAddress} doesn't look like a supported proxy or beacon`,
+      () => `Only transparent, UUPS, or beacon proxies or beacons can be used with the forceImport() function.`,
+    );
+  }
+}
+
 export async function assertNotProxy(provider: EthereumProvider, address: string) {
   if (await isTransparentOrUUPSProxy(provider, address)) {
     throw new UpgradesError(

package/src/utils/parse-type-id.test.ts.md

@@ -4,40 +4,56 @@ The actual snapshot is saved in `parse-type-id.test.ts.snap`.
 
 Generated by [AVA](https://avajs.dev).
 
-## t_array(t_function_internal_nonpayable(t_uint256)returns(t_address))10_storage
+## t_uint256
+
+> Snapshot 1
+
+    {
+      args: undefined,
+      head: 't_uint256',
+      id: 't_uint256',
+      rets: undefined,
+      tail: undefined,
+    }
+
+## t_enum(MyEnum)
 
 > Snapshot 1
 
     {
       args: [
         {
-          args: [
-            {
-              args: undefined,
-              head: 't_uint256',
-              id: 't_uint256',
-              rets: undefined,
-              tail: undefined,
-            },
-          ],
-          head: 't_function_internal_nonpayable',
-          id: 't_function_internal_nonpayable(t_uint256)returns(t_address)',
-          rets: [
-            {
-              args: undefined,
-              head: 't_address',
-              id: 't_address',
-              rets: undefined,
-              tail: undefined,
-            },
-          ],
+          args: undefined,
+          head: 'MyEnum',
+          id: 'MyEnum',
+          rets: undefined,
           tail: undefined,
         },
       ],
-      head: 't_array',
-      id: 't_array(t_function_internal_nonpayable(t_uint256)returns(t_address))10_storage',
+      head: 't_enum',
+      id: 't_enum(MyEnum)',
       rets: undefined,
-      tail: '10_storage',
+      tail: undefined,
+    }
+
+## t_struct(MyComplexStruct)storage
+
+> Snapshot 1
+
+    {
+      args: [
+        {
+          args: undefined,
+          head: 'MyComplexStruct',
+          id: 'MyComplexStruct',
+          rets: undefined,
+          tail: undefined,
+        },
+      ],
+      head: 't_struct',
+      id: 't_struct(MyComplexStruct)storage',
+      rets: undefined,
+      tail: 'storage',
     }
 
 ## t_array(t_uint256)3_storage
@@ -60,7 +76,7 @@ Generated by [AVA](https://avajs.dev).
       tail: '3_storage',
     }
 
-## t_function_internal_nonpayable(t_uint256)returns()
+## t_mapping(t_uint256,t_uint256)
 
 > Snapshot 1
 
@@ -73,10 +89,44 @@ Generated by [AVA](https://avajs.dev).
           rets: undefined,
           tail: undefined,
         },
+        {
+          args: undefined,
+          head: 't_uint256',
+          id: 't_uint256',
+          rets: undefined,
+          tail: undefined,
+        },
       ],
-      head: 't_function_internal_nonpayable',
-      id: 't_function_internal_nonpayable(t_uint256)returns()',
-      rets: [],
+      head: 't_mapping',
+      id: 't_mapping(t_uint256,t_uint256)',
+      rets: undefined,
+      tail: undefined,
+    }
+
+## t_mapping(unknown,t_uint256)
+
+> Snapshot 1
+
+    {
+      args: [
+        {
+          args: undefined,
+          head: 'unknown',
+          id: 'unknown',
+          rets: undefined,
+          tail: undefined,
+        },
+        {
+          args: undefined,
+          head: 't_uint256',
+          id: 't_uint256',
+          rets: undefined,
+          tail: undefined,
+        },
+      ],
+      head: 't_mapping',
+      id: 't_mapping(unknown,t_uint256)',
+      rets: undefined,
       tail: undefined,
     }
 
@@ -157,58 +207,12 @@ Generated by [AVA](https://avajs.dev).
       tail: undefined,
     }
 
-## t_mapping(t_uint256,t_uint256)
-
-> Snapshot 1
-
-    {
-      args: [
-        {
-          args: undefined,
-          head: 't_uint256',
-          id: 't_uint256',
-          rets: undefined,
-          tail: undefined,
-        },
-        {
-          args: undefined,
-          head: 't_uint256',
-          id: 't_uint256',
-          rets: undefined,
-          tail: undefined,
-        },
-      ],
-      head: 't_mapping',
-      id: 't_mapping(t_uint256,t_uint256)',
-      rets: undefined,
-      tail: undefined,
-    }
-
-## t_uint256
-
-> Snapshot 1
-
-    {
-      args: undefined,
-      head: 't_uint256',
-      id: 't_uint256',
-      rets: undefined,
-      tail: undefined,
-    }
-
-## t_mapping(unknown,t_uint256)
+## t_function_internal_nonpayable(t_uint256)returns()
 
 > Snapshot 1
 
     {
       args: [
-        {
-          args: undefined,
-          head: 'unknown',
-          id: 'unknown',
-          rets: undefined,
-          tail: undefined,
-        },
         {
           args: undefined,
           head: 't_uint256',
@@ -217,48 +221,44 @@ Generated by [AVA](https://avajs.dev).
           tail: undefined,
         },
       ],
-      head: 't_mapping',
-      id: 't_mapping(unknown,t_uint256)',
-      rets: undefined,
-      tail: undefined,
-    }
-
-## t_enum(MyEnum)
-
-> Snapshot 1
-
-    {
-      args: [
-        {
-          args: undefined,
-          head: 'MyEnum',
-          id: 'MyEnum',
-          rets: undefined,
-          tail: undefined,
-        },
-      ],
-      head: 't_enum',
-      id: 't_enum(MyEnum)',
-      rets: undefined,
+      head: 't_function_internal_nonpayable',
+      id: 't_function_internal_nonpayable(t_uint256)returns()',
+      rets: [],
       tail: undefined,
     }
 
-## t_struct(MyComplexStruct)storage
+## t_array(t_function_internal_nonpayable(t_uint256)returns(t_address))10_storage
 
 > Snapshot 1
 
     {
       args: [
         {
-          args: undefined,
-          head: 'MyComplexStruct',
-          id: 'MyComplexStruct',
-          rets: undefined,
+          args: [
+            {
+              args: undefined,
+              head: 't_uint256',
+              id: 't_uint256',
+              rets: undefined,
+              tail: undefined,
+            },
+          ],
+          head: 't_function_internal_nonpayable',
+          id: 't_function_internal_nonpayable(t_uint256)returns(t_address)',
+          rets: [
+            {
+              args: undefined,
+              head: 't_address',
+              id: 't_address',
+              rets: undefined,
+              tail: undefined,
+            },
+          ],
           tail: undefined,
         },
       ],
-      head: 't_struct',
-      id: 't_struct(MyComplexStruct)storage',
+      head: 't_array',
+      id: 't_array(t_function_internal_nonpayable(t_uint256)returns(t_address))10_storage',
       rets: undefined,
-      tail: 'storage',
+      tail: '10_storage',
     }

package/src/validate/query.ts

@@ -58,13 +58,17 @@ export function getStorageLayout(data: ValidationData, version: Version): Storag
 }
 
 export function unfoldStorageLayout(runData: ValidationRunData, contractName: string): StorageLayout {
-  const c = runData[contractName];
   const layout: StorageLayout = { storage: [], types: {} };
-  for (const name of [contractName].concat(c.inherit)) {
-    layout.storage.unshift(...runData[name].layout.storage);
-    Object.assign(layout.types, runData[name].layout.types);
+  if (runData[contractName].layout.flat) {
+    return { storage: runData[contractName].layout.storage, types: runData[contractName].layout.types };
+  } else {
+    const c = runData[contractName];
+    for (const name of [contractName].concat(c.inherit)) {
+      layout.storage.unshift(...runData[name].layout.storage);
+      Object.assign(layout.types, runData[name].layout.types);
+    }
+    return layout;
   }
-  return layout;
 }
 
 export function* findVersionWithoutMetadataMatches(

package/src/validate/run.ts

@@ -162,7 +162,12 @@ export function validate(solcOutput: SolcOutput, decodeSrc: SrcDecoder): Validat
           ...getLinkingErrors(contractDef, bytecode),
         ];
 
-        validation[contractDef.name].layout = extractStorageLayout(contractDef, decodeSrc, deref);
+        validation[contractDef.name].layout = extractStorageLayout(
+          contractDef,
+          decodeSrc,
+          deref,
+          solcOutput.contracts[source][contractDef.name].storageLayout,
+        );
         validation[contractDef.name].methods = [...findAll('FunctionDefinition', contractDef)]
           .filter(fnDef => ['external', 'public'].includes(fnDef.visibility))
           .map(fnDef => getFunctionSignature(fnDef, deref));

package/src/validate.test.ts

@@ -1,4 +1,4 @@
-import _test, { TestInterface } from 'ava';
+import _test, { TestFn } from 'ava';
 import { artifacts } from 'hardhat';
 
 import {
@@ -15,7 +15,7 @@ interface Context {
   validation: RunValidation;
 }
 
-const test = _test as TestInterface<Context>;
+const test = _test as TestFn<Context>;
 
 test.before(async t => {
   const contracts = [