Fix contract size and add custom override

contracts/TransformCustomSize.sol

@@ -0,0 +1,62 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8;
+
+struct OneAndAHalfSlot {
+  uint256 x;
+  uint128 y;
+}
+
+contract SizeDefault {
+  /// @custom:oz-upgrades-unsafe-allow state-variable-immutable state-variable-assignment
+  uint immutable w1 = block.number;
+  /// @custom:oz-upgrades-unsafe-allow state-variable-immutable state-variable-assignment
+  uint immutable w2 = block.timestamp;
+  uint immutable x; // slot 0 (after conversion to private)
+  uint constant y = 1;
+  uint224 z0; // slot 1
+  uint256 z1; // slot 2
+  uint32 z2; // slot 3
+  OneAndAHalfSlot s1; // slot 4&5
+  OneAndAHalfSlot s2; // slot 6&7
+  uint32 z3; // slot 8
+  uint32 z4; // slot 8
+  uint32 z5; // slot 8
+  uint64[5] a1; // slot 9&10
+  uint64[3] a2; // slot 11
+
+  constructor(uint _x) {
+    x = _x;
+  }
+  // gap should be 38 = 50 - 12
+}
+
+/// @custom:storage-size 128
+contract SizeOverride {
+  /// @custom:oz-upgrades-unsafe-allow state-variable-immutable state-variable-assignment
+  uint immutable w1 = block.number;
+  /// @custom:oz-upgrades-unsafe-allow state-variable-immutable  state-variable-assignment
+  uint immutable w2 = block.timestamp;
+  uint immutable x; // slot 0 (after conversion to private)
+  uint constant y = 1;
+  uint224 z0; // slot 1
+  uint256 z1; // slot 2
+  uint32 z2; // slot 3
+  OneAndAHalfSlot s1; // slot 4&5
+  OneAndAHalfSlot s2; // slot 6&7
+  uint32 z3; // slot 8
+  uint32 z4; // slot 8
+  uint32 z5; // slot 8
+  uint64[5] a1; // slot 9&10
+  uint64[3] a2; // slot 11
+
+  constructor(uint _x) {
+    x = _x;
+  }
+  // gap should be 116 = 128 - 12
+}
+
+/// @custom:storage-size 2
+contract SizeOverrideExact {
+  uint immutable s1 = block.number;
+  bool s2;
+}

src/transform-0.8.test.ts

@@ -13,6 +13,7 @@ import {
   transformConstructor,
 } from './transformations/transform-constructor';
 import { renameInheritdoc } from './transformations/rename-inheritdoc';
+import { addStorageGaps } from './transformations/add-storage-gaps';
 
 const test = _test as TestFn<Context>;
 
@@ -54,3 +55,13 @@ test('preserves immutable if allowed', t => {
   t.context.transform.apply(removeImmutable);
   t.snapshot(t.context.transform.results()[file]);
 });
+
+test('custom contract size', t => {
+  const file = 'contracts/TransformCustomSize.sol';
+  t.context.transform.apply(transformConstructor);
+  t.context.transform.apply(removeLeftoverConstructorHead);
+  t.context.transform.apply(removeStateVarInits);
+  t.context.transform.apply(removeImmutable);
+  t.context.transform.apply(addStorageGaps);
+  t.snapshot(t.context.transform.results()[file]);
+});

src/transform-0.8.test.ts.md

@@ -102,3 +102,100 @@ Generated by [AVA](https://avajs.dev).
         }␊
     }␊
     `
+
+## custom contract size
+
+> Snapshot 1
+
+    `// SPDX-License-Identifier: UNLICENSED␊
+    pragma solidity ^0.8;␊
+    ␊
+    struct OneAndAHalfSlot {␊
+      uint256 x;␊
+      uint128 y;␊
+    }␊
+    ␊
+    contract SizeDefault {␊
+      /// @custom:oz-upgrades-unsafe-allow state-variable-immutable state-variable-assignment␊
+      uint immutable w1 = block.number;␊
+      /// @custom:oz-upgrades-unsafe-allow state-variable-immutable state-variable-assignment␊
+      uint immutable w2 = block.timestamp;␊
+      uint x; // slot 0 (after conversion to private)␊
+      uint constant y = 1;␊
+      uint224 z0; // slot 1␊
+      uint256 z1; // slot 2␊
+      uint32 z2; // slot 3␊
+      OneAndAHalfSlot s1; // slot 4&5␊
+      OneAndAHalfSlot s2; // slot 6&7␊
+      uint32 z3; // slot 8␊
+      uint32 z4; // slot 8␊
+      uint32 z5; // slot 8␊
+      uint64[5] a1; // slot 9&10␊
+      uint64[3] a2; // slot 11␊
+    ␊
+      function __SizeDefault_init(uint _x) internal onlyInitializing {␊
+            __SizeDefault_init_unchained(_x);␊
+        }␊
+    ␊
+        function __SizeDefault_init_unchained(uint _x) internal onlyInitializing {␊
+        x = _x;␊
+      }␊
+      // gap should be 38 = 50 - 12␊
+    ␊
+        /**␊
+         * @dev This empty reserved space is put in place to allow future versions to add new␊
+         * variables without shifting down storage in the inheritance chain.␊
+         * See https://docs.openzeppelin.com/contracts/4.x/upgradeable#storage_gaps␊
+         */␊
+        uint256[38] private __gap;␊
+    }␊
+    ␊
+    /// @custom:storage-size 128␊
+    contract SizeOverride {␊
+      /// @custom:oz-upgrades-unsafe-allow state-variable-immutable state-variable-assignment␊
+      uint immutable w1 = block.number;␊
+      /// @custom:oz-upgrades-unsafe-allow state-variable-immutable  state-variable-assignment␊
+      uint immutable w2 = block.timestamp;␊
+      uint x; // slot 0 (after conversion to private)␊
+      uint constant y = 1;␊
+      uint224 z0; // slot 1␊
+      uint256 z1; // slot 2␊
+      uint32 z2; // slot 3␊
+      OneAndAHalfSlot s1; // slot 4&5␊
+      OneAndAHalfSlot s2; // slot 6&7␊
+      uint32 z3; // slot 8␊
+      uint32 z4; // slot 8␊
+      uint32 z5; // slot 8␊
+      uint64[5] a1; // slot 9&10␊
+      uint64[3] a2; // slot 11␊
+    ␊
+      function __SizeOverride_init(uint _x) internal onlyInitializing {␊
+            __SizeOverride_init_unchained(_x);␊
+        }␊
+    ␊
+        function __SizeOverride_init_unchained(uint _x) internal onlyInitializing {␊
+        x = _x;␊
+      }␊
+      // gap should be 116 = 128 - 12␊
+    ␊
+        /**␊
+         * @dev This empty reserved space is put in place to allow future versions to add new␊
+         * variables without shifting down storage in the inheritance chain.␊
+         * See https://docs.openzeppelin.com/contracts/4.x/upgradeable#storage_gaps␊
+         */␊
+        uint256[116] private __gap;␊
+    }␊
+    ␊
+    /// @custom:storage-size 2␊
+    contract SizeOverrideExact {␊
+        function __SizeOverrideExact_init() internal onlyInitializing {␊
+            __SizeOverrideExact_init_unchained();␊
+        }␊
+    ␊
+        function __SizeOverrideExact_init_unchained() internal onlyInitializing {␊
+            s1 = block.number;␊
+        }␊
+      uint s1;␊
+      bool s2;␊
+    }␊
+    `

src/transformations/add-storage-gaps.ts

@@ -1,66 +1,115 @@
 import { SourceUnit, ContractDefinition } from 'solidity-ast';
-import { findAll } from 'solidity-ast/utils';
+import { findAll, isNodeType } from 'solidity-ast/utils';
 
 import { formatLines } from './utils/format-lines';
+import { hasOverride } from '../utils/upgrades-overrides';
 import { getNodeBounds } from '../solc/ast-utils';
 import { StorageLayout } from '../solc/input-output';
 import { Transformation } from './type';
 import { TransformerTools } from '../transform';
+import { extractNatspec } from '../utils/extractNatspec';
+import { decodeTypeIdentifier } from '../utils/type-id';
 
-// 100 slots of 32 contractSize each
-const TARGET_SIZE = 32 * 50;
+// By default, make the contract a total of 50 slots (storage + gap)
+const DEFAULT_SLOT_COUNT = 50;
 
 export function* addStorageGaps(
   sourceUnit: SourceUnit,
   { getLayout }: TransformerTools,
 ): Generator<Transformation> {
   for (const contract of findAll('ContractDefinition', sourceUnit)) {
     if (contract.contractKind === 'contract') {
-      const gapSize = getGapSize(contract, getLayout(contract));
+      let targetSlots = DEFAULT_SLOT_COUNT;
+      for (const entry of extractNatspec(contract)) {
+        if (entry.title === 'custom' && entry.tag === 'storage-size') {
+          targetSlots = parseInt(entry.args);
+        }
+      }
 
-      const contractBounds = getNodeBounds(contract);
-      const start = contractBounds.start + contractBounds.length - 1;
+      const gapSize = targetSlots - getContractSlotCount(contract, getLayout(contract));
 
-      const text = formatLines(0, [
-        ``,
-        [
-          `/**`,
-          ` * @dev This empty reserved space is put in place to allow future versions to add new`,
-          ` * variables without shifting down storage in the inheritance chain.`,
-          ` * See https://docs.openzeppelin.com/contracts/4.x/upgradeable#storage_gaps`,
-          ` */`,
-          `uint256[${gapSize}] private __gap;`,
-        ],
-      ]);
+      if (gapSize > 0) {
+        const contractBounds = getNodeBounds(contract);
+        const start = contractBounds.start + contractBounds.length - 1;
 
-      yield {
-        kind: 'add-storage-gaps',
-        start,
-        length: 0,
-        text,
-      };
+        const text = formatLines(0, [
+          ``,
+          [
+            `/**`,
+            ` * @dev This empty reserved space is put in place to allow future versions to add new`,
+            ` * variables without shifting down storage in the inheritance chain.`,
+            ` * See https://docs.openzeppelin.com/contracts/4.x/upgradeable#storage_gaps`,
+            ` */`,
+            `uint256[${gapSize}] private __gap;`,
+          ],
+        ]);
+
+        yield {
+          kind: 'add-storage-gaps',
+          start,
+          length: 0,
+          text,
+        };
+      } else if (gapSize < 0) {
+        throw new Error(
+          `Contract ${contract.name} uses more then the ${targetSlots} reserved slots.`,
+        );
+      }
     }
   }
 }
 
-function getGapSize(contractNode: ContractDefinition, layout: StorageLayout): number {
-  const varIds = new Set([...findAll('VariableDeclaration', contractNode)].map(v => v.id));
-
-  if (layout === undefined) {
-    throw new Error('Storage layout is needed for this transformation');
+function getNumberOfBytesOfValueType(type: string) {
+  const details = type.match(/^t_(?<base>[a-z]+)(?<size>[\d]+)?$/);
+  switch (details?.groups?.base) {
+    case 'bool':
+    case 'byte':
+      return 1;
+    case 'address':
+      return 20;
+    case 'bytes':
+      return parseInt(details.groups.size, 10);
+    case 'int':
+    case 'uint':
+      return parseInt(details.groups.size, 10) / 8;
+    default:
+      throw new Error(`Unsupported value type: ${type}`);
   }
+}
 
-  const local = layout.storage.filter(l => varIds.has(l.astId));
-
-  let contractSize = 0;
+function getContractSlotCount(contractNode: ContractDefinition, layout: StorageLayout): number {
+  // This tracks both slot and offset:
+  // - slot   = Math.floor(contractSizeInBytes / 32)
+  // - offset = contractSizeInBytes % 32
+  let contractSizeInBytes = 0;
 
-  for (const l of local) {
-    const type = layout.types?.[l.type];
-    if (type === undefined) {
-      throw new Error(`Missing type information for ${type}`);
+  // don't use `findAll` here, we don't want to go recursive
+  for (const varDecl of contractNode.nodes.filter(isNodeType('VariableDeclaration'))) {
+    if (varDecl.mutability === 'constant') {
+      continue;
     }
-    contractSize += parseInt(type.numberOfBytes, 10);
+    if (varDecl.mutability === 'immutable' && hasOverride(varDecl, 'state-variable-immutable')) {
+      continue;
+    }
+
+    // try get type details
+    const typeIdentifier = decodeTypeIdentifier(varDecl.typeDescriptions.typeIdentifier ?? '');
+    const type = layout.types?.[typeIdentifier];
+
+    // size of current object from type details, or try to reconstruct it if
+    // they're not available try to reconstruct it, which can happen for
+    // immutable variables
+    const size = type
+      ? parseInt(type.numberOfBytes, 10)
+      : getNumberOfBytesOfValueType(typeIdentifier);
+
+    // used space in the current slot
+    const offset = contractSizeInBytes % 32;
+    // free space in the current slot (only if slot is dirty)
+    const free = (32 - offset) % 32;
+    // if the free space is not enough to fit the current object, then consume the free space to start at next slot
+    contractSizeInBytes += (size > free ? free : 0) + size;
   }
 
-  return Math.floor((TARGET_SIZE - contractSize) / 32);
+  return Math.ceil(contractSizeInBytes / 32);
 }

src/utils/extractNatspec.ts

@@ -0,0 +1,28 @@
+import { StructuredDocumentation } from 'solidity-ast';
+import { execall } from './execall';
+
+interface NatspecTag {
+  title: string;
+  tag: string;
+  args: string;
+}
+
+export function* extractNatspec(node: {
+  documentation?: string | StructuredDocumentation | null;
+}): Generator<NatspecTag> {
+  const doc =
+    typeof node.documentation === 'string' ? node.documentation : node.documentation?.text ?? '';
+
+  for (const { groups } of execall(
+    /^\s*(?:@(?<title>\w+)(?::(?<tag>[a-z][a-z-]*))? )?(?<args>(?:(?!^\s@\w+)[^])*)/m,
+    doc,
+  )) {
+    if (groups) {
+      yield {
+        title: groups.title ?? '',
+        tag: groups.tag ?? '',
+        args: groups.args ?? '',
+      };
+    }
+  }
+}