From fb245009debc40b11516f4aad675415a8881d365 Mon Sep 17 00:00:00 2001 From: agusduha Date: Wed, 17 Apr 2024 14:03:31 -0300 Subject: [PATCH] feat: improve formatting in templates --- src/smock-foundry.ts | 11 ++- .../partials/array-state-variable.hbs | 20 +++-- .../partials/external-or-public-function.hbs | 13 ++- src/templates/partials/internal-function.hbs | 84 ++++++++++--------- .../partials/mapping-state-variable.hbs | 51 +++++++---- src/templates/partials/state-variable.hbs | 24 ++++-- 6 files changed, 127 insertions(+), 76 deletions(-) diff --git a/src/smock-foundry.ts b/src/smock-foundry.ts index 966561a..8f042d8 100644 --- a/src/smock-foundry.ts +++ b/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); } diff --git a/src/templates/partials/array-state-variable.hbs b/src/templates/partials/array-state-variable.hbs index b6b6d57..73f4fcd 100644 --- a/src/templates/partials/array-state-variable.hbs +++ b/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}} diff --git a/src/templates/partials/external-or-public-function.hbs b/src/templates/partials/external-or-public-function.hbs index 5060bb1..fed97c7 100644 --- a/src/templates/partials/external-or-public-function.hbs +++ b/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}} + ) ); } diff --git a/src/templates/partials/internal-function.hbs b/src/templates/partials/internal-function.hbs index 5afbcd9..4efc9ae 100644 --- a/src/templates/partials/internal-function.hbs +++ b/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}} } diff --git a/src/templates/partials/mapping-state-variable.hbs b/src/templates/partials/mapping-state-variable.hbs index 6cd9796..d3de123 100644 --- a/src/templates/partials/mapping-state-variable.hbs +++ b/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}} diff --git a/src/templates/partials/state-variable.hbs b/src/templates/partials/state-variable.hbs index b0b0d46..3cd2972 100644 --- a/src/templates/partials/state-variable.hbs +++ b/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}}