add error with namespaces + @Custom:storage-size

Author: frangio

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## 0.3.27 (2023-08-29)
+
+- Throw error when using `@custom:storage-size` along with namespaced storage.
+
 ## 0.3.26 (2023-08-28)
 
 - Change location of initializer functions when the original contract doesn't have a constructor. Previously it would be the start of the contract, before state variables. It is now placed immediately before the first function of the contract, if the contract has functions.

contracts/namespaces-error-storage-size.sol

@@ -0,0 +1,6 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.20;
+
+/// @custom:storage-size 52
+contract C1 {
+}

src/transform-namespaces.test.ts

@@ -17,7 +17,7 @@ const test = _test as TestFn<Context>;
 interface Context {
   solcInput: SolcInput;
   solcOutput: SolcOutput;
-  transform: Transform;
+  transformFile: (file: string) => Transform;
 }
 
 test.serial.before('compile', async t => {
@@ -28,14 +28,27 @@ test.serial.before('compile', async t => {
 });
 
 test.beforeEach('transform', async t => {
-  t.context.transform = new Transform(t.context.solcInput, t.context.solcOutput);
+  t.context.transformFile = (file: string) =>
+    new Transform(t.context.solcInput, t.context.solcOutput, {
+      exclude: source => source !== file,
+    });
 });
 
 test('add namespace', t => {
   const file = 'contracts/namespaces.sol';
-  t.context.transform.apply(transformConstructor(() => true));
-  t.context.transform.apply(removeLeftoverConstructorHead);
-  t.context.transform.apply(removeStateVarInits);
-  t.context.transform.apply(addNamespaceStruct(() => true));
-  t.snapshot(t.context.transform.results()[file]);
+  const transform = t.context.transformFile(file);
+  transform.apply(transformConstructor(() => true));
+  transform.apply(removeLeftoverConstructorHead);
+  transform.apply(removeStateVarInits);
+  transform.apply(addNamespaceStruct(() => true));
+  t.snapshot(transform.results()[file]);
+});
+
+test('error with @custom:storage-size', t => {
+  const file = 'contracts/namespaces-error-storage-size.sol';
+  const transform = t.context.transformFile(file);
+  t.throws(
+    () => transform.apply(addNamespaceStruct(() => true)),
+    { message: 'Cannot combine namespaces with @custom:storage-size annotations (contracts/namespaces-error-storage-size.sol:5)' },
+  );
 });

src/transformations/add-namespace-struct.ts

@@ -8,6 +8,7 @@ import { isStorageVariable } from './utils/is-storage-variable';
 import { erc7201Location } from '../utils/erc7201';
 import { contractStartPosition } from './utils/contract-start-position';
 import { Node } from 'solidity-ast/node';
+import { extractContractStorageSize } from '../utils/natspec';
 
 export function getNamespaceStructName(contractName: string): string {
   return contractName + 'Storage';
@@ -22,6 +23,12 @@ export function addNamespaceStruct(include?: (source: string) => boolean) {
     const { error, resolver } = tools;
 
     for (const contract of findAll('ContractDefinition', sourceUnit)) {
+      const specifiesStorageSize = extractContractStorageSize(contract) !== undefined;
+
+      if (specifiesStorageSize) {
+        throw tools.error(contract, "Cannot combine namespaces with @custom:storage-size annotations");
+      }
+
       let start = contractStartPosition(contract, tools);
 
       let finished = false;

src/transformations/add-storage-gaps.ts

@@ -6,7 +6,7 @@ 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 { extractContractStorageSize } from '../utils/natspec';
 import { decodeTypeIdentifier } from '../utils/type-id';
 import { parseTypeId } from '../utils/parse-type-id';
 import { ASTResolver } from '../ast-resolver';
@@ -21,12 +21,7 @@ export function* addStorageGaps(
 ): Generator<Transformation> {
   for (const contract of findAll('ContractDefinition', sourceUnit)) {
     if (contract.contractKind === '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 targetSlots = extractContractStorageSize(contract) ?? DEFAULT_SLOT_COUNT;
 
       const gapSize = targetSlots - getContractSlotCount(contract, getLayout(contract), resolver);
 

src/utils/extractNatspec.ts renamed to src/utils/natspec.ts

@@ -1,4 +1,4 @@
-import { StructuredDocumentation } from 'solidity-ast';
+import { ContractDefinition, StructuredDocumentation } from 'solidity-ast';
 import { execall } from './execall';
 
 interface NatspecTag {
@@ -26,3 +26,13 @@ export function* extractNatspec(node: {
     }
   }
 }
+
+export function extractContractStorageSize(contract: ContractDefinition): number | undefined {
+  let targetSlots;
+  for (const entry of extractNatspec(contract)) {
+    if (entry.title === 'custom' && entry.tag === 'storage-size') {
+      targetSlots = parseInt(entry.args);
+    }
+  }
+  return targetSlots;
+}