brainfuck.cpp

#include <algorithm>
#include <array>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <sstream>
#include <stack>
#include <tuple>
#include <unordered_map>
#include <unordered_set>
#include <vector>

constexpr std::size_t MEMORY_SIZE {30000};

constexpr const char* TRANSPILED_C_HEADERS = R"(// Autogenerated, compile with: "gcc <file.c> -O2 -o <file.out>"
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>

#define MEMORY_SIZE 30000
#define MOD(a, b) (((a) % (b) + (b)) % (b))

uint8_t memory[MEMORY_SIZE] = {0};
size_t ptr = 0;

static inline void  CLEAR() { memory[ptr] = 0; }
static inline void OUTPUT() { putchar(memory[ptr]); }
static inline void  INPUT() { int ch = getchar(); if (ch != EOF) memory[ptr] = (uint8_t) ch; }

void SHIFTPTR(long shift) { 
    long ptr_ = (long) ptr + shift;
    if (ptr_ >= MEMORY_SIZE) ptr_ -= MEMORY_SIZE;
    else if (ptr_ < 0) ptr_ += MEMORY_SIZE;
    ptr = (size_t) ptr_;
}

void UPDATEVAL(long shift) { 
    long val = (long) memory[ptr] + shift;
    if (val > 255) val -= 256;
    else if (val < 0) val += 256;
    memory[ptr] = (uint8_t) val;
}

void SHIFTPTR_ZERO(long shift) { 
    while(memory[ptr] != 0) {
        long ptr_ = (long) ptr + shift;
        if (ptr_ >= MEMORY_SIZE) ptr_ -= MEMORY_SIZE;
        else if (ptr_ < 0) ptr_ += MEMORY_SIZE;
        ptr = (size_t) ptr_;
    }
}

void UPDATE_BY_CURR(long shift, long factor) { 
    long ptr_ = (long) ptr + shift;
    if (ptr_ >= MEMORY_SIZE) ptr_ -= MEMORY_SIZE;
    else if (ptr_ < 0) ptr_ += MEMORY_SIZE;
    size_t copyTo = (size_t) ptr_;
    memory[copyTo] = (uint8_t) MOD((long) memory[copyTo] + factor * memory[ptr], 255);
}

int main() {
    // Inserted dynamically from compiled brainfuck instructions
)";

constexpr const char* TRANSPILED_C_FOOTERS = R"(
    return 0;
}
)";

constexpr const char* HELP_MESSAGE = R"(
Brainfuck CLI Tool - Brainfuck interpreter, compiler & transpiler.

Usage:
    brainfuck                  Spawns an interactive Brainfuck shell.
    brainfuck          <file>  Compile a .bf file into bytecode (.bfc).
    brainfuck <OPTION> <file>  Perform a specific operation based on the provided option.

Options:
    -p, --profile              Profile raw .bf code, execute without compiling or transpiling.
    -t, --transpile            Transpile .bfc bytecode to a C source file (.bfc.c).
    -e, --exec                 Execute compiled .bfc code directly.
    -h, --help                 Display this help message.

Examples:
    brainfuck    test.bf       Compiles test.bf into test.bf.bfc.
    brainfuck    test.bfc      Error: Already compiled. Use -t or -e instead.
    brainfuck -t test.bfc      Transpiles test.bfc to test.bfc.c.
    brainfuck -e test.bfc      Execute the compiled bytecode.
    brainfuck -p test.bf       Profile and execute raw Brainfuck source code.

Behavior:
    - All flags except for the ones above are ignored.
    - Input files are validated by their extensions:
        * .bf  : Source code files to be compiled.
        * .bfc : Compiled bytecode for execution or transpilation.

Note:
    Flags cannot be combined. Use one mode (profile, transpile, or execute) at a time.

)";

// Helper to check if 'arg' is present in args vector
bool hasOpt(const std::string &arg, std::vector<std::string> &args) {
    return std::find(args.begin(), args.end(), arg) != args.end();
}

// Helper to exit with error message to console
void exitWithError(const std::string &message) {
    std::cerr << "Error: " << message << "\n";
    std::exit(1);
}

class BrainFuck {
    private:
        static std::unordered_set<char> BRAINFUCK_CHARS;
        enum OPCODE { 
            UPDATEVAL, SHIFTPTR, OUTPUT, INPUT, LSTART, LEND, 
            CLEAR, SHIFTPTR_ZERO, UPDATE_BY_CURR
        };

        using INSTRUCTION_T = std::tuple<OPCODE, long, long>;

        std::array<std::uint8_t, MEMORY_SIZE> memory {0};
        std::size_t ptr {0};

        struct HashPair {
            std::size_t operator() (const std::pair<std::size_t, std::size_t> &p) const {
                std::hash<std::size_t> Hash;
                return Hash(p.first) ^ (Hash(p.second) << 1);
            }
        };

        template <typename T>
        inline static T mod(T val, T modulus) {
            return (val % modulus + modulus) % modulus;
        }

        // Validate string as BF and print an error message at point of error (if any)
        static bool validateCode(const std::string &code, bool showErr = true) {
            std::size_t lineNo {1};
            std::stack<std::pair<std::size_t, std::size_t>> loopStk;
            for (std::size_t i {0}; i < code.size(); i++) {
                if (code[i] == '\n') lineNo++;
                else if (code[i] == '[')
                    loopStk.push({lineNo, i + 1});
                else if (code[i] == ']') {
                    if (!loopStk.empty()) loopStk.pop();
                    else {
                        if (showErr) {
                            std::cerr << "Error: Unexpected closing bracket in line " 
                                      << lineNo << " char " << i + 1 << "\n";
                        }

                        return false;
                    }
                }
            }

            if (!loopStk.empty() && showErr) {
                std::cerr << "Error: Unclosed bracket in line " 
                          << loopStk.top().first << " char " << loopStk.top().second << "\n";
            }

            return loopStk.empty();
        }

        // Read file as raw text
        static std::string readTextFile(const std::string &fileName) {
            std::ifstream ifs {fileName};
            if (!ifs) exitWithError("Unable to open file for reading raw: " + fileName);
            std::ostringstream oss; 
            oss << ifs.rdbuf();
            return oss.str();
        }

        // Read file as compiled bytecode
        static std::vector<INSTRUCTION_T> readByteCodeFile(const std::string &fileName) {
            std::ifstream ifs {fileName, std::ios::binary};
            if (!ifs) exitWithError("Unable to open binary file for reading: " + fileName);

            // Read the headers
            char header[9] {};
            std::size_t instructionLen;
            ifs.read(reinterpret_cast<char *>(&header), 9);
            ifs.read(reinterpret_cast<char *>(&instructionLen), sizeof(std::size_t));

            // Check header match
            if (std::strcmp(header, "BRAINFUCK") != 0) 
                exitWithError("Bytecode supplied is corrupted. Please recompile.");

            std::vector<INSTRUCTION_T> instructions;
            instructions.reserve(instructionLen);
            INSTRUCTION_T instruction;
            while (ifs.read(reinterpret_cast<char*>(&instruction), sizeof(instruction)))
                instructions.emplace_back(instruction);

            // Check if size still matches
            if (instructions.size() != instructionLen) exitWithError("Bytecode supplied is corrupted. Please recompile.");

            return instructions;
        }

        // Extract substr excluding comments, compress if applicable
        static std::string getLoopRepr(std::size_t start, std::size_t end, const std::string& code) {
            std::ostringstream oss;
            std::pair<char, std::size_t> curr{'\0', 0};
            for (std::size_t pos {start}; pos <= end; pos++) {
                char ch {code.at(pos)}; 
                if (BRAINFUCK_CHARS.find(ch) == BRAINFUCK_CHARS.end()) 
                    continue;
                else if (curr.first == ch && ch != '[' && ch != ']' && ch != '.' && ch != ',') 
                    curr.second++;
                else {
                    if (curr.first != '\0') oss << curr.first << (curr.second > 1? std::to_string(curr.second): "");
                    curr = {ch, 1};
                } 
            }

            oss << curr.first << (curr.second > 1? std::to_string(curr.second): "");
            return oss.str();
        }

        /*
         * - Check if starting and ending points are the same
         * - Loop is non nested without any compound op codes
         * - We decrement startPos memory by 1 
         * - For eg: [<3+2<1+>4-]
         */
        static bool validateSimpleDistributionLoop(
            std::size_t start, std::size_t end, 
            std::vector<INSTRUCTION_T> &instructions,
            std::unordered_map<long, long> &distributeTo
        ) {
            // Keep track of curr relative pos
            long shift {0};
            for (std::size_t pos {start + 1}; pos < end; pos++) {
                OPCODE opcode; long val1, val2; 
                std::tie(opcode, val1, val2) = instructions[pos];

                // Move curr pointer
                if (opcode == OPCODE::SHIFTPTR) shift += val1;

                // Update value by a factor
                else if (opcode == OPCODE::UPDATEVAL) distributeTo[shift] += val1;

                // Disallow any other opcode
                else return false;
            }

            // Check if we are back to where we start and 
            // we are doing a simple decrement
            bool status {shift == 0 && distributeTo[0] == -1};
            return status;
        }

        void _logByteCodeInstructions(std::vector<INSTRUCTION_T> &instructions, const std::string &fileName) {
            std::ofstream ofs {fileName};
            if (!ofs) exitWithError("Unable to open file for writing log: " + fileName);
            for (INSTRUCTION_T &instruction: instructions)
                ofs << std::get<0>(instruction) << "," << std::get<1>(instruction) << "\n";
        }

    public:
        // Execute raw string instruction - assumes code is already validated
        void interpretCode(const std::string &code, const std::string &logFileName = "") {
            // Profiling loops with a counter
            std::unordered_map<std::pair<std::size_t, std::size_t>, std::size_t, HashPair> loopCounter;
            bool logOutput {!logFileName.empty()};

            std::size_t pos {0};
            std::stack<std::size_t> loopStk;
            while (pos < code.size()) {
                char ch  {code.at(pos)};
                switch (ch) {
                    case '+':
                        memory[ptr] = memory[ptr] == 255? 0: memory[ptr] + 1;
                        break;
                    case '-':
                        memory[ptr] = memory[ptr] == 0? 255: memory[ptr] - 1;
                        break;
                    case '>':
                        ptr = ptr == MEMORY_SIZE - 1? 0: ptr + 1;
                        break;
                    case '<':
                        ptr = ptr == 0? MEMORY_SIZE - 1: ptr - 1;
                        break;
                    case '.':
                        std::putchar(memory[ptr]);
                        break;
                    case ',':
                        memory[ptr] = static_cast<std::uint8_t>(std::getchar());
                        break;
                    case '[':
                        if (memory[ptr] != 0) loopStk.push(pos);
                        else {
                            int openBrackets {1};
                            while (openBrackets > 0) {
                                ++pos;
                                if (code.at(pos) == ']') openBrackets--;
                                else if (code.at(pos) == '[') openBrackets++;
                            }
                        }
                        break;
                    case ']':
                        if (logOutput) loopCounter[{loopStk.top(), pos}]++;
                        if (memory[ptr] == 0) loopStk.pop();
                        else pos = loopStk.top();
                        break;
                }

                // Go to the next instruction
                pos++;
            }

            // Output loop profiling output
            if (logOutput) {
                // Get the loops string repr into a map
                std::unordered_map<std::string, std::size_t> loopCounterStrRep;
                for (const std::pair<const std::pair<std::size_t, std::size_t>, std::size_t> &kv: loopCounter)
                    loopCounterStrRep[getLoopRepr(kv.first.first, kv.first.second, code)] += kv.second;

                // Write to a vector for sorting
                std::vector<std::pair<std::string, std::size_t>> loopCounterVec;
                for (const std::pair<const std::string, std::size_t> &kv: loopCounterStrRep)
                    loopCounterVec.emplace_back(kv.first, kv.second);
                
                // Sort the vector
                std::sort(loopCounterVec.begin(), loopCounterVec.end(), []
                    (std::pair<std::string, std::size_t> &v1, std::pair<std::string, std::size_t> &v2) {
                        return v1.second > v2.second;
                    }
                );

                // Write the profile output
                std::ofstream log {logFileName};
                if (!log) exitWithError("Unable to open file for logging: " + logFileName);
                for (std::pair<std::string, std::size_t> &v: loopCounterVec)
                    log << v.first << "," << v.second << "\n";
            }
        }

        // Execute bytecode (binary file, extension is not explicity checked)
        void executeByteCode(const std::string &fileName) {
            // Read file as bytecode
            std::vector<INSTRUCTION_T> instructions {readByteCodeFile(fileName)};

            // Execute bytecode, profile loops if logFileName is set
            OPCODE opcode; long val1, val2;
            for (std::size_t pos {0}; pos < instructions.size(); pos++) {
                std::tie(opcode, val1, val2) = instructions[pos];
                switch (opcode) {

                    case OPCODE::UPDATEVAL:
                        memory[ptr] = static_cast<std::uint8_t>(mod(static_cast<long>(memory[ptr]) + val1, 255l));
                        break;

                    case OPCODE::SHIFTPTR:
                        ptr = static_cast<std::size_t>(mod(static_cast<long>(ptr) + val1, static_cast<long>(MEMORY_SIZE)));
                        break;

                    case OPCODE::INPUT:
                        memory[ptr] = static_cast<std::uint8_t>(std::getchar());
                        break;

                    case OPCODE::OUTPUT:
                        std::putchar(memory[ptr]);
                        break;

                    case OPCODE::LSTART:
                        pos = memory[ptr] == 0? static_cast<std::size_t>(val1): pos;
                        break;

                    case OPCODE::LEND:
                        pos = memory[ptr] == 0? pos: static_cast<std::size_t>(val1);
                        break;

                    case OPCODE::CLEAR:
                        memory[ptr]  = 0;
                        break;

                    case OPCODE::SHIFTPTR_ZERO:
                        while(memory[ptr] != 0) 
                            ptr = static_cast<std::size_t>(mod(static_cast<long>(ptr) + val1, static_cast<long>(MEMORY_SIZE)));
                        break;

                    case OPCODE::UPDATE_BY_CURR: {
                        // Where the destination position is
                        std::size_t copyTo {static_cast<std::size_t>(mod(static_cast<long>(ptr) + val1, static_cast<long>(MEMORY_SIZE)))};

                        // Updated val to write
                        long updatedVal {static_cast<long>(memory[copyTo])};
                        updatedVal += val2 * static_cast<long>(memory[ptr]);
                        updatedVal = mod(updatedVal, 255l);

                        // Update val at destination
                        memory[copyTo] = static_cast<std::uint8_t>(updatedVal);
                    } break;

                }
            }
        }

        // Compiles the raw code from file into bytecode - "c" is appended to filename (assumes file ext .bf)
        static void compileFile(const std::string &fileName) {
            // Read as raw file
            std::string code {readTextFile(fileName)};

            // Validate the code and exit if not valid
            if (!validateCode(code)) std::exit(1);

            // Process and write the instructions to bytecode
            std::stack<std::size_t> stk;
            std::vector<INSTRUCTION_T> instructions;
            for (const char ch: code) {
                if (ch == '+' || ch == '-' || ch == '>' || ch == '<') {
                    long   shiftVal {ch == '+' || ch == '>'? 1l: -1l};
                    OPCODE currOpCode {ch == '+' || ch == '-'? OPCODE::UPDATEVAL: OPCODE::SHIFTPTR};
                    if (!instructions.empty() && std::get<0>(instructions.back()) == currOpCode)
                        std::get<1>(instructions.back()) += shiftVal;
                    else
                        instructions.emplace_back(currOpCode, shiftVal, 0l);
                } else if (ch == '.' || ch == ',') {
                    OPCODE currOpCode {ch == '.'? OPCODE::OUTPUT: OPCODE::INPUT};
                    instructions.emplace_back(currOpCode, 1l, 0l);
                } else if (ch == '[') {
                    stk.push(instructions.size());
                    instructions.emplace_back(OPCODE::LSTART, -1, 0l);
                } else if (ch == ']') {
                    std::size_t loopStartPos {stk.top()}, loopEndPos {instructions.size()}; stk.pop();
                    std::get<1>(instructions[loopStartPos]) = static_cast<long>(loopEndPos);
                    instructions.emplace_back(OPCODE::LEND, static_cast<long>(loopStartPos), 0l);

                    // When we are processing simple distribution 
                    // loops, we would make use of this guy
                    std::unordered_map<long, long> distributeTo;

                    // If loop contains only UPDATEVAL, it is a clear
                    if (loopEndPos - loopStartPos == 2 && std::get<0>(instructions[loopStartPos + 1]) == OPCODE::UPDATEVAL) {
                        instructions.pop_back(); instructions.pop_back(); instructions.pop_back();
                        instructions.emplace_back(OPCODE::CLEAR, 1l, 0l);
                    }

                    // If loop contains only SHIFTPTR, it is shift until nonzero
                    else if (loopEndPos - loopStartPos == 2 && std::get<0>(instructions[loopStartPos + 1]) == OPCODE::SHIFTPTR) {
                        instructions.pop_back(); instructions.pop_back(); instructions.pop_back();
                        instructions.emplace_back(OPCODE::SHIFTPTR_ZERO, std::get<1>(instructions[loopStartPos + 1]), 0l);
                    }

                    // For eg: [<3+2<1+>4-]
                    else if (validateSimpleDistributionLoop(loopStartPos, loopEndPos, instructions, distributeTo)) {
                        // Remove all instructions inside the loop
                        while (instructions.size() > loopStartPos) 
                            instructions.pop_back();

                        // Insert INCR_BY_CURR, DECR_BY_CURR opcodes
                        for (const std::pair<const long, long> &kv: distributeTo) {
                            long shift, factor;
                            std::tie(shift, factor) = kv;
                            if (factor != 0 && shift != 0)
                                instructions.emplace_back(OPCODE::UPDATE_BY_CURR, shift, factor);
                        }

                        // Clear curr cell
                        instructions.emplace_back(OPCODE::CLEAR, 1l, 0l);
                    }

                }
            }

            // Open new binary file for writing bytecode
            std::ofstream ofs{fileName + "c", std::ios::binary};
            if (!ofs) exitWithError("Unable to open file for writing bytecode: " + fileName);

            // Write the headers
            std::size_t instructionLen {instructions.size()};
            ofs.write("BRAINFUCK", 9);
            ofs.write(reinterpret_cast<const char*>(&instructionLen), sizeof(std::size_t));

            // Write the instruction set
            for (INSTRUCTION_T &instruction: instructions) {
                OPCODE opcode {std::get<0>(instruction)};
                if (opcode == OPCODE::SHIFTPTR || opcode == OPCODE::SHIFTPTR_ZERO || opcode == OPCODE::UPDATE_BY_CURR)
                    std::get<1>(instruction) = std::get<1>(instruction) % static_cast<long>(MEMORY_SIZE);
                else if (opcode == OPCODE::UPDATEVAL)
                    std::get<1>(instruction) = std::get<1>(instruction) % 255;
                ofs.write(reinterpret_cast<const char*>(&instruction), sizeof(instruction));
            }
        }

        void transpileByteCode(const std::string &fileName) {
            // Read file as bytecode
            std::vector<INSTRUCTION_T> instructions {readByteCodeFile(fileName)};

            // Write C code to a temp string
            std::ofstream ofs {fileName + ".c"};
            ofs << TRANSPILED_C_HEADERS;

            // Insert dynamically from the instruction set
            std::size_t nestingLevel {1};
            for (const INSTRUCTION_T &instruction: instructions) {
                OPCODE opcode; long val1, val2;
                std::tie(opcode, val1, val2) = instruction;
                ofs << std::string(nestingLevel, '\t');
                switch (opcode) {
                    case OPCODE::SHIFTPTR:
                        ofs << "SHIFTPTR(" << val1 << ");";
                        break;
                    case OPCODE::UPDATEVAL:
                        ofs << "UPDATEVAL(" << val1 << ");";
                        break;
                    case OPCODE::CLEAR:
                        ofs << "CLEAR();";
                        break;
                    case OPCODE::OUTPUT:
                        ofs << "OUTPUT();";
                        break;
                    case OPCODE::INPUT:
                        ofs << "INPUT();";
                        break;
                    case OPCODE::SHIFTPTR_ZERO:
                        ofs << "SHIFTPTR_ZERO(" << val1 << ");";
                        break;
                    case OPCODE::UPDATE_BY_CURR:
                        ofs << "UPDATE_BY_CURR(" << val1 << "," << val2 << ");";
                        break;
                    case OPCODE::LSTART:
                        nestingLevel++;
                        ofs << "while (memory[ptr]) {";
                        break;
                    case OPCODE::LEND:
                        nestingLevel--;
                        ofs << "}";
                        break;
                }
                ofs << "\n";
            }

            ofs << TRANSPILED_C_FOOTERS;
        }

        void shell() {
            std::cout << "Brainfuck Interpreter. Hit Ctrl+C to exit.\n";
            bool continueLoop {true};
            while (continueLoop) {
                std::cout << "BF> ";

                // Continue reading input if line ends with '\'
                std::string code, line;
                while ((continueLoop = static_cast<bool>(std::getline(std::cin, line)))) {
                    code += line;
                    if (!line.empty() && line.back() == '\\') 
                        code.pop_back(); // remove '\' at end
                    else break;
                }

                // If Interupt or Ctrl + D is hit, don't execute code
                if (!continueLoop) std::cout << "\n";
                else if (validateCode(code, true)) 
                    interpretCode(code);
            }
        }

        void executeTextFile(const std::string &fileName) {
            std::string code {readTextFile(fileName)};
            if (validateCode(code)) interpretCode(code, fileName + ".log");
        }
};

// Init static variables
std::unordered_set<char> BrainFuck::BRAINFUCK_CHARS {'+', '-', '>', '<', ',', '.', '[', ']'};

int main(int argc, char **argv) {
    BrainFuck bf;
    std::vector<std::string> args {argv, argv + argc};

    if (hasOpt("-h", args) || hasOpt("--help", args))
        std::cout << HELP_MESSAGE;

    else if (args.size() == 1)
        bf.shell();        

    else {
        std::string fileName {args.back()};
        bool profileFlag   {hasOpt("-p", args) || hasOpt("--profile",   args)};
        bool transpileFlag {hasOpt("-t", args) || hasOpt("--transpile", args)};
        bool execFlag      {hasOpt("-e", args) || hasOpt("--exec", args)};
        std::string ext {std::filesystem::path(fileName).extension()};

        // Fail on incorrect flag combination or missing file
        if (!std::filesystem::exists(fileName) || !std::filesystem::is_regular_file(fileName))
            exitWithError("No such file: " + fileName);
        else if (profileFlag && transpileFlag)
            exitWithError("Cannot transpile & profile at the same time.");
        else if (transpileFlag && execFlag)
            exitWithError("Can only Transpile OR Execute the bytecode.");
        else if (ext != ".bfc" && ext != ".bf")
            exitWithError("Unsupported extension: " + (ext.empty()? "*EMPTY*": ext));

        // Profile raw string by executing it
        else if (profileFlag) {
            if (ext == ".bfc") exitWithError("Cannot profile a bytecode.");   
            std::cout << "Code will be auto executed. Check: " << fileName << ".log" << ".\n";
            bf.executeTextFile(fileName);
        } 

        // Transpile or execute compiled bytecode
        else if (transpileFlag || execFlag) {
            if (ext != ".bfc") exitWithError("Must be compiled before transpiling or executing it.\nTry: `brainfuck <file.bf>`.");
            else if (execFlag) bf.executeByteCode(fileName);
            else bf.transpileByteCode(fileName);
        } 

        // No options set, compile it
        else {
            if (ext != ".bfc") bf.compileFile(fileName);
            else exitWithError("Already compiled. Use -t or -e instead.");
        }
    }

    return 0;
}