fix: merge internal functions templates

defi-wonderland · Apr 29, 2024 · d6365e5 · d6365e5
1 parent 2476e3e
commit d6365e5
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 53 deletions.
diff --git a/src/context.ts b/src/context.ts
@@ -46,8 +46,8 @@ export function internalFunctionContext(node: FunctionDefinition): InternalFunct
   }
 
   // Check if the function is view or pure
-  const isView = node.stateMutability === FunctionStateMutability.View;
   const isPure = node.stateMutability === FunctionStateMutability.Pure;
+  const isView = node.stateMutability === FunctionStateMutability.View || isPure;
 
   // Save the internal function information
   return {

diff --git a/src/templates/partials/internal-function.hbs b/src/templates/partials/internal-function.hbs
@@ -6,8 +6,8 @@ function mock_call_{{functionName}}({{parameters}}) public {
   );
 }
 
-function {{functionName}}({{inputs}}) internal override {{#if outputs}}returns ({{outputs}}){{/if}} {
-  (bool _success, bytes memory _data) = address(this).call(abi.encodeWithSignature('{{signature}}'{{#if inputs}}, {{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}{{/if}}));
+function {{functionName}}{{#if isPure}}Helper{{/if}}({{inputs}}) internal {{#if isView}}view{{/if}} {{#unless isPure}}override{{/unless}} {{#if outputs}}returns ({{outputs}}){{/if}} {
+  (bool _success, bytes memory _data) = address(this).{{#if isView}}static{{/if}}call(abi.encodeWithSignature('{{signature}}'{{#if inputs}}, {{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}{{/if}}));
 
   if (_success) return abi.decode(_data, ({{#each outputTypes}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}));
 
@@ -23,3 +23,21 @@ function call_{{functionName}}({{inputs}}) public {{#if outputs}}returns ({{outp
 function expectCall_{{functionName}}({{#if inputs}}{{inputs}}{{/if}}) public {
   vm.expectCall(address(this), abi.encodeWithSignature('{{signature}}'{{#if inputs}}, {{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}{{/if}}));
 }
+
+{{#if isPure}}
+
+  function _{{functionName}}CastToPure(function({{#each inputTypes}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}) internal view {{#if outputs}}returns ({{explicitOutputTypes}}){{/if}} fnIn) 
+    internal 
+    pure 
+    returns (function({{#each inputTypes}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}) internal pure {{#if outputs}}returns ({{explicitOutputTypes}}){{/if}} fnOut)
+  {
+    assembly {
+      fnOut := fnIn
+    }
+  }
+
+  function {{functionName}}({{inputs}}) internal pure override {{#if outputs}}returns ({{outputs}}){{/if}} {
+    return _{{functionName}}CastToPure({{functionName}}Helper)({{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}});
+  }
+
+{{/if}}
diff --git a/src/templates/partials/internal-view-function.hbs b/src/templates/partials/internal-view-function.hbs
diff --git a/src/utils.ts b/src/utils.ts
@@ -192,7 +192,6 @@ export async function renderNodeMock(node: ASTNode): Promise<string> {
     constructor: constructorContext,
     'external-or-public-function': externalOrPublicFunctionContext,
     'internal-function': internalFunctionContext,
-    'internal-view-function': internalFunctionContext,
     import: importContext,
   };
 
@@ -222,11 +221,7 @@ export function partialName(node: ASTNode): string {
     } else if (node.visibility === 'external' || node.visibility === 'public') {
       return 'external-or-public-function';
     } else if (node.visibility === 'internal' && node.virtual) {
-      if (node.stateMutability === 'view' || node.stateMutability === 'pure') {
-        return 'internal-view-function';
-      } else {
-        return 'internal-function';
-      }
+      return 'internal-function';
     }
   } else if (node instanceof ImportDirective) {
     return 'import';

diff --git a/test/unit/context/internalFunctionContext.spec.ts b/test/unit/context/internalFunctionContext.spec.ts
@@ -198,9 +198,10 @@ describe('internalFunctionContext', () => {
   it('determines whether the function is view or not', () => {
     for (const stateMutability in FunctionStateMutability) {
       const isView = FunctionStateMutability[stateMutability] === FunctionStateMutability.View;
+      const isPure = FunctionStateMutability[stateMutability] === FunctionStateMutability.Pure;
       const node = mockFunctionDefinition({ ...defaultAttributes, stateMutability: FunctionStateMutability[stateMutability] });
       const context = internalFunctionContext(node);
-      expect(context.isView).to.be.equal(isView);
+      expect(context.isView).to.be.equal(isView || isPure);
     }
   });