refactor

Author: zemse

src/HuffConfig.sol

@@ -10,6 +10,12 @@ contract HuffConfig {
     /// @notice Initializes cheat codes in order to use ffi to compile Huff contracts
     Vm public constant vm = Vm(address(bytes20(uint160(uint256(keccak256("hevm cheat code"))))));
 
+    /// @notice Enum that declares two possible values for a CodeType parameter
+    enum CodeType {
+        CREATION,
+        RUNTIME
+    }
+
     /// @notice Struct that represents a constant to be passed to the `-c` flag
     struct Constant {
         string key;
@@ -144,82 +150,20 @@ contract HuffConfig {
 
     /// @notice Get the creation bytecode of a contract
     function creation_code(string memory file) public payable returns (bytes memory bytecode) {
-        binary_check();
-
-        // Split the file into its parts
-        strings.slice memory s = file.toSlice();
-        strings.slice memory delim = "/".toSlice();
-        string[] memory parts = new string[](s.count(delim) + 1);
-        for (uint256 i = 0; i < parts.length; i++) {
-            parts[i] = s.split(delim).toString();
-        }
-
-        // Get the system time with our script
-        string[] memory time = new string[](1);
-        time[0] = "./lib/foundry-huff/scripts/rand_bytes.sh";
-        bytes memory retData = vm.ffi(time);
-        string memory rand_bytes = bytes32ToString(keccak256(abi.encode(bytes32(retData))));
-
-        // Re-concatenate the file with a "__TEMP__" prefix
-        string memory tempFile = parts[0];
-        if (parts.length <= 1) {
-            tempFile = string.concat("__TEMP__", rand_bytes, tempFile);
-        } else {
-            for (uint256 i = 1; i < parts.length - 1; i++) {
-                tempFile = string.concat(tempFile, "/", parts[i]);
-            }
-            tempFile = string.concat(tempFile, "/", "__TEMP__", rand_bytes, parts[parts.length - 1]);
-        }
-
-        // Paste the code in a new temp file
-        string[] memory create_cmds = new string[](3);
-        // TODO: create_cmds[0] = "$(find . -name \"file_writer.sh\")";
-        create_cmds[0] = "./lib/foundry-huff/scripts/file_writer.sh";
-        create_cmds[1] = string.concat("src/", tempFile, ".huff");
-        create_cmds[2] = string.concat(code, "\n");
-        vm.ffi(create_cmds);
-
-        // Append the real code to the temp file
-        string[] memory append_cmds = new string[](3);
-        append_cmds[0] = "./lib/foundry-huff/scripts/read_and_append.sh";
-        append_cmds[1] = string.concat("src/", tempFile, ".huff");
-        append_cmds[2] = string.concat("src/", file, ".huff");
-        vm.ffi(append_cmds);
-
-        /// Create a list of strings with the commands necessary to compile Huff contracts
-        string[] memory cmds = new string[](5);
-        if (const_overrides.length > 0) {
-            cmds = new string[](6 + const_overrides.length);
-            cmds[5] = "-c";
-
-            Constant memory cur_const;
-            for (uint256 i; i < const_overrides.length; i++) {
-                cur_const = const_overrides[i];
-                cmds[6 + i] = string.concat(cur_const.key, "=", cur_const.value);
-            }
-        }
-
-        cmds[0] = "huffc";
-        cmds[1] = string(string.concat("src/", tempFile, ".huff"));
-        cmds[2] = "-b";
-        cmds[3] = "-e";
-        cmds[4] = get_evm_version();
-
-        /// @notice compile the Huff contract and return the bytecode
-        bytecode = vm.ffi(cmds);
-
-        // Clean up temp files
-        string[] memory cleanup = new string[](2);
-        cleanup[0] = "rm";
-        cleanup[1] = string.concat("src/", tempFile, ".huff");
-
-        // set `msg.sender` for upcoming create context
-        vm.prank(deployer);
-
-        vm.ffi(cleanup);
+        return get_code(file, CodeType.CREATION);
     }
 
+    /// @notice Get the runtime bytecode of a contract
     function runtime_code(string memory file) public payable returns (bytes memory bytecode) {
+        return get_code(file, CodeType.RUNTIME);
+    }
+
+    /// @notice Get the bytecode of a contract
+    function get_code(string memory file, CodeType codeType)
+        public
+        payable
+        returns (bytes memory bytecode)
+    {
         binary_check();
 
         // Split the file into its parts
@@ -277,7 +221,11 @@ contract HuffConfig {
 
         cmds[0] = "huffc";
         cmds[1] = string(string.concat("src/", tempFile, ".huff"));
-        cmds[2] = "-r";
+        if (codeType == CodeType.CREATION) {
+            cmds[2] = "-b";
+        } else {
+            cmds[2] = "-r";
+        }
         cmds[3] = "-e";
         cmds[4] = get_evm_version();