sp1-jit 6.2.0

JIT compilation for SP1 trace generation
Documentation
#![allow(clippy::fn_to_numeric_cast)]

use super::{TranspilerBackend, CONTEXT};
use crate::{
    DebugFn, EcallHandler, ExternFn, JitFunction, JitMemory, RiscOperand, RiscRegister,
    RiscvTranspiler,
};
use dynasmrt::{
    dynasm,
    x64::{Assembler, Rq},
    DynasmApi,
};
use hashbrown::HashMap;
use std::io;

impl RiscvTranspiler for TranspilerBackend {
    fn new(
        program_size: usize,
        memory_size: usize,
        max_trace_size: u64,
        pc_start: u64,
        pc_base: u64,
        clk_bump: u64,
    ) -> Result<Self, std::io::Error> {
        if pc_start < pc_base {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                "pc_start must be greater than pc_base",
            ));
        }

        let mut this = Self {
            inner: Assembler::new()?,
            jump_table: Vec::with_capacity(program_size),
            memory_size,
            has_instructions: false,
            pc_base,
            pc_start,
            // Register a dummy ecall handler.
            ecall_handler: super::ecallk as _,
            control_flow_instruction_inserted: false,
            instruction_started: false,
            branch_generated: false,
            clk_bump,
            max_trace_size,
            may_early_exit: false,
            pc_current: pc_base,
            reg_values: HashMap::new(),
            labels: HashMap::new(),
            program_size,
        };

        // Handle calling conventions and save anything were gonna clobber.
        this.prologue();

        Ok(this)
    }

    fn register_ecall_handler(&mut self, handler: EcallHandler) {
        self.ecall_handler = handler;
    }

    fn start_instr(&mut self) {
        // We dont want to compile without a single jumpdest, otherwise we will sigsegv.
        self.has_instructions = true;

        // If the instruction has already started, then we are in a bad state.
        if self.instruction_started {
            panic!("start_instr called without calling end_instr");
        }

        // Push the offset of the jumpdest for this instruction.
        let offset = self.inner.offset();
        self.jump_table.push(offset.0);

        // We are now "within" an instruction.
        self.instruction_started = true;
    }

    fn end_instr(&mut self) {
        if self.control_flow_instruction_inserted {
            // Control flow instructions might have multiple branch targets,
            // as a result, we let them call `end_branch` directly. Here we
            // only handle bookkeeping tasks for control flow instructions.
            if !self.branch_generated {
                panic!("No branch has been generated, maybe a JIT mistake?");
            }

            // When basic block ends, reset all transpile-time register values, as they
            // would be incorrect for the next basic block.
            self.reg_values.clear();
        } else {
            // Add the base amount of cycles for the instruction.
            self.bump_clk();

            // We only bump / update PC when:
            // * A control flow instruction is executing.
            // * Before ecall and unimp calls extern functions.
            // * When trace is complete, execution suspends.
            // For normal, sequential operations, there is no need to bump PC.

            // We dont have a control flow insruction so we need to bump the pc first.
            if self.may_early_exit {
                self.exit_if_trace_exceeds(self.max_trace_size);
            }
        }

        // Transpiling is done for current instruction
        self.pc_current += 4;

        self.may_early_exit = false;
        self.control_flow_instruction_inserted = false;
        self.instruction_started = false;
        self.branch_generated = false;
    }

    fn finalize<M: JitMemory>(mut self) -> io::Result<JitFunction<M>> {
        self.epilogue();

        let code = self.inner.finalize().expect("failed to finalize x86 backend");

        debug_assert!(code.size() > 0, "Got empty x86 code buffer");

        JitFunction::new(code, self.jump_table, self.memory_size, self.pc_start)
    }

    fn call_extern_fn(&mut self, fn_ptr: ExternFn) {
        // Load the JitContext pointer into the argument register.
        dynasm! {
            self;
            .arch x64;
            mov rdi, Rq(CONTEXT)
        };

        self.call_extern_fn_raw(fn_ptr as _);
    }

    fn inspect_register(&mut self, reg: RiscRegister, handler: DebugFn) {
        // Load into the argument register for the function call.
        self.emit_risc_operand_load(RiscOperand::Register(reg), Rq::RDI as u8);

        // Call the handler with the value of the register.
        self.call_extern_fn_raw(handler as _);
    }

    fn inspect_immediate(&mut self, imm: u64, handler: DebugFn) {
        dynasm! {
            self;
            .arch x64;

            mov rdi, imm as i32
        }

        self.call_extern_fn_raw(handler as _);
    }
}