jit.cpp

#include <AK/DistinctNumeric.h>
#include <AK/NonnullOwnPtr.h>
#include <AK/Vector.h>
#include <LibMain/Main.h>
#include <sys/mman.h>

// JS::Bytecode::Executable (jitme)
// 1:
//[   0] Store $5
//[  20] LoadImmediate 0
//[  40] SetLocal 0
//[  60] Load $5
//[  80] LoadImmediate undefined
//[  a0] Store $6
//[  c0] Jump @4
// 2:
// 3:
//[   0] LoadImmediate undefined
//[  20] Jump @5
// 4:
//[   0] GetLocal 0
//[  20] Store $7
//[  40] LoadImmediate 100000000
//[  60] LessThan $7
//[  80] JumpConditional true:@3 false:@6
// 5:
//[   0] Store $6
//[  20] GetLocal 0
//[  40] Increment
//[  58] SetLocal 0
//[  78] Jump @4
// 6:
//[   0] Load $6
//[  20] Jump @2

struct Value {
    u64 value { 0 };
};

AK_TYPEDEF_DISTINCT_NUMERIC_GENERAL(size_t, VMRegister);
AK_TYPEDEF_DISTINCT_NUMERIC_GENERAL(size_t, VMLocal);

struct Instruction {
    enum class Type {
        LoadImmediate,
        Load,
        Store,
        SetLocal,
        GetLocal,
        Increment,
        LessThan,
        Jump,
        JumpConditional,
        Exit,
    };

    Type type {};

protected:
    explicit Instruction(Type type)
        : type(type)
    {
    }
};

struct BasicBlock {
    Vector<NonnullOwnPtr<Instruction>> instructions;

    // Offset into the instruction stream where this code block starts.
    size_t offset { 0 };

    // Offsets into the instruction stream where we have RIP-relative jump offsets to here that need patching.
    Vector<size_t> jumps_to_here;

    template<typename T, typename... Args>
    void append(Args&&... args)
    {
        instructions.empend(make<T>(forward<Args>(args)...));
    }

    void dump() const
    {
        for (size_t i = 0; i < instructions.size(); ++i) {
            dbgln("    [{}] {}", i, to_underlying(instructions[i]->type));
        }
    }
};

struct LoadImmediate : public Instruction {
    LoadImmediate(Value value)
        : Instruction(Type::LoadImmediate)
        , value(value)
    {
    }

    Value value;
};

struct Store : public Instruction {
    Store(VMRegister reg)
        : Instruction(Type::Store)
        , reg(reg)
    {
    }

    VMRegister reg;
};

struct Load : public Instruction {
    Load(VMRegister reg)
        : Instruction(Type::Load)
        , reg(reg)
    {
    }

    VMRegister reg;
};

struct SetLocal : public Instruction {
    SetLocal(VMLocal local)
        : Instruction(Type::SetLocal)
        , local(local)
    {
    }

    VMLocal local;
};

struct GetLocal : public Instruction {
    GetLocal(VMLocal local)
        : Instruction(Type::GetLocal)
        , local(local)
    {
    }

    VMLocal local;
};

struct Increment : public Instruction {
    Increment()
        : Instruction(Type::Increment)
    {
    }
};

struct LessThan : public Instruction {
    LessThan(VMRegister lhs)
        : Instruction(Type::LessThan)
        , lhs(lhs)
    {
    }

    VMRegister lhs;
};

struct Exit : public Instruction {
    Exit()
        : Instruction(Type::Exit)
    {
    }
};

struct Jump : public Instruction {
    Jump(BasicBlock& target)
        : Instruction(Type::Jump)
        , target(target)
    {
    }

    BasicBlock& target;
};

struct JumpConditional : public Instruction {
    JumpConditional(BasicBlock& true_target, BasicBlock& false_target)
        : Instruction(Type::JumpConditional)
        , true_target(true_target)
        , false_target(false_target)
    {
    }

    BasicBlock& true_target;
    BasicBlock& false_target;
};

struct VM {
    Vector<Value> registers;
    Vector<Value> locals;

    void dump() const
    {
        dbgln("Registers:");
        for (size_t i = 0; i < registers.size(); ++i) {
            dbgln("    [{}] {}", i, registers[i].value);
        }
        dbgln("Locals:");
        for (size_t i = 0; i < locals.size(); ++i) {
            dbgln("    [{}] {}", i, locals[i].value);
        }
    }
};

struct Program {
    Vector<NonnullOwnPtr<BasicBlock>> blocks;

    BasicBlock& make_block()
    {
        blocks.append(make<BasicBlock>());
        return *blocks.last();
    }

    void dump() const
    {
        for (size_t i = 0; i < blocks.size(); ++i) {
            dbgln("Block {}:", i + 1);
            blocks[i]->dump();
        }
    }
};

struct Executable {
    void* code { nullptr };
    size_t code_size { 0 };

    void run(VM& vm)
    {
        // RDI: VM&
        // RSI: Value* registers
        // RDX: Value* locals
        typedef void (*JittedFunction)(VM&, Value* registers, Value* locals);
        (*(JittedFunction)code)(vm, vm.registers.data(), vm.locals.data());
    }
};

struct Assembler {
    Assembler(Vector<u8>& output)
        : m_output(output)
    {
    }

    Vector<u8>& m_output;

    enum class Reg {
        GPR0 = 0, // RAX
        GPR1 = 1, // RCX

        RegisterArrayBase = 6, // RSI
        LocalsArrayBase = 2,   // RDX
    };

    struct Operand {
        enum class Type {
            Reg,
            Imm64,
            Mem64BaseAndOffset,
        };

        Type type {};

        Reg reg {};
        u64 offset_or_immediate { 0 };

        static Operand Register(Reg reg)
        {
            Operand operand;
            operand.type = Type::Reg;
            operand.reg = reg;
            return operand;
        }

        static Operand Imm64(u64 imm64)
        {
            Operand operand;
            operand.type = Type::Imm64;
            operand.offset_or_immediate = imm64;
            return operand;
        }

        static Operand Mem64BaseAndOffset(Reg base, u64 offset)
        {
            Operand operand;
            operand.type = Type::Mem64BaseAndOffset;
            operand.reg = base;
            operand.offset_or_immediate = offset;
            return operand;
        }
    };

    void mov(Operand dst, Operand src)
    {
        if (dst.type == Operand::Type::Reg && src.type == Operand::Type::Reg) {
            emit8(0x48);
            emit8(0x89);
            emit8(0xc0 | (to_underlying(dst.reg) << 3) | to_underlying(src.reg));
            return;
        }

        if (dst.type == Operand::Type::Reg && src.type == Operand::Type::Imm64) {
            emit8(0x48);
            emit8(0xb8 | to_underlying(dst.reg));
            emit64(src.offset_or_immediate);
            return;
        }

        if (dst.type == Operand::Type::Mem64BaseAndOffset && src.type == Operand::Type::Reg) {
            emit8(0x48);
            emit8(0x89);
            emit8(0x80 | (to_underlying(src.reg) << 3) | to_underlying(dst.reg));
            emit32(dst.offset_or_immediate);
            return;
        }

        if (dst.type == Operand::Type::Reg && src.type == Operand::Type::Mem64BaseAndOffset) {
            emit8(0x48);
            emit8(0x8b);
            emit8(0x80 | (to_underlying(dst.reg) << 3) | to_underlying(src.reg));
            emit32(src.offset_or_immediate);
            return;
        }

        VERIFY_NOT_REACHED();
    }

    void emit8(u8 value)
    {
        m_output.append(value);
    }

    void emit32(u32 value)
    {
        m_output.append((value >> 0) & 0xff);
        m_output.append((value >> 8) & 0xff);
        m_output.append((value >> 16) & 0xff);
        m_output.append((value >> 24) & 0xff);
    }

    void emit64(u64 value)
    {
        m_output.append((value >> 0) & 0xff);
        m_output.append((value >> 8) & 0xff);
        m_output.append((value >> 16) & 0xff);
        m_output.append((value >> 24) & 0xff);
        m_output.append((value >> 32) & 0xff);
        m_output.append((value >> 40) & 0xff);
        m_output.append((value >> 48) & 0xff);
        m_output.append((value >> 56) & 0xff);
    }

    void load_immediate64(Reg dst, u64 imm)
    {
        mov(Operand::Register(dst), Operand::Imm64(imm));
    }

    void store_vm_register(VMRegister dst, Reg src)
    {
        mov(Operand::Mem64BaseAndOffset(Reg::RegisterArrayBase, dst.value() * sizeof(Value)), Operand::Register(src));
    }

    void load_vm_register(Reg dst, VMRegister src)
    {
        mov(Operand::Register(dst), Operand::Mem64BaseAndOffset(Reg::RegisterArrayBase, src.value() * sizeof(Value)));
    }

    void store_vm_local(VMLocal dst, Reg src)
    {
        mov(Operand::Mem64BaseAndOffset(Reg::LocalsArrayBase, dst.value() * sizeof(Value)), Operand::Register(src));
    }

    void load_vm_local(Reg dst, VMLocal src)
    {
        mov(Operand::Register(dst), Operand::Mem64BaseAndOffset(Reg::LocalsArrayBase, src.value() * sizeof(Value)));
    }

    void increment(Reg dst)
    {
        emit8(0x48);
        emit8(0xff);
        emit8(0xc0 | to_underlying(dst));
    }

    void less_than(Reg dst, Reg src)
    {
        // cmp src, dst
        emit8(0x48);
        emit8(0x39);
        emit8(0xc0 | (to_underlying(src) << 3) | to_underlying(dst));

        // setl dst
        emit8(0x0f);
        emit8(0x9c);
        emit8(0xc0 | to_underlying(dst));

        // movzx dst, dst
        emit8(0x48);
        emit8(0x0f);
        emit8(0xb6);
        emit8(0xc0 | (to_underlying(dst) << 3) | to_underlying(dst));
    }

    void jump(BasicBlock& target)
    {
        // jmp target (RIP-relative 32-bit offset)
        emit8(0xe9);
        target.jumps_to_here.append(m_output.size());
        emit32(0xdeadbeef);
    }

    void jump_conditional(Reg reg, BasicBlock& true_target, BasicBlock& false_target)
    {
        // if reg == 0, jump to false_target, else jump to true_target
        // cmp reg, 0
        emit8(0x48);
        emit8(0x83);
        emit8(0xf8);
        emit8(0x00 | to_underlying(reg));

        // jz false_target (RIP-relative 32-bit offset)
        emit8(0x0f);
        emit8(0x84);
        false_target.jumps_to_here.append(m_output.size());
        emit32(0xdeadbeef);

        // jmp true_target (RIP-relative 32-bit offset)
        jump(true_target);
    }

    void exit()
    {
        // ret
        emit8(0xc3);
    }
};

struct JIT {
    Vector<u8> m_output;
    Assembler m_assembler { m_output };

    void compile_load_immediate(LoadImmediate const& instruction)
    {
        m_assembler.load_immediate64(Assembler::Reg::GPR0, instruction.value.value);
        m_assembler.store_vm_register(VMRegister { 0 }, Assembler::Reg::GPR0);
    }

    void compile_load(Load const& instruction)
    {
        m_assembler.load_vm_register(Assembler::Reg::GPR0, instruction.reg);
        m_assembler.store_vm_register(VMRegister { 0 }, Assembler::Reg::GPR0);
    }

    void compile_store(Store const& instruction)
    {
        m_assembler.load_vm_register(Assembler::Reg::GPR0, VMRegister { 0 });
        m_assembler.store_vm_register(instruction.reg, Assembler::Reg::GPR0);
    }

    void compile_get_local(GetLocal const& instruction)
    {
        m_assembler.load_vm_local(Assembler::Reg::GPR0, instruction.local);
        m_assembler.store_vm_register(VMRegister { 0 }, Assembler::Reg::GPR0);
    }

    void compile_set_local(SetLocal const& instruction)
    {
        m_assembler.load_vm_register(Assembler::Reg::GPR0, VMRegister { 0 });
        m_assembler.store_vm_local(instruction.local, Assembler::Reg::GPR0);
    }

    void compile_increment(Increment const&)
    {
        m_assembler.load_vm_register(Assembler::Reg::GPR0, VMRegister { 0 });
        m_assembler.increment(Assembler::Reg::GPR0);
        m_assembler.store_vm_register(VMRegister { 0 }, Assembler::Reg::GPR0);
    }

    void compile_less_than(LessThan const& instruction)
    {
        m_assembler.load_vm_register(Assembler::Reg::GPR0, instruction.lhs);
        m_assembler.load_vm_register(Assembler::Reg::GPR1, VMRegister { 0 });

        m_assembler.less_than(Assembler::Reg::GPR0, Assembler::Reg::GPR1);

        m_assembler.store_vm_register(VMRegister { 0 }, Assembler::Reg::GPR0);
    }

    void compile_jump(Jump const& instruction)
    {
        m_assembler.jump(instruction.target);
    }

    void compile_jump_conditional(JumpConditional const& instruction)
    {
        m_assembler.load_vm_register(Assembler::Reg::GPR0, VMRegister { 0 });
        m_assembler.jump_conditional(Assembler::Reg::GPR0, instruction.true_target, instruction.false_target);
    }

    void compile_exit(Exit const&)
    {
        m_assembler.exit();
    }

    static Executable compile(Program const& program)
    {
        JIT jit;

        for (auto& block : program.blocks) {
            block->offset = jit.m_output.size();
            for (auto& instruction : block->instructions) {
                switch (instruction->type) {
                case Instruction::Type::LoadImmediate:
                    jit.compile_load_immediate(static_cast<LoadImmediate&>(*instruction));
                    break;
                case Instruction::Type::Load:
                    jit.compile_load(static_cast<Load&>(*instruction));
                    break;
                case Instruction::Type::Store:
                    jit.compile_store(static_cast<Store&>(*instruction));
                    break;
                case Instruction::Type::SetLocal:
                    jit.compile_set_local(static_cast<SetLocal&>(*instruction));
                    break;
                case Instruction::Type::GetLocal:
                    jit.compile_get_local(static_cast<GetLocal&>(*instruction));
                    break;
                case Instruction::Type::Increment:
                    jit.compile_increment(static_cast<Increment&>(*instruction));
                    break;
                case Instruction::Type::LessThan:
                    jit.compile_less_than(static_cast<LessThan&>(*instruction));
                    break;
                case Instruction::Type::Jump:
                    jit.compile_jump(static_cast<Jump&>(*instruction));
                    break;
                case Instruction::Type::JumpConditional:
                    jit.compile_jump_conditional(static_cast<JumpConditional&>(*instruction));
                    break;
                case Instruction::Type::Exit:
                    jit.compile_exit(static_cast<Exit&>(*instruction));
                    break;
                default:
                    VERIFY_NOT_REACHED();
                }
            }
        }

        for (auto& block : program.blocks) {
            for (auto& jump : block->jumps_to_here) {
                auto offset = block->offset - jump - 4;
                jit.m_output[jump + 0] = (offset >> 0) & 0xff;
                jit.m_output[jump + 1] = (offset >> 8) & 0xff;
                jit.m_output[jump + 2] = (offset >> 16) & 0xff;
                jit.m_output[jump + 3] = (offset >> 24) & 0xff;
            }
        }

        write(STDOUT_FILENO, jit.m_output.data(), jit.m_output.size());

        auto* executable_memory = mmap(nullptr, jit.m_output.size(), PROT_READ | PROT_WRITE | PROT_EXEC, MAP_ANONYMOUS | MAP_PRIVATE, 0, 0);
        VERIFY(executable_memory != MAP_FAILED);

        memcpy(executable_memory, jit.m_output.data(), jit.m_output.size());

        return Executable {
            .code = executable_memory,
            .code_size = jit.m_output.size(),
        };
    }
};

ErrorOr<int> serenity_main(Main::Arguments)
{
    auto program = Program {};
    auto& block1 = program.make_block();
    auto& block2 = program.make_block();
    auto& block3 = program.make_block();
    auto& block4 = program.make_block();
    auto& block5 = program.make_block();
    auto& block6 = program.make_block();

    block1.append<Store>(VMRegister { 5 });
    block1.append<LoadImmediate>(Value(0));
    block1.append<SetLocal>(VMLocal(0));
    block1.append<Load>(VMRegister(5));
    block1.append<LoadImmediate>(Value(0));
    block1.append<Store>(VMRegister(6));
    block1.append<Jump>(block4);

    block2.append<Exit>();

    block3.append<LoadImmediate>(Value(0));
    block3.append<Jump>(block5);

    block4.append<GetLocal>(VMLocal(0));
    block4.append<Store>(VMRegister(7));
    block4.append<LoadImmediate>(Value(100000000));
    block4.append<LessThan>(VMRegister(7));
    block4.append<JumpConditional>(block3, block6);

    block5.append<Store>(VMRegister(6));
    block5.append<GetLocal>(VMLocal(0));
    block5.append<Increment>();
    block5.append<SetLocal>(VMLocal(0));
    block5.append<Jump>(block4);

    block6.append<Load>(VMRegister(6));
    block6.append<Jump>(block2);

    program.dump();

    auto vm = VM {};
    vm.registers.resize(8);
    vm.locals.resize(1);

    auto executable = JIT::compile(program);
    executable.run(vm);

#if 0
    auto* current_block = &block1;
    size_t instruction_index = 0;

    for (;;) {
        if (instruction_index >= current_block->instructions.size()) {
            break;
        }
        auto& instruction = current_block->instructions[instruction_index];
        switch (instruction->type) {
        case Instruction::Type::LoadImmediate:
            vm.registers[0].value = static_cast<LoadImmediate&>(*instruction).value.value;
            break;
        case Instruction::Type::Load:
            vm.registers[0].value = vm.registers[static_cast<Load&>(*instruction).reg.value()].value;
            break;
        case Instruction::Type::Store:
            vm.registers[static_cast<Store&>(*instruction).reg.value()].value = vm.registers[0].value;
            break;
        case Instruction::Type::SetLocal:
            vm.locals[static_cast<SetLocal&>(*instruction).local.value()].value = vm.registers[0].value;
            break;
        case Instruction::Type::GetLocal:
            vm.registers[0].value = vm.locals[static_cast<GetLocal&>(*instruction).local.value()].value;
            break;
        case Instruction::Type::Increment:
            vm.registers[0].value += 1;
            break;
        case Instruction::Type::LessThan:
            vm.registers[0].value = vm.registers[static_cast<LessThan&>(*instruction).lhs.value()].value < vm.registers[0].value;
            break;
        case Instruction::Type::Jump:
            current_block = &static_cast<Jump&>(*instruction).target;
            instruction_index = 0;
            continue;
        case Instruction::Type::JumpConditional:
            if (vm.registers[0].value) {
                current_block = &static_cast<JumpConditional&>(*instruction).true_target;
            } else {
                current_block = &static_cast<JumpConditional&>(*instruction).false_target;
            }
            instruction_index = 0;
            continue;
        case Instruction::Type::Exit:
            vm.dump();
            return 0;
        default:
            VERIFY_NOT_REACHED();
        }

        ++instruction_index;
    }
#endif

    vm.dump();

    return 0;
}