Disable transpilation of contracts with @Custom:stateless in peerProject mode (#132)

Amxx · frangio · web-flow · commit eb97aef3bc7e · 2023-09-26T22:08:06.000-03:00
Co-authored-by: Francisco &lt;fg@frang.io&gt;
diff --git a/contracts/project/SomeContract.sol b/contracts/project/SomeContract.sol
@@ -3,6 +3,7 @@ pragma solidity ^0.8.0;
 
 import { ISomeInterface } from "./ISomeInterface.sol";
 import { SomeLibrary } from "./SomeLibrary.sol";
+import { SomeStatelessContract } from "./SomeStatelessContract.sol";
 
 interface ISomeContract is ISomeInterface {}
 
@@ -17,11 +18,11 @@ struct SomeStruct {
 contract SomeContract is ISomeContract {
     SomeStruct s;
 
-    function someFunction() public override returns (bool) {
+    function someFunction() public pure override returns (bool) {
         return false;
     }
 
-    function someOtherFunction() public override returns (bool) {
+    function someOtherFunction() public pure override returns (bool) {
         return true;
     }
 
diff --git a/contracts/project/SomeStatelessContract.sol b/contracts/project/SomeStatelessContract.sol
@@ -0,0 +1,5 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.0;
+
+/// @custom:stateless
+contract SomeStatelessContract {}
diff --git a/src/index.ts b/src/index.ts
@@ -3,13 +3,13 @@ import fs from 'fs';
 import { mapValues } from 'lodash';
 import { minimatch } from 'minimatch';
 
-import { Node } from 'solidity-ast/node';
 import { matcher } from './utils/matcher';
 import { renamePath, isRenamed } from './rename';
 import { SolcOutput, SolcInput } from './solc/input-output';
-import { Transform, TransformData } from './transform';
+import { Transform } from './transform';
 import { generateWithInit } from './generate-with-init';
 import { findAlreadyInitializable } from './find-already-initializable';
+import { preparePeerProject } from './prepare-peer-project';
 
 import { fixImportDirectives } from './transformations/fix-import-directives';
 import { renameIdentifiers } from './transformations/rename-identifiers';
@@ -52,8 +52,6 @@ interface TranspileOptions {
   peerProject?: string;
 }
 
-type NodeTransformData = { node: Node; data: Partial<TransformData> };
-
 function getExtraOutputPaths(
   paths: Paths,
   options: TranspileOptions = {},
@@ -73,64 +71,16 @@ function getExtraOutputPaths(
   return outputPaths;
 }
 
-function getExcludeAndImportPathsForPeer(
-  solcOutput: SolcOutput,
-  peerProject: string,
-): [Set<string>, NodeTransformData[]] {
-  const data: NodeTransformData[] = [];
-  const exclude: Set<string> = new Set();
-
-  for (const [source, { ast }] of Object.entries(solcOutput.sources)) {
-    let shouldExclude = true;
-    for (const node of ast.nodes) {
-      switch (node.nodeType) {
-        case 'ContractDefinition': {
-          if (node.contractKind === 'contract') {
-            shouldExclude = false;
-          } else {
-            const importFromPeer = path.join(peerProject, source);
-            data.push({ node, data: { importFromPeer } });
-          }
-          break;
-        }
-        case 'EnumDefinition':
-        case 'ErrorDefinition':
-        case 'FunctionDefinition':
-        case 'StructDefinition':
-        case 'UserDefinedValueTypeDefinition':
-        case 'VariableDeclaration': {
-          const importFromPeer = path.join(peerProject, source);
-          data.push({ node, data: { importFromPeer } });
-          break;
-        }
-        case 'ImportDirective':
-        case 'PragmaDirective':
-        case 'UsingForDirective': {
-          break;
-        }
-      }
-    }
-    if (shouldExclude) {
-      exclude.add(source);
-    }
-  }
-
-  return [exclude, data];
-}
-
 export async function transpile(
   solcInput: SolcInput,
   solcOutput: SolcOutput,
   paths: Paths,
   options: TranspileOptions = {},
 ): Promise<OutputFile[]> {
-  const nodeData: NodeTransformData[] = [];
-
   const outputPaths = getExtraOutputPaths(paths, options);
   const alreadyInitializable = findAlreadyInitializable(solcOutput, options.initializablePath);
 
   const excludeSet = new Set([...alreadyInitializable, ...Object.values(outputPaths)]);
-  const softExcludeSet = new Set();
   const excludeMatch = matcher(options.exclude ?? []);
 
   const namespaceInclude = (source: string) => {
@@ -139,27 +89,12 @@ export async function transpile(
     return namespaced && !namespaceExclude.some(p => minimatch(source, p));
   };
 
-  // if partial transpilation, extract the list of soft exclude, and the peer import paths.
-  if (options.peerProject !== undefined) {
-    const [peerSoftExcludeSet, importFromPeerData] = getExcludeAndImportPathsForPeer(
-      solcOutput,
-      options.peerProject,
-    );
-    peerSoftExcludeSet.forEach(source => softExcludeSet.add(source));
-    nodeData.push(...importFromPeerData);
-  }
-
   const transform = new Transform(solcInput, solcOutput, {
-    exclude: source =>
-      excludeSet.has(source) || (excludeMatch(source) ?? isRenamed(source))
-        ? 'hard'
-        : softExcludeSet.has(source)
-        ? 'soft'
-        : false,
+    exclude: source => excludeSet.has(source) || (excludeMatch(source) ?? isRenamed(source)),
   });
 
-  for (const { node, data } of nodeData) {
-    Object.assign(transform.getData(node), data);
+  if (options.peerProject !== undefined) {
+    preparePeerProject(transform, options.peerProject);
   }
 
   transform.apply(renameIdentifiers);
diff --git a/src/partial-transform.test.ts.snap b/src/partial-transform.test.ts.snap
diff --git a/src/prepare-peer-project.test.ts b/src/prepare-peer-project.test.ts
@@ -1,5 +1,6 @@
 import _test, { TestFn } from 'ava';
 import hre from 'hardhat';
+import path from 'path';
 
 import { getBuildInfo } from './test-utils/get-build-info';
 import { OutputFile, transpile } from '.';
@@ -8,6 +9,7 @@ import { SolcOutput } from './solc/input-output';
 const test = _test as TestFn<Context>;
 
 interface Context {
+  inputs: string[];
   files: OutputFile[];
 }
 
@@ -20,23 +22,31 @@ test.serial.before('compile', async t => {
   const solcOutput = buildInfo.output as SolcOutput;
   const exclude = Object.keys(solcOutput.sources).filter(path => !path.startsWith(projectDir));
 
+  t.context.inputs = Object.keys(solcInput.sources);
   t.context.files = await transpile(solcInput, solcOutput, hre.config.paths, {
     exclude,
     peerProject,
   });
 });
 
-for (const fileName of ['ISomeInterface.sol', 'SomeLibrary.sol']) {
+for (const fileName of ['ISomeInterface.sol', 'SomeLibrary.sol', 'SomeStatelessContract.sol']) {
   test(`do not transpile ${fileName}`, t => {
     const file = t.context.files.find(f => f.fileName === fileName);
+    // source file exists
+    t.true(t.context.inputs.includes(path.join(projectDir, fileName)));
+    // transpiled file does not exist
     t.is(file, undefined, 'file should not be transpiled');
   });
 }
 
 for (const fileName of ['SomeContract.sol', 'SomeOtherContract.sol']) {
   test(`transpile ${fileName}`, t => {
     const file = t.context.files.find(f => f.fileName === fileName);
+    // source file exists
+    t.true(t.context.inputs.includes(path.join(projectDir, fileName)));
+    // transpiled file exists
     t.not(file, undefined, 'file not found');
+    // snapshot
     t.snapshot(file);
   });
 }
diff --git a/src/prepare-peer-project.test.ts.md b/src/prepare-peer-project.test.ts.md
@@ -16,6 +16,7 @@ Generated by [AVA](https://avajs.dev).
       ␊
       import {ISomeInterface} from "@openzeppelin/contracts/project/ISomeInterface.sol";␊
       import {SomeLibrary} from "@openzeppelin/contracts/project/SomeLibrary.sol";␊
+      import {SomeStatelessContract} from "@openzeppelin/contracts/project/SomeStatelessContract.sol";␊
       import {Initializable} from "../Initializable.sol";␊
       ␊
       import { ISomeContract } from "@openzeppelin/contracts/project/SomeContract.sol";␊
@@ -34,11 +35,11 @@ Generated by [AVA](https://avajs.dev).
       ␊
           function __SomeContract_init_unchained() internal onlyInitializing {␊
           }␊
-          function someFunction() public override returns (bool) {␊
+          function someFunction() public pure override returns (bool) {␊
               return false;␊
           }␊
       ␊
-          function someOtherFunction() public override returns (bool) {␊
+          function someOtherFunction() public pure override returns (bool) {␊
               return true;␊
           }␊
       ␊
diff --git a/src/prepare-peer-project.test.ts.snap b/src/prepare-peer-project.test.ts.snap
diff --git a/src/prepare-peer-project.ts b/src/prepare-peer-project.ts
@@ -0,0 +1,65 @@
+import path from 'path';
+
+import { Transform } from './transform';
+import { extractContractStorageSize, extractContractStateless } from './utils/natspec';
+import { isStorageVariable } from './transformations/utils/is-storage-variable';
+
+export function preparePeerProject(transform: Transform, peerProject: string) {
+  for (const ast of transform.asts()) {
+    let shouldExclude = true;
+    for (const node of ast.nodes) {
+      switch (node.nodeType) {
+        case 'ContractDefinition': {
+          if (node.contractKind === 'contract') {
+            if (!extractContractStateless(node)) {
+              shouldExclude = false;
+              break;
+            }
+            if (extractContractStorageSize(node) !== undefined) {
+              throw transform.error(
+                node,
+                'Contract marked as stateless should not have an associated storage size',
+              );
+            }
+            for (const decl of node.nodes) {
+              if (
+                decl.nodeType == 'VariableDeclaration' &&
+                isStorageVariable(decl, transform.resolver)
+              ) {
+                throw transform.error(
+                  node,
+                  'Contract marked as stateless should not contain storage variable declarations',
+                );
+              }
+              if (decl.nodeType == 'FunctionDefinition' && decl.kind == 'constructor') {
+                throw transform.error(
+                  node,
+                  'Contract marked as stateless should not have a constructor',
+                );
+              }
+            }
+          }
+          transform.getData(node).importFromPeer = path.join(peerProject, ast.absolutePath);
+          break;
+        }
+        case 'EnumDefinition':
+        case 'ErrorDefinition':
+        case 'FunctionDefinition':
+        case 'StructDefinition':
+        case 'UserDefinedValueTypeDefinition':
+        case 'VariableDeclaration': {
+          transform.getData(node).importFromPeer = path.join(peerProject, ast.absolutePath);
+          break;
+        }
+        case 'ImportDirective':
+        case 'PragmaDirective':
+        case 'UsingForDirective': {
+          break;
+        }
+      }
+    }
+    if (shouldExclude) {
+      transform.exclude(ast.absolutePath);
+    }
+  }
+}
diff --git a/src/transform-namespaces.test.ts b/src/transform-namespaces.test.ts
@@ -30,7 +30,7 @@ test.serial.before('compile', async t => {
 test.beforeEach('transform', async t => {
   t.context.transformFile = (file: string) =>
     new Transform(t.context.solcInput, t.context.solcOutput, {
-      exclude: source => (source !== file ? 'hard' : false),
+      exclude: source => source !== file,
     });
 });
 
diff --git a/src/transform.test.ts b/src/transform.test.ts
@@ -39,13 +39,13 @@ test.serial.before('compile', async t => {
 
   t.context.transformFile = (file: string) =>
     new Transform(t.context.solcInput, t.context.solcOutput, {
-      exclude: source => (source !== file ? 'hard' : false),
+      exclude: source => source !== file,
     });
 });
 
 test.beforeEach('transform', async t => {
   t.context.transform = new Transform(t.context.solcInput, t.context.solcOutput, {
-    exclude: source => (source.startsWith('contracts/invalid/') ? 'hard' : false),
+    exclude: source => source.startsWith('contracts/invalid/'),
   });
 });
 
@@ -191,7 +191,7 @@ test('fix new statement in var init', t => {
 test('exclude', t => {
   const file = 'contracts/TransformInitializable.sol';
   const transform = new Transform(t.context.solcInput, t.context.solcOutput, {
-    exclude: s => (s === file ? 'hard' : false),
+    exclude: s => s === file,
   });
   // eslint-disable-next-line require-yield
   transform.apply(function* (s) {
diff --git a/src/transform.ts b/src/transform.ts
@@ -45,7 +45,7 @@ interface TransformState {
 }
 
 interface TransformOptions {
-  exclude?: (source: string) => 'hard' | 'soft' | false;
+  exclude?: (source: string) => boolean;
 }
 
 export class Transform {
@@ -62,7 +62,7 @@ export class Transform {
   constructor(input: SolcInput, output: SolcOutput, options: TransformOptions = {}) {
     this.decodeSrc = srcDecoder(output);
     this.getLayout = layoutGetter(output);
-    this.resolver = new ASTResolver(output, source => options.exclude?.(source) === 'hard');
+    this.resolver = new ASTResolver(output, options.exclude);
 
     for (const source in input.sources) {
       if (options.exclude?.(source)) {
@@ -186,6 +186,10 @@ export class Transform {
   asts(): SourceUnit[] {
     return Object.values(this.state).map(s => s.ast);
   }
+
+  exclude(source: string) {
+    delete this.state[source];
+  }
 }
 
 function insertSortedAndValidate(
diff --git a/src/utils/natspec.ts b/src/utils/natspec.ts
@@ -14,7 +14,7 @@ export function* extractNatspec(node: {
     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,
+    /^\s*(?:@(?<title>\w+)(?::(?<tag>[a-z][a-z-]*))?)?(?: (?<args>(?:(?!^\s@\w+)[^])*))?/m,
     doc,
   )) {
     if (groups) {
@@ -36,3 +36,12 @@ export function extractContractStorageSize(contract: ContractDefinition): number
   }
   return targetSlots;
 }
+
+export function extractContractStateless(contract: ContractDefinition): boolean {
+  for (const entry of extractNatspec(contract)) {
+    if (entry.title === 'custom' && entry.tag === 'stateless') {
+      return true;
+    }
+  }
+  return false;
+}
