unwind.rs

extern crate gimli;
extern crate libc;

use crate::codegen::CodePtr;
use std::collections::HashMap;

#[cfg(target_arch = "x86_64")]
use crate::backend::x86_64::Reg;
#[cfg(target_arch = "x86_64")]
use crate::backend::x86_64::unwind::*;

#[cfg(target_arch = "aarch64")]
use crate::backend::arm64::Reg;
#[cfg(target_arch = "aarch64")]
use crate::backend::arm64::unwind::*;


#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
pub enum CStackSetupRule {
    // Blocks that have the CalledFromC CIE are YJIT entry points; they are called
    // directly from the CRuby code.
    CalledFromC,
    // Blocks that have the NormalJumpFromJITCode CIE are targets of jumps from YJIT
    // generated code; the C-stack setup has already been done by an entry point.
    NormalJumpFromJITCode,
}

// Mirrors https://sourceware.org/binutils/docs/as/CFI-directives.html#g_t_002ecfi_005fstartproc-_005bsimple_005d
#[derive(Copy, Clone, Debug)]
#[allow(dead_code)]
pub enum CFIDirective {
    StartProc(CStackSetupRule, CodePtr),
    EndProc(),
    DefCFA(Reg, i32),
    DefCFARegister(Reg),
    DefCFAOffset(i32),
    AdjustCFAOffset(i32),
    Offset(Reg, i32),
    RelOffset(Reg, i32),
    Restore(Reg),
}

#[derive(Copy, Clone, Debug)]
pub struct CFIDirectiveWithOffset {
    pub offset: usize,
    pub directive: CFIDirective,
}

pub struct UnwindInfoManager {
    cies: HashMap<CStackSetupRule, gimli::write::CommonInformationEntry>,
    cie_indicies: HashMap<CStackSetupRule, usize>,
    write_buf: gimli::write::EndianVec<gimli::LittleEndian>,
    live_buf: Vec<DwarfBufAlignT>,
    fdes: Vec<(CStackSetupRule, gimli::write::FrameDescriptionEntry)>,
    next_fde_vec_index: usize,
    next_live_buf_write_index: usize,
    handler_registered: bool,
}

impl UnwindInfoManager {
    pub fn new() -> Self {
        Self {
            cies: Self::cies_for_all_stack_rules(),
            cie_indicies: HashMap::default(),
            write_buf: gimli::write::EndianVec::new(gimli::LittleEndian),
            live_buf: Vec::default(),
            fdes: Vec::default(),
            next_fde_vec_index: 0,
            next_live_buf_write_index: 0,
            handler_registered: false,
        }
    }

    pub fn add_unwind_info(&mut self, directives_and_offsets: &[CFIDirectiveWithOffset]) {
        let mut fde_stack_rule = CStackSetupRule::NormalJumpFromJITCode; // default stack rule
        let mut fde_current_cfa_offset = Self::initial_cfa_offset_for_rule(fde_stack_rule);
        let mut fde_start_addr: Option<CodePtr> = None;
        let mut fde_insns = Vec::<(u32, gimli::write::CallFrameInstruction)>::new();
        let mut start_offset = 0;
        let mut end_offset = 0;
        for diroff in directives_and_offsets.iter() {
            match diroff.directive {
                CFIDirective::StartProc(stack_rule, start_addr) => {
                    fde_stack_rule = stack_rule;
                    fde_current_cfa_offset = Self::initial_cfa_offset_for_rule(fde_stack_rule);
                    fde_start_addr = Some(start_addr);
                    start_offset = diroff.offset;
                },
                CFIDirective::EndProc() => {
                    end_offset = diroff.offset;
                },
                CFIDirective::DefCFA(reg, cfa_offset) => {
                    fde_insns.push((diroff.offset as u32, gimli::write::CallFrameInstruction::Cfa(
                        gimli::Register(reg.reg_no.into()), cfa_offset,
                    )));
                    fde_current_cfa_offset = cfa_offset;
                },
                CFIDirective::DefCFAOffset(cfa_offset) => {
                    fde_insns.push((diroff.offset as u32, gimli::write::CallFrameInstruction::CfaOffset(
                        cfa_offset,
                    )));
                    fde_current_cfa_offset = cfa_offset;
                },
                CFIDirective::DefCFARegister(reg) => {
                    fde_insns.push((diroff.offset as u32, gimli::write::CallFrameInstruction::CfaRegister(
                        gimli::Register(reg.reg_no.into())
                    )));
                },
                CFIDirective::AdjustCFAOffset(cfa_offset_adj) => {
                    fde_current_cfa_offset = fde_current_cfa_offset + cfa_offset_adj;
                    fde_insns.push((diroff.offset as u32, gimli::write::CallFrameInstruction::CfaOffset(
                        fde_current_cfa_offset,
                    )));
                },
                CFIDirective::Offset(reg, cfa_offset) => {
                    fde_insns.push((diroff.offset as u32, gimli::write::CallFrameInstruction::Offset(
                        gimli::Register(reg.reg_no.into()), cfa_offset,
                    )));
                },
                CFIDirective::RelOffset(reg, cfa_offset_rel) => {
                    fde_insns.push((diroff.offset as u32, gimli::write::CallFrameInstruction::Offset(
                        gimli::Register(reg.reg_no.into()), fde_current_cfa_offset + cfa_offset_rel,
                    )));
                }
                CFIDirective::Restore(reg) => {
                    fde_insns.push((diroff.offset as u32, gimli::write::CallFrameInstruction::Restore(
                        gimli::Register(reg.reg_no.into())
                    )));
                }
            };
        }
        let mut fde = gimli::write::FrameDescriptionEntry::new(
            gimli::write::Address::Constant(fde_start_addr.expect(".cfi_start_proc directive missing!").into_u64()),
            (end_offset - start_offset) as u32,
        );
        for (offset, instruction) in fde_insns.into_iter() {
            fde.add_instruction(offset, instruction)
        }
        self.fdes.push((fde_stack_rule, fde));
    }

    pub fn flush_and_register(&mut self) -> () {
        let copy_from = self.write_buf.slice().len();
        // Write all FDE's that haven't been written yet.
        for (stack_rule, fde) in self.fdes[self.next_fde_vec_index as usize..].iter() {
            let cie = self.cies.get(stack_rule).expect("all CStackSetupRules should have CIE's");
            let cie_index = match self.cie_indicies.get(stack_rule) {
                Some(offset) => *offset,
                None => {
                    // Also write the CIE if that hasn't been written yet.
                    let o = self.write_buf.slice().len();
                    cie.write(&mut self.write_buf, true).unwrap();
                    self.cie_indicies.insert(*stack_rule, o);
                    o
                }
            };
            fde.write(&mut self.write_buf, true, cie_index, cie).unwrap();
        }
        let copy_length = self.write_buf.slice().len() - copy_from;
        
        // De-register the handler
        if self.handler_registered {
            unsafe { __deregister_frame(Self::buf_as_u8_ptr_mut(&mut self.live_buf)) };
            self.handler_registered = false
        }
        
        // Copy the new unwind info into the live buf
        let size_of_empty_fde = std::mem::size_of::<u32>();
        // The size of the buffer needs to be the current size NOT including the terminator + the new FDE's to add + the terminator.
        let alloc_length = self.next_live_buf_write_index + copy_length + size_of_empty_fde;
        Self::resize_and_align_buf(&mut self.live_buf, alloc_length);

        let live_buf_slice = Self::buf_as_u8_slice_mut(&mut self.live_buf);
        let new_dwarf_src_slice = &self.write_buf.slice()[copy_from..];
        let new_dwarf_dst_slice = &mut live_buf_slice[self.next_live_buf_write_index..(self.next_live_buf_write_index + new_dwarf_src_slice.len())];
        new_dwarf_dst_slice.copy_from_slice(new_dwarf_src_slice);
        self.next_live_buf_write_index += new_dwarf_src_slice.len();
        // Copy in the terminator too.
        let terminator_bytes = (0 as u32).to_le_bytes();
        let terminator_dst_slice = &mut live_buf_slice[self.next_live_buf_write_index..(self.next_live_buf_write_index + terminator_bytes.len())];
        terminator_dst_slice.copy_from_slice(terminator_bytes.as_slice());

        // Re-register the handler.
        self.dump_dwarf_info();
        unsafe { __register_frame(Self::buf_as_u8_ptr_mut(&mut self.live_buf)) };
        self.handler_registered = true
    }

    fn resize_and_align_buf(membuf: &mut Vec<DwarfBufAlignT>, new_size: usize) -> () {
        let new_size_in_elements = ceil_div(new_size, std::mem::size_of::<DwarfBufAlignT>());
        if membuf.len() < new_size {
            // needs growth
            membuf.extend(vec![DwarfBufAlignT::default(); new_size_in_elements - membuf.len()]);
        } else if membuf.capacity() > new_size {
            // needs shrinkage
            membuf.truncate(new_size_in_elements);
        }
    }

    unsafe fn buf_as_u8_ptr_mut(membuf: &mut Vec<DwarfBufAlignT>) -> *mut u8 {
        std::mem::transmute::<*mut DwarfBufAlignT, *mut u8>(membuf.as_mut_ptr())
    }

    unsafe fn buf_as_u8_ptr(membuf: &Vec<DwarfBufAlignT>) -> *const u8 {
        std::mem::transmute::<*const DwarfBufAlignT, *const u8>(membuf.as_ptr())
    }

    fn buf_as_u8_slice_mut(membuf: &mut Vec<DwarfBufAlignT>) -> &mut [u8] {
        unsafe {
            let ptr_as_u8 = Self::buf_as_u8_ptr_mut(membuf);
            std::slice::from_raw_parts_mut(ptr_as_u8, membuf.len() * std::mem::size_of::<DwarfBufAlignT>())
        }
    }

    fn buf_as_u8_slice(membuf: &Vec<DwarfBufAlignT>) -> &[u8] {
        unsafe {
            let ptr_as_u8 = Self::buf_as_u8_ptr(membuf);
            std::slice::from_raw_parts(ptr_as_u8, membuf.len() * std::mem::size_of::<DwarfBufAlignT>())
        }
    }

    fn dump_dwarf_info(&self) {
        static mut N: usize = 0;
        use std::io::prelude::*;

        let fname = unsafe { format!("eh_frame_{}", N) };
        let mut file = std::fs::File::create(fname).unwrap();
        file.write_all(&Self::buf_as_u8_slice(&self.live_buf)).unwrap();
        unsafe { N += 1 };
    }
}

fn ceil_div(a: usize, b: usize) -> usize {
    a / b + (if a % b == 0 { 0 } else { 1 } )
}

extern "C" {
    pub fn __register_frame(fde: *mut libc::c_uchar);
    pub fn __deregister_frame(fde: *mut libc::c_uchar);
}