main.cpp

// 2025 Invalid Units
// x86 Unix Just In time Compiler
// Translates a list of intermediate Opcode instructions into real x86 instructions

#include <sys/mman.h>
#include <cerrno>
#include <cstring>
#include <iostream>
#include <csignal>
#include <csetjmp>
#include <cstdint>
#include <fstream>
#include <functional>
extern "C" {
    #include <unistd.h>
}


#include <string>
#include <sstream>
#include <cstdlib>
#include <exception>
#include <bit>
#include <bitset>
#include <iomanip>

constexpr bool is_big_endian(void) {
    union {
        uint32_t i;
        char c[4];
    } bint = {0x01020304};

    return bint.c[0] == 1;
};

void local_to_network(char *val, size_t size) {
    if (!is_big_endian()) return;

    size_t l = 0;
    size_t r = size - 1;
    while (l < r) {
        char temp = val[l];
        val[l] = val[r];
        val[r] = temp; 

        l++;
        r--;
    };
}

void network_to_local(char *val, size_t size) {
    local_to_network(val, size);
}


enum class Opcode: unsigned char {

    __BEG_ARITH = 0,
    ADD,
    SUB,
    MUL,
    IMUL, // multiply signed
    DIV,
    IDIV,
    MOD,
    __END_ARITH,


    PA,
    __BEG_IMM,
    IMM8,
    IMM16,
    IMM32,
    IMM64,
    __END_IMM,


    RET
};



#ifdef _WIN32
#define __JIT_CALL __cdecl
#endif

#ifdef __GNUC__
#define __JIT_CALL
#endif



struct stackframe_t {
    stackframe_t *behind = nullptr;

    enum {
        TYPE_UNDEFINED,
        TYPE_STACK,
        TYPE_REGISTER,
        TYPE_IMM8_16_32,
        TYPE_STACK_PTR,
    } type = TYPE_UNDEFINED;
    size_t value_size;

    union {
        char reg = 0;
        char immidiete_value[4];
    };
};

int compile_instruction(const char *&instruction, const char * const end, 
                        struct stackframe_t * &stack, std::ostream& out, std::ostream& error_out
                        ) {
    Opcode opcode = Opcode(instruction[0]);

    // Because these registers may be overidden by math instructions we cannot risk them being modified
    char push_override_registers[] = {
        0,
        2,
        4 + 16 // ah
    };

    char general_registers[] = {
        // ax, al, ah are used for math operations. 
        // 0, 
        
        1, // cx
        
        // TODO: preserve bx so it can be used.
        // bx, can be used but must be preserved.
        // 2,


        // dx is used for arithmetic. 
        //3, // dx,

        8, // r8
        9, // r9
        10, // r10
        11, // r11
        12, // r12
        13, // r13
        14, // r14
        15, // r15


        /* 16, */ // ah
    };
    
    std::bitset<16> registers_in_use;
    const auto isregisteravailable = [&](char reg) -> int {
        if (registers_in_use[reg % 16]) return false;


        bool register_in_use = false;
        struct stackframe_t *top_stack = stack;
        while (top_stack != nullptr) {
            
            if (top_stack->type == stackframe_t::TYPE_REGISTER && (top_stack->reg % 16) == (reg % 16)) {
                register_in_use = true;
                break;
            }

            top_stack = top_stack->behind;
        }

        return !register_in_use;
    };
    
    const auto getavailableregister = [&](char preferred_register) -> int {
        if (preferred_register != -1) 
            if (isregisteravailable(preferred_register)) return preferred_register;

        for (int i = 0; i < sizeof(general_registers)/sizeof(char); i++) {
            int reg = general_registers[i];
            if (reg == preferred_register) continue; // we already checked it.

            if (isregisteravailable(reg)) return reg;
        };
        return -1;
    };





    const int opcode_int = (int)opcode;


    const auto push = [&]() -> stackframe_t* {
        stackframe_t *top = new stackframe_t;
        top->behind = stack;
        stack = top;
        return top;
    };

    const auto pop = [&]() -> stackframe_t* {
        if (stack == nullptr) {
            return nullptr;
        }

        stackframe_t *top = stack;
        stack = stack->behind;
        return top; // Caller is responsible for deleting the returned stackframe.
    };

    const auto writerex = [&](unsigned char src, unsigned char dst, size_t size) {
        unsigned char rex = 0;

        if (size == 1) {
            bool a = src >= 4 && src <= 7;
            bool b = dst >= 4 && dst <= 7;

            if (a || b) rex = rex | 0b01000000;
            if ((src > 15 && b) || (dst > 15 && a)) {
                error_out << "Warning: you cannot encode both ah-dh and spl-dil at the same time." << std::endl;       
            }
        }

        if (src >= 8 && src <= 15) rex = rex | 0b01000100; // REX prefix for r8-r15 src registers
        if (dst >= 8 && dst <= 15) rex = rex | 0b01000001; // REX prefix for r8-r15 dst registers
        if (size == 8) rex = rex | 0b01001000;  // REX prefix for 64-bit operand size
        if (rex != 0) out.put((char)rex); 
    };

    const auto writemodrm = [&out](unsigned char reg,  unsigned char rm) {
        reg = reg % 8;
        rm = rm % 8;

        // 2 bits (registermode), 3 bits reg, 3 bit reg2e
        unsigned char modrm = 0b11000000 | ((reg << 3) & 0b111000) | (rm & 0b111); // mod=11 (register mode), reg=dest, rm=src
        out.put((char)modrm);
    };


    const auto move_reg = [&](unsigned char dst,  unsigned char src, size_t size) {
        if (size == 2) out.put((char)0x66);  // prefix for 16 bit instructions.
        writerex(src, dst, size);
        switch (size) {
            case 1:  // 8-bit move (byte)
                out.put((char)0x88);  // mov r/m8, r8
                break;
            case 2:  // 16-bit move (word)
            case 4:  // 32-bit move (dword)
                out.put((char)0x89);  // mov r/m32, r32
                break;
            case 8:  // 64-bit move (qword)
                out.put((char)0x89);  // mov r/m64, r64
                break;
            default:
                error_out << "Invalid size " << size << " when generating mov instruction" << std::endl;
        }
        writemodrm(src, dst);
    };

    const auto put_imm_into_reg = [&](void *value, unsigned char reg, size_t size) {
        unsigned char offset = (size == 1)? 0xb0 : 0xb8;
        if (size == 2) out.put((char)0x66);  // prefix for 16 bit instructions.
        writerex(0, reg, size);
        out.put((char)offset + reg % 8); // mov r8, imm8
        out.write((char*)value, size);
    };

    const auto pop_reg = [&](unsigned char &reg, bool dont_move) -> stackframe_t* {
        stackframe_t *top = pop();
        if (top == nullptr) return top;

        if (top->type == stackframe_t::TYPE_STACK) {
            if (reg > 8) out.put((char)0x41); // rex for r8-r15
            out.put((char)(0x58 + reg % 8));
        } else if (top->type == stackframe_t::TYPE_REGISTER) {
            if (dont_move) {
                reg = top->reg;
            } else if (reg != top->reg) {
                move_reg(reg, top->reg, top->value_size);
                top->reg = reg;
            }
        } else if (top->type == stackframe_t::TYPE_IMM8_16_32) {
            put_imm_into_reg(top->immidiete_value, reg, top->value_size);
        } else {
            // we can't pop into the register.
            // return stackframe to it's rightful place.
            top->behind = stack;
            stack = top;
            return nullptr;
        }

        // DO NOT access back after popping  
        return top; // top should be freed after deleteion;
    };
    


    // cant use auto because of recursion :'(
    const std::function<stackframe_t*(unsigned char reg, size_t size, bool dont_keep_reg)> push_reg = [&](unsigned char reg, size_t size, bool dont_keep_reg) -> stackframe_t* {
        if (stack != nullptr)
            if (stack->type == stackframe_t::TYPE_REGISTER) {
                for (int i = 0; i < sizeof(push_override_registers)/sizeof(*push_override_registers); i++) {
                    // eax and edx are overwriten by division. 
                    // If we push another, we are likely going to do a division.
                    // in that case, just reset the register.

                    // If we do two division in a row?? IDK... more testing is required.
                    if (push_override_registers[i] == stack->reg) {
                        unsigned char unkept_reg = stack->reg;
                        auto a = pop_reg(unkept_reg, true);
                        push_reg(unkept_reg, a->value_size, true);
                        delete a; 
                    }
                }
        }


        stackframe_t *ret = push();
        ret->value_size = size;

         // try to use registers instead.
        if (int new_reg = getavailableregister(dont_keep_reg? -1 : reg); new_reg != -1) {
            stack->type = stackframe_t::TYPE_REGISTER;
            stack->reg = new_reg;
            // move value into new_reg.
            if (reg != new_reg)
                move_reg(stack->reg, reg, size);
            return ret;
        }  
        
        // push to the execution stack.
        stack->type = stackframe_t::TYPE_STACK;
        stack->reg = 0;
        if (reg > 8) out.put((char)0x41);
        out.put((char)(0x50 + reg % 8));
        return stack;
    };


    if (opcode > Opcode::__BEG_IMM && opcode < Opcode::__END_IMM) {
        int err = 0;
        int size = 0;
        switch (opcode) {
            case Opcode::IMM8: size = 1; break;
            case Opcode::IMM16: size = 2; break;
            case Opcode::IMM32: size = 4; break;
            case Opcode::IMM64: size = 8; break;
            default:
                error_out << "Unimplemented immediate size." << std::endl;
                return 1;
        }
        
        instruction += 1 + size;
        if (instruction >= end) {
            error_out << "Ill-formed instruction" << std::endl;
            return 1;
        }
        char *data = new char[size];
        memcpy(data, instruction - size, size);
        local_to_network(data, size);
        
        if (size == 8) {
            if (int new_reg = getavailableregister(-1); new_reg != -1) {
                push_reg(new_reg, 8, false);
                put_imm_into_reg(data, stack->reg, 8);
            } else {
                push_reg(0, 8, false);
                put_imm_into_reg(data, 0, 8); // mov, rax, imm64
            }
        } else if (size == 4 || size == 2 || size == 1) {
            push();
            stack->value_size = size;
            stack->type = stackframe_t::TYPE_IMM8_16_32;
            memcpy(stack->immidiete_value, data, size);
        } else {
            error_out << "Unimplemented immediate size." << std::endl;
            err = 1;
        }

        delete[] data;
        return err;
    
    }


    if (opcode == Opcode::PA) {
        instruction += 2;
        if (instruction >= end) {
            error_out << "Ill-formed instruction" << std::endl;
            return 1;
        }
        unsigned char argument = instruction[-1];
        // assuming all arguments are integer(ish) 
        unsigned char integer_param_reg[] = {
            6,
            7,
            8,
            9,
        };

        constexpr size_t integer_param_reg_size = (sizeof(integer_param_reg)/sizeof(char));

        if (argument >= integer_param_reg_size) {
            error_out << "Ill-formed instruction" << std::endl;
            return 1;
        }


        // assume extended integers for now. more metadata about the argument is needed.
        push_reg(integer_param_reg[argument], 1, false);
        return 0;
    }
                            
    if (opcode_int > int(Opcode::__BEG_ARITH) && opcode_int < int(Opcode::__END_ARITH)) {
        instruction += 1;


        int err = 0;
        bool handle_as_immidiete = false;
        bool multiplicationlike = opcode_int >= int(Opcode::MUL) && opcode_int <= int(Opcode::MOD);
        bool signed_multiply = opcode_int == int(Opcode::IMUL) || opcode_int == int(Opcode::IDIV);

        unsigned char instruction_offset = 0;
        unsigned char instruction_extension = 0;

        /*
        unsigned char rega = getavailableregister(multiplicationlike? -1 : 0);
        auto a = pop_reg(rega, true); // if it's multiplication like. our dividend must be in eax. 
        registers_in_use[rega % 16] = true;

        unsigned char regb = multiplicationlike? 0 : getavailableregister(-1);
        stackframe_t *b = nullptr;
        if (stack->type == stackframe_t::TYPE_IMM8_16_32 && rega == 0 && !multiplicationlike) {
            b = pop();
            handle_as_immidiete = true;
        } else {
            b = pop_reg(regb, !multiplicationlike);
            handle_as_immidiete = false;
        }
        */

        stackframe_t *rhs = nullptr;
        unsigned char rhs_reg = 255; // -1 indicates that we don't have a register for it and it's an immidiete.
        if (stack->type == stackframe_t::TYPE_IMM8_16_32 && !multiplicationlike) {
            rhs = pop();
        } else {
            rhs_reg = getavailableregister(-1); 
            rhs = pop_reg(rhs_reg, !multiplicationlike);
        }

        unsigned char lhs_reg = multiplicationlike? 0 : getavailableregister(0);
        auto lhs = pop_reg(lhs_reg, !multiplicationlike);

        if (lhs != nullptr)
        if (lhs_reg != 0 && rhs->type == stackframe_t::TYPE_IMM8_16_32) {
            rhs_reg = getavailableregister(-1);
            put_imm_into_reg(rhs->immidiete_value, rhs_reg, rhs->value_size);
            rhs->reg = rhs_reg;
            rhs->type = stackframe_t::TYPE_REGISTER;
        }


        size_t size = 0; 
        if (lhs == nullptr || rhs == nullptr) {
            error_out << "Ill-formed arithmatic. Missing or invalid Opperands" << std::endl;
            err = 1;
            goto arith_end;
        }

        if (lhs->value_size != rhs->value_size) {
            error_out << "Opperands must be of the same size." << std::endl;
            err = 1;
            goto arith_end;
        }

        size = lhs->value_size;
        if (size == 8 || size == 4 || size == 2 || size == 1) ;
        else {
            error_out << "Bad type size.";
            err = 1;
            goto arith_end;
        }


        // during multiplicationlike instructions
        // if we are using dx as an operand. move it to a different register. dx is overriden.
        // a is always ax so we don't have to consider it.
        // we can skip both of these steps if the size is one.

        // we actually don't care if it's overridden (because we are popping it off.)
        // it's just that the multiplication like 

        if (((rhs_reg == 2 && size != 1) || (rhs_reg == 4 && size == 1)) && multiplicationlike) {
            if (int reg = getavailableregister(-1); reg != -1) {
                move_reg(reg, rhs_reg, rhs->value_size);
                rhs->reg = reg;
                rhs_reg = reg;
            } else {
                error_out << "Couldn't find a valid register for register b" << std::endl;
                err = 2;
                goto arith_end;
            }
        }

        // if (multiplicationlike) {
        //     if (signed_multiply) {
        //         if (size == 2) out.put((char)0x66);  // prefix for 16 bit instructions.
        //         writerex(0, 0, size);
        //         out.put(size == 1? (char)0x98 : 0x99); // cbw or cdq 
        //     } else {
        //         if (size != 1) {
        //             // erase edx
        //             writerex(push_override_registers[1], push_override_registers[1], size);
        //             if (size == 2) out.put((char)0x66);  // prefix for 16 bit instructions.  
        //             out.put((char)0x31); // xor
        //             writemodrm(push_override_registers[1], push_override_registers[1]);
        //         } else {
        //             // erase ah
        //             writerex(push_override_registers[2], push_override_registers[2], size);
        //             out.put((char)0x30); // xor
        //             writemodrm(push_override_registers[2], push_override_registers[2]);
        //         }
        //     }
        // }

        
        switch (opcode) {
            case Opcode::ADD: instruction_offset = 0x1; break;
            case Opcode::SUB: instruction_offset = 0x29; break;
            case Opcode::MUL: instruction_offset = 0xf7; instruction_extension = 4; break;
            case Opcode::IMUL: instruction_offset = 0xf7; instruction_extension = 5; break;

            case Opcode::MOD:
            case Opcode::DIV: instruction_offset = 0xf7; instruction_extension = 6; break;
            case Opcode::IDIV: instruction_offset = 0xf7; instruction_extension = 7; break;

            default:
                error_out << "Unimplemented arithmetic. " << std::endl;
                err = 2;
                goto arith_end;
        }

        if (size == 1 ) {
            // subtract 1 to get 8bit versions. 
            // weirdly enough this applies for all of the arithmetic instructions. ¯\_(ツ)_/¯
            instruction_offset -= 1;
        }

        
        if (size == 2) out.put((char)0x66);  // prefix for 16 bit instructions.
        
        writerex(rhs_reg, lhs_reg, size);


        
        if (rhs_reg == 255) {
            instruction_offset += 4;
        }

        out.put((char)(instruction_offset));

        if (rhs_reg == 255) {
            out.write(rhs->immidiete_value, rhs->value_size);
        } else if (multiplicationlike) {
            writemodrm(instruction_extension, rhs_reg);
        } else {
            writemodrm(lhs_reg, rhs_reg);
        }
        
        
        
        if (opcode == Opcode::MOD) {
            if (size == 1) {
                push_reg(push_override_registers[2], size, false); // push ah as it contains our dividor.
            } else {
                push_reg(push_override_registers[1], size, false); // push dx as it contains our dividor.
            }
        } else {
            push_reg(lhs_reg, size, false);
        }

        arith_end:
            if (lhs != nullptr) delete lhs; 
            if (rhs != nullptr) delete rhs;
            return err;
    }


    if (opcode == Opcode::RET) {
        instruction += 1;

        unsigned char eax = 0;
        auto stack_frame = pop_reg(eax, false); // pop into eax
        delete stack_frame;
        out.put(0xc3); // ret48 b9 02 00 00 00 00 00 00 00 49 b8 04 00 00 00 00 00 00 00 48 89 c1 48 31 d2 49 f7 f0 c3 
        return 0;
    }


    error_out << "Ill-formed or unexpected instruction" << std::endl;
    return 1;
                           
}

int compile_jit(const char *begin, size_t size, std::ostream& out, std::ostream& error_out) {
    const char *end = begin + size;
    struct stackframe_t *stack = nullptr;
    int err = 0;

    while (intptr_t(begin) < intptr_t(end) && err == 0) { 
        err = compile_instruction(begin, end, stack, out, error_out); 
    }
    if (err != 0) {
        error_out << "Failed to compile opcodes." << std::endl;
    }

    stackframe_t *top = stack;
    while (top != nullptr) {
        err = 1;
        stackframe_t *to_delete = top;
        top = top->behind;
        delete to_delete; 
    }

    return err;
}


#include <functional>
static jmp_buf exec_jmp;
static std::ostream *current_error_out = nullptr;
template<typename ReturnType, typename... Args> int execute (void* executable, Args... args,
        std::function<void(ReturnType)> return_callback = nullptr, std::ostream *error_out = &std::cerr) {
    current_error_out = error_out;
    auto sighandler = [] (int sig) {
        if (current_error_out != nullptr) {
            (*current_error_out) << "Executing mapped memory raised signal " << strsignal(sig) << std::endl; 
            (*current_error_out) << " (" << sig << ")" << std::endl;
        }

        longjmp(exec_jmp, 1);
    };

    constexpr int possibly_raised_sigs[] = {
        SIGSEGV,
        SIGILL,
        SIGFPE
    };

    for (int i = 0; i < sizeof(possibly_raised_sigs)/sizeof(int); i++) 
        signal(possibly_raised_sigs[i], sighandler);

    if (int i = setjmp(exec_jmp); i == 0) {
        try {
            using FuncType = ReturnType (__JIT_CALL *)(Args...);
            FuncType func = reinterpret_cast<FuncType>(executable);
            if (return_callback == nullptr) {
                func(std::forward<Args>(args)...);
            } else return_callback(std::forward<ReturnType>(func(std::forward<Args>(args)...)));
        } catch (const std::exception& e) {
            if (current_error_out != nullptr) {
                (*current_error_out)  << "C++ Exception during JIT execution: " << e.what() << std::endl;
            }
        } catch (...) {
            if (current_error_out != nullptr) {
                (*current_error_out)  << "Unknown exception during JIT execution. " << std::endl;
            }
        };

        longjmp(exec_jmp, 2); // The function may have manipulated our stack, so we should longjump to recover it.
    } else {
        for (int i = 0; i < sizeof(possibly_raised_sigs)/sizeof(int); i++) 
            signal(possibly_raised_sigs[i], SIG_DFL);


         // i == 1 if error i == 2 if success
         if (i != 2) {
            return i;
        }

        return 0;
    }

    // unreachable
    return 1;
}




#define ERRNO_ASSERT(a) if (!a) {  \
    std::cout << "Failed to run executable memory.'" << std::endl; \
    std::cout << std::strerror(errno) << " (" << errno << ")" << std::endl; \
    return 1; \
}



int main() {

    // Compiling the program.
    
    std::stringstream ss;

    const char my_prog[] = {char(Opcode::IMM16), char(25), char(0), char(Opcode::IMM16), char(0), char(0), char(Opcode::IMUL), char(Opcode::RET)}; 
    if (compile_jit(my_prog, sizeof(my_prog)/sizeof(char), ss, std::cout) != 0) {
        std::cout << "Failed to compile program JIT." << std::endl;
        return 1;
    }

    std::string mystr = ss.str();
    size_t size = mystr.size();

    std::ofstream outfile("myfile", std::ios_base::binary);
    outfile << ss.rdbuf();
    outfile.close();

    for (unsigned char c : mystr) {
        std::cout << std::hex << std::setw(2) << std::setfill('0') << (int)c << " ";
    }
    std::cout << std::endl;
    

    // Preparing execution memory
    int err = 0;
    void *execmemory = nullptr;

    execmemory = mmap(nullptr, size, PROT_WRITE | PROT_READ, MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
    ERRNO_ASSERT((execmemory != (void*)MAP_FAILED))

    err = mlock(execmemory, size); // prevent memory from being paged during execution.
    ERRNO_ASSERT(err == 0)
    memset(execmemory, 0, size);
    memcpy(execmemory, mystr.data(), size);
    if (memcmp(execmemory, mystr.data(), size) != 0) {
        return 1;
    }


    // Prepare for execution.
    err = mprotect(execmemory, size, PROT_EXEC);
    ERRNO_ASSERT(err == 0)

    // Executing
    ERRNO_ASSERT(err == 0)
    err = execute<int16_t, int16_t, int16_t>(execmemory, 2, 4, [](int16_t ret) {
        std::cout << "JIT execution returened: " << std::dec << int(ret) << std::endl;
    });

    if (err != 0)
    if (current_error_out != nullptr) {
        (*current_error_out) << "Failed to execute JIT code, Error code: " << err << std::endl;
    }


        
    err = munlock(execmemory, size);
    ERRNO_ASSERT(err == 0);
    ERRNO_ASSERT(err == 0);

    err = munmap(execmemory, size); // free memory
    ERRNO_ASSERT(err == 0);

    #undef ERRNO_ASSERT
}