Automatic checking of storage compatibility of structs and enums

packages/core/contracts/test/Storage.sol

@@ -128,3 +128,56 @@ contract StorageUpgrade_Delete_V1 {
 contract StorageUpgrade_Delete_V2 {
   address x2;
 }
+
+contract StorageUpgrade_Struct_V1 {
+    struct Struct {
+        uint x;
+        string s1;
+        string s2;
+    }
+    Struct data;
+}
+
+contract StorageUpgrade_Struct_V2_Ok {
+    struct StructRenamed {
+        uint x;
+        string s1;
+        string s2;
+    }
+    StructRenamed data;
+}
+
+contract StorageUpgrade_Struct_V2_Bad {
+    struct StructRenamed {
+        uint x;
+        string s2;
+    }
+    StructRenamed data;
+}
+
+contract StorageUpgrade_Enum_V1 {
+    enum Enum {
+        A,
+        B
+    }
+    Enum data;
+}
+
+contract StorageUpgrade_Enum_V2_Ok {
+    enum Enum {
+        A,
+        B,
+        C
+    }
+    Enum data;
+}
+
+contract StorageUpgrade_Enum_V2_Bad {
+    enum Enum {
+        A,
+        C,
+        B
+    }
+    Enum data;
+}
+

packages/core/package.json

@@ -20,7 +20,7 @@
     "prepare:contracts": "hardhat compile",
     "prepublishOnly": "bash scripts/copy-artifacts.sh",
     "test": "tsc -b && hardhat compile --force && ava",
-    "test:watch": "hardhat compile --force && fgbg 'ava --watch' 'tsc -b --watch'",
+    "test:watch": "hardhat compile --force && fgbg 'ava --watch' 'tsc -b --watch' --",
     "version": "node ../../scripts/bump-changelog.js"
   },
   "devDependencies": {
@@ -51,6 +51,6 @@
     "fp-ts": "^2.7.1",
     "io-ts": "^2.2.9",
     "proper-lockfile": "^4.1.1",
-    "solidity-ast": "^0.4.4"
+    "solidity-ast": "^0.4.15"
   }
 }

packages/core/src/ast-dereferencer.ts

@@ -0,0 +1,42 @@
+import { findAll } from 'solidity-ast/utils';
+import type { Node, NodeType, NodeTypeMap } from 'solidity-ast/node';
+
+import { curry2 } from './utils/curry';
+import type { SolcOutput } from './solc-api';
+
+// An ASTDereferencer is a function that looks up an AST node given its id, in all of the source files involved in a
+// solc run. It will generally be used together with the AST property `referencedDeclaration` (found in Identifier,
+// UserDefinedTypeName, etc.) to look up a variable definition or type definition.
+
+export interface ASTDereferencer {
+  <T extends NodeType>(nodeTypes: T[]): (id: number) => NodeTypeMap[T];
+  <T extends NodeType>(nodeTypes: T[], id: number): NodeTypeMap[T];
+}
+
+export function astDereferencer(solcOutput: SolcOutput): ASTDereferencer {
+  const asts = Array.from(Object.values(solcOutput.sources), s => s.ast);
+  const cache = new Map<number, Node>();
+
+  function deref<T extends NodeType>(nodeTypes: T[], id: number): NodeTypeMap[T] {
+    const cached = cache.get(id);
+
+    if (cached) {
+      if ((nodeTypes as NodeType[]).includes(cached.nodeType)) {
+        return cached as NodeTypeMap[T];
+      }
+    }
+
+    for (const ast of asts) {
+      for (const node of findAll(nodeTypes, ast)) {
+        if (node.id === id) {
+          cache.set(id, node);
+          return node;
+        }
+      }
+    }
+
+    throw new Error(`No node with id ${id} of type ${nodeTypes}`);
+  }
+
+  return curry2(deref);
+}

packages/core/src/storage.test.ts

@@ -4,10 +4,11 @@ import { findAll } from 'solidity-ast/utils';
 import { artifacts } from 'hardhat';
 
 import { SolcOutput } from './solc-api';
+import { astDereferencer } from './ast-dereferencer';
 import { extractStorageLayout, getStorageUpgradeErrors, stabilizeTypeIdentifier, StorageLayout } from './storage';
 
 interface Context {
-  contracts: Record<string, ContractDefinition>;
+  extractStorageLayout: (contract: string) => ReturnType<typeof extractStorageLayout>;
 }
 
 const test = _test as TestInterface<Context>;
@@ -18,57 +19,91 @@ test.before(async t => {
     throw new Error('Build info not found');
   }
   const solcOutput: SolcOutput = buildInfo.output;
-  t.context.contracts = {};
+  const contracts: Record<string, ContractDefinition> = {};
   for (const def of findAll('ContractDefinition', solcOutput.sources['contracts/test/Storage.sol'].ast)) {
-    t.context.contracts[def.name] = def;
+    contracts[def.name] = def;
   }
+  const deref = astDereferencer(solcOutput);
+  t.context.extractStorageLayout = name => extractStorageLayout(contracts[name], dummyDecodeSrc, deref);
 });
 
 const dummyDecodeSrc = () => 'file.sol:1';
 
 test('Storage1', t => {
-  const contract = 'Storage1';
-  const def = t.context.contracts[contract];
-  const layout = extractStorageLayout(def, dummyDecodeSrc);
+  const layout = t.context.extractStorageLayout('Storage1');
   t.snapshot(stabilizeStorageLayout(layout));
 });
 
 test('Storage2', t => {
-  const contract = 'Storage2';
-  const def = t.context.contracts[contract];
-  const layout = extractStorageLayout(def, dummyDecodeSrc);
+  const layout = t.context.extractStorageLayout('Storage2');
   t.snapshot(stabilizeStorageLayout(layout));
 });
 
 test('storage upgrade equal', t => {
-  const v1 = extractStorageLayout(t.context.contracts['StorageUpgrade_Equal_V1'], dummyDecodeSrc);
-  const v2 = extractStorageLayout(t.context.contracts['StorageUpgrade_Equal_V2'], dummyDecodeSrc);
+  const v1 = t.context.extractStorageLayout('StorageUpgrade_Equal_V1');
+  const v2 = t.context.extractStorageLayout('StorageUpgrade_Equal_V2');
   const comparison = getStorageUpgradeErrors(v1, v2);
   t.deepEqual(comparison, []);
 });
 
 test('storage upgrade append', t => {
-  const v1 = extractStorageLayout(t.context.contracts['StorageUpgrade_Append_V1'], dummyDecodeSrc);
-  const v2 = extractStorageLayout(t.context.contracts['StorageUpgrade_Append_V2'], dummyDecodeSrc);
+  const v1 = t.context.extractStorageLayout('StorageUpgrade_Append_V1');
+  const v2 = t.context.extractStorageLayout('StorageUpgrade_Append_V2');
   const comparison = getStorageUpgradeErrors(v1, v2);
   t.deepEqual(comparison, []);
 });
 
 test('storage upgrade delete', t => {
-  const v1 = extractStorageLayout(t.context.contracts['StorageUpgrade_Delete_V1'], dummyDecodeSrc);
-  const v2 = extractStorageLayout(t.context.contracts['StorageUpgrade_Delete_V2'], dummyDecodeSrc);
+  const v1 = t.context.extractStorageLayout('StorageUpgrade_Delete_V1');
+  const v2 = t.context.extractStorageLayout('StorageUpgrade_Delete_V2');
   const comparison = getStorageUpgradeErrors(v1, v2);
-  t.deepEqual(comparison, [
-    {
+  t.like(comparison, {
+    length: 1,
+    0: {
       kind: 'delete',
       original: {
         contract: 'StorageUpgrade_Delete_V1',
         label: 'x1',
-        type: 't_uint256',
-        src: 'file.sol:1',
+        type: {
+          id: 't_uint256',
+        },
       },
     },
-  ]);
+  });
+});
+
+test('storage upgrade with structs', t => {
+  const v1 = t.context.extractStorageLayout('StorageUpgrade_Struct_V1');
+
+  const v2_Ok = t.context.extractStorageLayout('StorageUpgrade_Struct_V2_Ok');
+  t.deepEqual(getStorageUpgradeErrors(v1, v2_Ok), []);
+
+  const v2_Bad = t.context.extractStorageLayout('StorageUpgrade_Struct_V2_Bad');
+  t.like(getStorageUpgradeErrors(v1, v2_Bad), {
+    length: 1,
+    0: {
+      kind: 'typechange',
+      original: { label: 'data' },
+      updated: { label: 'data' },
+    },
+  });
+});
+
+test('storage upgrade with enums', t => {
+  const v1 = t.context.extractStorageLayout('StorageUpgrade_Enum_V1');
+
+  const v2_Ok = t.context.extractStorageLayout('StorageUpgrade_Enum_V2_Ok');
+  t.deepEqual(getStorageUpgradeErrors(v1, v2_Ok), []);
+
+  const v2_Bad = t.context.extractStorageLayout('StorageUpgrade_Enum_V2_Bad');
+  t.like(getStorageUpgradeErrors(v1, v2_Bad), {
+    length: 1,
+    0: {
+      kind: 'typechange',
+      original: { label: 'data' },
+      updated: { label: 'data' },
+    },
+  });
 });
 
 function stabilizeStorageLayout(layout: StorageLayout) {