feat: improve formatting in templates

src/smock-foundry.ts

@@ -11,6 +11,7 @@ import {
 } from './utils';
 import path from 'path';
 import { ensureDir } from 'fs-extra';
+import { exec } from 'child_process';
 
 /**
  * Generates the mock contracts
@@ -22,7 +23,8 @@ export async function generateMockContracts(
   mocksDirectory: string,
   ignoreDirectories: string[],
 ) {
-  await emptySmockDirectory(mocksDirectory);
+  const mocksPath = path.resolve(rootPath, mocksDirectory);
+  await emptySmockDirectory(mocksPath);
   const contractTemplate = getContractTemplate();
 
   try {
@@ -88,12 +90,15 @@ export async function generateMockContracts(
       // Generate SmockHelper contract
       const smockHelperTemplate = await getSmockHelperTemplate();
       const smockHelperCode: string = smockHelperTemplate({});
-      const helperPath = path.resolve(rootPath, mocksDirectory);
-      writeFileSync(`${helperPath}/SmockHelper.sol`, smockHelperCode);
+      writeFileSync(`${mocksPath}/SmockHelper.sol`, smockHelperCode);
 
       console.log('Mock contracts generated successfully');
 
       await compileSolidityFilesFoundry(rootPath, mocksDirectory);
+
+      // Format the generated files
+      console.log('Formatting generated files...');
+      exec(`forge fmt --root ${rootPath} ${mocksPath}`);
     } catch (error) {
       console.error(error);
     }

src/templates/partials/array-state-variable.hbs

@@ -1,19 +1,27 @@
 function set_{{setFunction.functionName}}({{setFunction.arrayType}} _{{setFunction.paramName}}) public {
   {{#if isStructArray}}
-  for (uint256 _i; _i < _{{setFunction.paramName}}.length; ++_i) {
-    {{setFunction.functionName}}.push(_{{setFunction.paramName}}[_i]);
-  }
+    for (uint256 _i; _i < _{{setFunction.paramName}}.length; ++_i) {
+      {{setFunction.functionName}}.push(_{{setFunction.paramName}}[_i]);
+    }
   {{else}}
-  {{setFunction.functionName}} = _{{setFunction.paramName}};
+    {{setFunction.functionName}} = _{{setFunction.paramName}};
   {{/if}}
 }
-{{#unless isInternal}}
 
+{{#unless isInternal}}
 function mock_call_{{mockFunction.functionName}}(uint256 _index, {{mockFunction.baseType}} _value) public {
   vm.mockCall(
     address(this),
     abi.encodeWithSignature('{{mockFunction.functionName}}(uint256)', _index),
-    abi.encode({{#if isStructArray}}{{#each mockFunction.structFields}}_value.{{this}}{{#unless @last}}, {{/unless}}{{/each}}{{else}}_value{{/if}})
+    abi.encode(
+      {{#if isStructArray}}
+        {{#each mockFunction.structFields}}
+          _value.{{this}}{{#unless @last}}, {{/unless}}
+        {{/each}}
+      {{else}}
+        _value
+      {{/if}}
+    )
   );
 }
 {{/unless}}

src/templates/partials/external-or-public-function.hbs

@@ -1,11 +1,16 @@
 {{#unless implemented}}
-function {{functionName}}({{inputs}}) {{visibility}} {{stateMutability}} override{{#if overrides}}{{overrides}}{{/if}} {{#if outputs}}returns ({{outputs}}){{/if}} {}
-
+  function {{functionName}}({{inputs}}) {{visibility}} {{stateMutability}} override {{#if overrides}}{{overrides}}{{/if}} {{#if outputs}}returns ({{outputs}}){{/if}} {}
 {{/unless}}
+
 function mock_call_{{functionName}}({{parameters}}) public {
   vm.mockCall(
     address(this),
-    abi.encodeWithSignature('{{signature}}'{{#if inputs}}, {{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}{{/if}}),
-    abi.encode({{#each outputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}})
+    abi.encodeWithSignature(
+      '{{signature}}'
+      {{#if inputs}}, {{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}{{/if}}
+    ),
+    abi.encode(
+      {{#each outputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}
+    )
   );
 }

src/templates/partials/internal-function.hbs

@@ -1,54 +1,62 @@
 {{#if isView}}
-{{#if outputs}}
-struct {{functionName}}Output {
-  {{#each outputTypes}}
-  {{this}} {{lookup ../outputNames @index}};
-  {{/each}}
-}
+  {{#if outputs}}
+    struct {{functionName}}Output {
+      {{#each outputTypes}}
+        {{this}} {{lookup ../outputNames @index}};
+      {{/each}}
+    }
 
-mapping(bytes32 => {{functionName}}Output) private {{functionName}}Outputs;
-{{/if}}
-bytes32[] private {{functionName}}InputHashes;
+    mapping(bytes32 => {{functionName}}Output) private {{functionName}}Outputs;
+  {{/if}}
 
+  bytes32[] private {{functionName}}InputHashes;
 {{/if}}
+
 function mock_call_{{functionName}}({{parameters}}) public {
   {{#if isView}}
-  bytes32 _key = keccak256(abi.encode({{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}));
-  {{#if outputs}}
-  {{functionName}}Outputs[_key] = {{functionName}}Output({{#each outputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}});
-  {{/if}}
-  for (uint256 _i; _i < {{functionName}}InputHashes.length; ++_i) {
-    if (_key == {{functionName}}InputHashes[_i]) return;
-  }
-  {{functionName}}InputHashes.push(_key);
+    bytes32 _key = keccak256(abi.encode({{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}));
+    {{#if outputs}}
+      {{functionName}}Outputs[_key] = {{functionName}}Output({{#each outputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}});
+    {{/if}}
+
+    for (uint256 _i; _i < {{functionName}}InputHashes.length; ++_i) {
+      if (_key == {{functionName}}InputHashes[_i]) return;
+    }
+
+    {{functionName}}InputHashes.push(_key);
   {{else}}
-  vm.mockCall(
-    address(this),
-    abi.encodeWithSignature('{{signature}}'{{#if inputs}}, {{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}{{/if}}),
-    abi.encode({{#each outputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}})
-  );
+    vm.mockCall(
+      address(this),
+      abi.encodeWithSignature('{{signature}}'{{#if inputs}}, {{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}{{/if}}),
+      abi.encode({{#each outputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}})
+    );
   {{/if}}
 }
 
 function {{functionName}}({{inputs}}) internal {{#if isView}}view {{/if}}override {{#if outputs}}returns ({{outputs}}){{/if}} {
   {{#if isView}}
-  bytes32 _key = keccak256(abi.encode({{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}));
-  for (uint256 _i; _i < {{functionName}}InputHashes.length; ++_i) {
-    if (_key == {{functionName}}InputHashes[_i]) {
-      {{#if outputs}}
-      {{functionName}}Output memory _output = {{functionName}}Outputs[_key];
-      {{/if}}
-      return ({{#each outputNames}}_output.{{this}}{{#unless @last}}, {{/unless}}{{/each}});
+    bytes32 _key = keccak256(abi.encode({{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}));
+
+    for (uint256 _i; _i < {{functionName}}InputHashes.length; ++_i) {
+      if (_key == {{functionName}}InputHashes[_i]) {
+        {{#if outputs}}
+          {{functionName}}Output memory _output = {{functionName}}Outputs[_key];
+        {{/if}}
+
+        return ({{#each outputNames}}_output.{{this}}{{#unless @last}}, {{/unless}}{{/each}});
+      }
     }
-  }
-  {{#if implemented}}
-  return super.{{functionName}}({{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}});
-  {{/if}}
+    
+    {{#if implemented}}
+      return super.{{functionName}}({{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}});
+    {{/if}}
   {{else}}
-  (bool _success, bytes memory _data) = address(this).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}}));
-  {{#if implemented}}
-  else return super.{{functionName}}({{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}});
-  {{/if}}
+    (bool _success, bytes memory _data) = address(this).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}}));
+
+    {{#if implemented}}
+      else return super.{{functionName}}({{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}});
+    {{/if}}
   {{/if}}
 }

src/templates/partials/mapping-state-variable.hbs

@@ -1,21 +1,38 @@
 {{#unless hasNestedMapping}}
-function set_{{setFunction.functionName}}({{#each setFunction.keyTypes}}{{this}} _key{{@index}}, {{/each}}{{setFunction.valueType}} _value) public {
-  {{#if isStructArray}}
-  for (uint256 _i; _i < _value.length; ++_i) {
-    {{setFunction.functionName}}{{#each setFunction.keyTypes}}[_key{{@index}}]{{/each}}.push(_value[_i]);
+  function set_{{setFunction.functionName}}(
+    {{#each setFunction.keyTypes}}{{this}} _key{{@index}}, {{/each}}
+    {{setFunction.valueType}} _value
+  ) public {
+    {{#if isStructArray}}
+      for (uint256 _i; _i < _value.length; ++_i) {
+        {{setFunction.functionName}}
+          {{#each setFunction.keyTypes}}[_key{{@index}}]{{/each}}.push(_value[_i]);
+      }
+    {{else}}
+      {{setFunction.functionName}}
+        {{#each setFunction.keyTypes}}[_key{{@index}}]{{/each}} = _value;
+    {{/if}}
   }
-  {{else}}
-  {{setFunction.functionName}}{{#each setFunction.keyTypes}}[_key{{@index}}]{{/each}} = _value;
-  {{/if}}
-}
-{{#unless isInternal}}
 
-function mock_call_{{mockFunction.functionName}}({{#each mockFunction.keyTypes}}{{this}} _key{{@index}}, {{/each}}{{#if isArray}}uint256 _index, {{mockFunction.baseType}}{{else}}{{mockFunction.valueType}}{{/if}} _value) public {
-  vm.mockCall(
-    address(this),
-    abi.encodeWithSignature('{{mockFunction.functionName}}({{#each mockFunction.keyTypes}}{{this}}{{#unless @last}},{{/unless}}{{/each}}{{#if isArray}},uint256{{/if}})'{{#each mockFunction.keyTypes}}, _key{{@index}}{{/each}}{{#if isArray}}, _index{{/if}}),
-    abi.encode({{#if isStruct}}{{#each mockFunction.structFields}}_value.{{this}}{{#unless @last}}, {{/unless}}{{/each}}{{else}}_value{{/if}})
-  );
-}
-{{/unless}}
+  {{#unless isInternal}}
+    function mock_call_{{mockFunction.functionName}}(
+      {{#each mockFunction.keyTypes}}{{this}} _key{{@index}}, {{/each}}
+      {{#if isArray}}uint256 _index, {{mockFunction.baseType}}{{else}}{{mockFunction.valueType}}{{/if}} _value
+    ) public {
+      vm.mockCall(
+        address(this),
+        abi.encodeWithSignature(
+          '{{mockFunction.functionName}}({{#each mockFunction.keyTypes}}{{this}}{{#unless @last}},{{/unless}}{{/each}}{{#if isArray}},uint256{{/if}})'
+          {{#each mockFunction.keyTypes}}, _key{{@index}}{{/each}}{{#if isArray}}, _index{{/if}}
+        ),
+        abi.encode(
+          {{#if isStruct}}
+            {{#each mockFunction.structFields}}_value.{{this}}{{#unless @last}}, {{/unless}}{{/each}}
+          {{else}}
+            _value
+          {{/if}}
+        )
+      );
+    }
+  {{/unless}}
 {{/unless}}

src/templates/partials/state-variable.hbs

@@ -1,13 +1,21 @@
 function set_{{setFunction.functionName}}({{setFunction.paramType}} _{{setFunction.paramName}}) public {
   {{setFunction.paramName}} = _{{setFunction.paramName}};
 }
-{{#unless isInternal}}
 
-function mock_call_{{mockFunction.functionName}}({{mockFunction.paramType}} _value) public {
-  vm.mockCall(
-    address(this),
-    abi.encodeWithSignature('{{mockFunction.functionName}}()'),
-    abi.encode({{#if isStruct}}{{#each mockFunction.structFields}}_value.{{this}}{{#unless @last}}, {{/unless}}{{/each}}{{else}}_value{{/if}})
-  );
-}
+{{#unless isInternal}}
+  function mock_call_{{mockFunction.functionName}}({{mockFunction.paramType}} _value) public {
+    vm.mockCall(
+      address(this),
+      abi.encodeWithSignature('{{mockFunction.functionName}}()'),
+      abi.encode(
+        {{#if isStruct}}
+          {{#each mockFunction.structFields}}
+            _value.{{this}}{{#unless @last}}, {{/unless}}
+          {{/each}}
+        {{else}}
+          _value
+        {{/if}}
+      )
+    );
+  }
 {{/unless}}