refactor: typeFix(...)

src/get-constructor.ts

@@ -1,4 +1,5 @@
 import { ContractDefinitionNode, FunctionDefinitionNode, VariableDeclarationNode } from './types';
+import { typeFix } from './utils';
 
 /**
  * Return the constructor of the contract
@@ -21,31 +22,34 @@ export const getConstructor = (contractNode: ContractDefinitionNode): string =>
   const parameters: VariableDeclarationNode[] = constructorNode.parameters.parameters ? constructorNode.parameters.parameters : [];
 
   // We save the parameters in an array with their types and storage location
-  const mockConstructorParameters: string[] = [];
-  // We save the parameters names in an other array
+  const constructorParameters: string[] = [];
+  // We save the parameters names in another array
   const parameterNames: string[] = [];
 
+  let parameterIndex = 0;
   parameters.forEach((parameter) => {
+    // We remove the 'contract ' string from the type name if it exists
+    const typeName: string = typeFix(parameter.typeDescriptions.typeString);
+    const paramName: string = parameter.name == '' ? `_param${parameterIndex}` : parameter.name;
+
     // If the storage location is memory or calldata then we keep it
     const storageLocation =
       parameter.storageLocation === 'memory' || parameter.storageLocation === 'calldata' ? `${parameter.storageLocation} ` : '';
 
-    // We remove the 'contract ' string from the type name if it exists
-    const typeName: string = parameter.typeDescriptions.typeString.replace(/contract |struct |enum /g, '');
-
     // We create the string that will be used in the constructor signature
-    const parameterString = `${typeName} ${storageLocation}${parameter.name}`;
+    const parameterString = `${typeName} ${storageLocation}${paramName}`;
 
     // We save the strings in the arrays
-    mockConstructorParameters.push(parameterString);
-    parameterNames.push(parameter.name);
+    constructorParameters.push(parameterString);
+    parameterNames.push(paramName);
+    parameterIndex++;
   });
 
   // We join the arrays with a comma and a space between them
-  const mockConstructorParametersString = mockConstructorParameters.join(', ');
+  const constructorParametersString = constructorParameters.join(', ');
   const parameterNamesString = parameterNames.join(', ');
 
   // Create the constructor signature and return it
-  const constructorSignature = `constructor(${mockConstructorParametersString}) ${contractName}(${parameterNamesString}) {}`;
+  const constructorSignature = `constructor(${constructorParametersString}) ${contractName}(${parameterNamesString}) {}`;
   return constructorSignature;
 };

src/get-external-functions.ts

@@ -1,4 +1,5 @@
 import { ContractDefinitionNode, FunctionDefinitionNode, VariableDeclarationNode, ExternalFunctionOptions } from './types';
+import { typeFix } from './utils';
 
 /**
  * Returns the information of the external function for the mock contract
@@ -29,25 +30,27 @@ export const getExternalMockFunctions = (contractNode: ContractDefinitionNode):
 
     // We save the parameters in an array with their types and storage location
     const functionParameters: string[] = [];
-    // We save the parameters names in an other array
-    const parameterNames: string[] = [];
     // We save the types in an array to use them in order to create the signature
     const parameterTypes: string[] = [];
+    // We save the parameters names in another array
+    const parameterNames: string[] = [];
+
     let parameterIndex = 0;
     parameters.forEach((parameter: VariableDeclarationNode) => {
+      // We remove the 'contract ' string from the type name if it exists
+      const typeName: string = typeFix(parameter.typeDescriptions.typeString);
+      const paramName: string = parameter.name == '' ? `_param${parameterIndex}` : parameter.name;
+
       // If the storage location is memory or calldata then we keep it
       const storageLocation =
         parameter.storageLocation === 'memory' || parameter.storageLocation === 'calldata' ? `${parameter.storageLocation} ` : '';
 
-      // We remove the 'contract ' string from the type name if it exists
-      const typeName: string = parameter.typeDescriptions.typeString.replace(/contract |struct |enum /g, '');
-
-      const paramName: string = parameter.name == '' ? `_param${parameterIndex}` : parameter.name;
       // We create the string that will be used in the constructor signature
       const parameterString = `${typeName} ${storageLocation}${paramName}`;
 
-      parameterTypes.push(typeName);
+      // We save the strings in the arrays
       functionParameters.push(parameterString);
+      parameterTypes.push(typeName);
       parameterNames.push(paramName);
       parameterIndex++;
     });
@@ -60,16 +63,17 @@ export const getExternalMockFunctions = (contractNode: ContractDefinitionNode):
     const functionReturnParameters: string[] = [];
     // We save the return parameters names in an other array
     const returnParameterNames: string[] = [];
+
     parameterIndex = 0;
     returnParameters.forEach((parameter: VariableDeclarationNode) => {
+      // We remove the 'contract ' string from the type name if it exists
+      const typeName: string = typeFix(parameter.typeDescriptions.typeString);
+      const returnName: string = parameter.name == '' ? `_return${parameterIndex}` : parameter.name;
+
       // If the storage location is memory or calldata then we keep it
       const storageLocation =
         parameter.storageLocation === 'memory' || parameter.storageLocation === 'calldata' ? `${parameter.storageLocation} ` : '';
 
-      // We remove the 'contract ' string from the type name if it exists
-      const typeName: string = parameter.typeDescriptions.typeString.replace(/contract |struct |enum /g, '');
-
-      const returnName: string = parameter.name == '' ? `_return${parameterIndex}` : parameter.name;
       // We create the string that will be used in the constructor signature
       const parameterString = `${typeName} ${storageLocation}${returnName}`;
 

src/get-internal-functions.ts

@@ -1,4 +1,5 @@
 import { ContractDefinitionNode, FunctionDefinitionNode, VariableDeclarationNode, InternalFunctionOptions } from './types';
+import { typeFix } from './utils';
 
 /**
  * Returns the information of the internal function for the mock contract
@@ -33,12 +34,13 @@ export const getInternalMockFunctions = (contractNode: ContractDefinitionNode):
     let parameterIndex = 0;
     parameters.forEach((parameter: VariableDeclarationNode) => {
       // We remove the 'contract ' string from the type name if it exists
-      const typeName: string = parameter.typeDescriptions.typeString.replace(/contract |struct |enum /g, '');
+      const typeName: string = typeFix(parameter.typeDescriptions.typeString);
       const paramName: string = parameter.name == '' ? `_param${parameterIndex}` : parameter.name;
 
       // If the storage location is memory or calldata then we keep it
       const storageLocation =
         parameter.storageLocation === 'memory' || parameter.storageLocation === 'calldata' ? `${parameter.storageLocation} ` : '';
+
       // We create the string that will be used in the constructor signature
       const parameterString = `${typeName} ${storageLocation}${paramName}`;
 
@@ -63,15 +65,17 @@ export const getInternalMockFunctions = (contractNode: ContractDefinitionNode):
     parameterIndex = 0;
     returnParameters.forEach((parameter: VariableDeclarationNode) => {
       // We remove the 'contract ' string from the type name if it exists
-      const typeName: string = parameter.typeDescriptions.typeString.replace(/contract |struct |enum /g, '');
+      const typeName: string = typeFix(parameter.typeDescriptions.typeString);
       const paramName: string = parameter.name == '' ? `_returnParam${parameterIndex}` : parameter.name;
 
       // If the storage location is memory or calldata then we keep it
       const storageLocation =
         parameter.storageLocation === 'memory' || parameter.storageLocation === 'calldata' ? `${parameter.storageLocation} ` : '';
+
       // We create the string that will be used in the constructor signature
       const parameterString = `${typeName} ${storageLocation}${paramName}`;
 
+      // We save the strings in the arrays
       functionReturnParameters.push(parameterString);
       returnParameterTypes.push(typeName);
       returnParameterNames.push(paramName);

src/get-variables-functions.ts

@@ -6,7 +6,7 @@ import {
   MappingStateVariableOptions,
   StateVariablesOptions,
 } from './types';
-import { typeFix } from './utils';
+import { typeFix, explicitTypeStorageLocation } from './utils';
 
 /**
  * Returns all the mock functions information for state variables
@@ -65,13 +65,13 @@ function getArrayFunction(arrayNode: VariableDeclarationNode): ArrayStateVariabl
   const arrayName: string = arrayNode.name;
 
   // Array type
-  const arrayType: string = typeFix(arrayNode.typeDescriptions.typeString).replace(/contract |struct |enum /g, '');
+  const arrayType: string = typeFix(explicitTypeStorageLocation(arrayNode.typeDescriptions.typeString));
 
   // Base type
-  const baseType: string = typeFix(arrayNode.typeName.baseType.typeDescriptions.typeString).replace(/contract |struct |enum /g, '');
+  const baseType: string = typeFix(explicitTypeStorageLocation(arrayNode.typeName.baseType.typeDescriptions.typeString));
 
   // Struct flag
-  const isStruct: boolean = arrayNode.typeDescriptions.typeString.includes('struct ');
+  const isStruct: boolean = arrayNode.typeDescriptions.typeString.startsWith('struct ');
 
   // If the array is internal we don't create mockCall for it
   const isInternal: boolean = arrayNode.visibility == 'internal';
@@ -112,22 +112,22 @@ function getMappingFunction(mappingNode: VariableDeclarationNode): MappingStateV
   const keyTypes: string[] = [];
 
   do {
-    const keyType: string = typeFix(mappingTypeNameNode.keyType.typeDescriptions.typeString).replace(/contract |struct |enum /g, '');
+    const keyType: string = typeFix(explicitTypeStorageLocation(mappingTypeNameNode.keyType.typeDescriptions.typeString));
     keyTypes.push(keyType);
     mappingTypeNameNode = mappingTypeNameNode.valueType;
   } while (mappingTypeNameNode.typeDescriptions.typeString.startsWith('mapping'));
 
   // Value type
-  const valueType: string = typeFix(mappingTypeNameNode.typeDescriptions.typeString).replace(/contract |struct |enum /g, '');
+  const valueType: string = typeFix(explicitTypeStorageLocation(mappingTypeNameNode.typeDescriptions.typeString));
 
   // Array flag
   const isArray: boolean = valueType.includes('[]');
 
   // Base type
-  const baseType: string = isArray ? typeFix(mappingTypeNameNode.baseType.typeDescriptions.typeString).replace(/contract |struct |enum /g, '') : valueType;
+  const baseType: string = isArray ? typeFix(explicitTypeStorageLocation(mappingTypeNameNode.baseType.typeDescriptions.typeString)) : valueType;
 
   // Struct array flag
-  const isStructArray: boolean = isArray && mappingTypeNameNode.typeDescriptions.typeString.includes('struct ');
+  const isStructArray: boolean = isArray && mappingTypeNameNode.typeDescriptions.typeString.startsWith('struct ');
 
   // If the mapping is internal we don't create mockCall for it
   const isInternal: boolean = mappingNode.visibility == 'internal';
@@ -164,7 +164,7 @@ function getBasicStateVariableFunction(variableNode: VariableDeclarationNode): B
   const variableName: string = variableNode.name;
 
   // remove spec type leading string
-  const variableType: string = typeFix(variableNode.typeDescriptions.typeString).replace(/contract |struct |enum /g, '');
+  const variableType: string = typeFix(explicitTypeStorageLocation(variableNode.typeDescriptions.typeString));
 
   // If the variable is internal we don't create mockCall for it
   const isInternal: boolean = variableNode.visibility == 'internal';

src/types.ts

@@ -159,7 +159,5 @@ export interface InternalFunctionOptions {
   isView: boolean;
 }
 
-export const memoryTypes = ['string', 'bytes', 'mapping'];
-
-export const arrayRegex = /(\w+)\[\]/;
-export const structRegex = /struct (\w+)/;
+export const userDefinedTypes = ['contract', 'enum', 'struct'];
+export const explicitTypes = ['string', 'bytes', 'mapping', 'struct'];

src/utils.ts

@@ -1,4 +1,4 @@
-import { arrayRegex, memoryTypes, structRegex } from './types';
+import { userDefinedTypes, explicitTypes } from './types';
 import { resolve, join, relative, dirname } from 'path';
 import { readFileSync, readdirSync, statSync } from 'fs';
 import { exec } from 'child_process';
@@ -40,17 +40,26 @@ export const lowercaseFirstLetter = (str: string): string => {
 };
 
 /**
- * Returns the string with the type fixed, for example if the type is string it will return `string memory`
- * @param str The string to fix the type
- * @returns Returns the string with the type fixed
+ * Fixes user-defined types
+ * @param type The string of the type to fix
+ * @returns The string with the type fixed
  */
-export const typeFix = (str: string): string => {
-  if (memoryTypes.includes(str) || arrayRegex.exec(str)) {
-    return `${str} memory`;
-  } else if (structRegex.exec(str)) {
-    return `${str} memory`;
+export const typeFix = (type: string): string => {
+  const regExp = new RegExp(`^(${userDefinedTypes.join('|')}) `);
+  return type.replace(regExp, '');
+};
+
+/**
+ * Explicits a type's storage location, if required
+ * @param type The string of the type to explicit
+ * @returns The string with the type explicit
+ */
+export const explicitTypeStorageLocation = (type: string): string => {
+  const regExp = new RegExp(`^(${explicitTypes.join('|')})\\b`);
+  if (regExp.test(type) || type.includes('[]')) {
+    return `${type} memory`;
   } else {
-    return str;
+    return type;
   }
 };