Improve detection of storage gap usage (#731)

Author: ericglau

packages/core/CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Unreleased
 
 - Support storage gaps named with `__gap_*`. ([#732](https://github.com/OpenZeppelin/openzeppelin-upgrades/pull/732))
+- Improve detection of storage gap usage. ([#731](https://github.com/OpenZeppelin/openzeppelin-upgrades/pull/731))
 
 ## 1.22.0 (2023-01-31)
 

packages/core/contracts/test/Storage.sol

@@ -858,4 +858,45 @@ contract StorageUpgrade_CustomGap_V2_Bad_Changed_Gap_Name {
     uint256 c;
     uint256[23] __gap_part2; // name stayed the same
     uint256 z;
+}
+
+contract Parent1 {
+    string[5] s;
+    uint256[45] private __gap;
+}
+
+contract Parent2 {
+    bool b;
+    uint256[49] private __gap;
+}
+
+contract StorageUpgrade_ConsumeAndAddGap_V1 is Parent1 {
+    uint256[50] private __gap;
+}
+
+contract StorageUpgrade_ConsumeAndAddGap_V2 is Parent1, Parent2 {
+}
+
+contract StorageUpgrade_ConsumeAndAddGap_V3 is Parent1, Parent2 {
+    uint256 a;
+    uint256[49] private __gap;
+}
+
+contract StorageUpgrade_ConsumeAndAddGap_V3b is Parent1, Parent2 {
+    uint256[50] private __gap;
+}
+
+contract V1Storage {
+    uint256[50] private __gap;
+}
+
+contract StorageUpgrade_ConsumeAndAddGap_Storage_V1 is V1Storage, Parent1 {
+}
+
+contract V2Storage {
+    uint256 abc;
+    uint256[49] private __gap;
+}
+
+contract StorageUpgrade_ConsumeAndAddGap_Storage_V2 is V2Storage, Parent1, Parent2 {
 }

packages/core/src/storage/compare.ts

@@ -91,10 +91,10 @@ export function storageFieldEnd(field: StorageField): bigint | undefined {
   return begin + BigInt(numberOfBytes);
 }
 
-const LAYOUTCHANGE_COST = 0;
+const LAYOUTCHANGE_COST = 1;
 const FINISHGAP_COST = 1;
 const SHRINKGAP_COST = 0;
-const TYPECHANGE_COST = 0;
+const TYPECHANGE_COST = 1;
 
 export class StorageLayoutComparator {
   hasAllowedUncheckedCustomTypes = false;

packages/core/src/storage/report-gap.test.ts

@@ -31,6 +31,12 @@ const testContracts = [
   'contracts/test/Storage.sol:StorageUpgrade_CustomGap_V2_Ok_Switched_Gaps',
   'contracts/test/Storage.sol:StorageUpgrade_CustomGap_V2_Bad_Switched_Gaps',
   'contracts/test/Storage.sol:StorageUpgrade_CustomGap_V2_Bad_Changed_Gap_Name',
+  'contracts/test/Storage.sol:StorageUpgrade_ConsumeAndAddGap_V1',
+  'contracts/test/Storage.sol:StorageUpgrade_ConsumeAndAddGap_V2',
+  'contracts/test/Storage.sol:StorageUpgrade_ConsumeAndAddGap_V3',
+  'contracts/test/Storage.sol:StorageUpgrade_ConsumeAndAddGap_V3b',
+  'contracts/test/Storage.sol:StorageUpgrade_ConsumeAndAddGap_Storage_V1',
+  'contracts/test/Storage.sol:StorageUpgrade_ConsumeAndAddGap_Storage_V2',
 ];
 
 test.before(async t => {
@@ -172,3 +178,29 @@ test('custom gap - insert var, did not shrink gaps, changed first gap name', t =
   t.false(report.ok);
   t.snapshot(report.explain());
 });
+
+test('consume entire gap and add new gap of different size', t => {
+  const v1 = t.context.extractStorageLayout('StorageUpgrade_ConsumeAndAddGap_V1');
+  const v2 = t.context.extractStorageLayout('StorageUpgrade_ConsumeAndAddGap_V2');
+  const v3 = t.context.extractStorageLayout('StorageUpgrade_ConsumeAndAddGap_V3');
+
+  t.true(getReport(v1, v2).ok);
+  t.true(getReport(v2, v3).ok);
+  t.true(getReport(v1, v3).ok);
+});
+
+test('consume entire gap and add new gap of same size', t => {
+  const v1 = t.context.extractStorageLayout('StorageUpgrade_ConsumeAndAddGap_V1');
+  const v2 = t.context.extractStorageLayout('StorageUpgrade_ConsumeAndAddGap_V2');
+  const v3b = t.context.extractStorageLayout('StorageUpgrade_ConsumeAndAddGap_V3b');
+
+  t.true(getReport(v2, v3b).ok);
+  t.true(getReport(v1, v3b).ok);
+});
+
+test('consume partial gap and add new gap, storage contract pattern', t => {
+  const v1 = t.context.extractStorageLayout('StorageUpgrade_ConsumeAndAddGap_Storage_V1');
+  const v2 = t.context.extractStorageLayout('StorageUpgrade_ConsumeAndAddGap_Storage_V2');
+
+  t.true(getReport(v1, v2).ok);
+});