Parsing LLVM stackmaps.

lib/llvm-backend/src/backend.rs

@@ -1,3 +1,4 @@
+use super::stackmap::{self, StackmapRegistry};
 use crate::intrinsics::Intrinsics;
 use inkwell::{
     memory_buffer::MemoryBuffer,
@@ -25,6 +26,7 @@ use wasmer_runtime_core::{
     },
     cache::Error as CacheError,
     module::ModuleInfo,
+    state::ModuleStateMap,
     structures::TypedIndex,
     typed_func::{Wasm, WasmTrapInfo},
     types::{LocalFuncIndex, SigIndex},
@@ -76,6 +78,10 @@ extern "C" {
     ) -> LLVMResult;
     fn module_delete(module: *mut LLVMModule);
     fn get_func_symbol(module: *mut LLVMModule, name: *const c_char) -> *const vm::Func;
+    fn llvm_backend_get_stack_map_ptr(module: *const LLVMModule) -> *const u8;
+    fn llvm_backend_get_stack_map_size(module: *const LLVMModule) -> usize;
+    fn llvm_backend_get_code_ptr(module: *const LLVMModule) -> *const u8;
+    fn llvm_backend_get_code_size(module: *const LLVMModule) -> usize;
 
     fn throw_trap(ty: i32);
 
@@ -231,10 +237,15 @@ pub struct LLVMBackend {
     module: *mut LLVMModule,
     #[allow(dead_code)]
     buffer: Arc<Buffer>,
+    msm: Option<ModuleStateMap>,
 }
 
 impl LLVMBackend {
-    pub fn new(module: Module, _intrinsics: Intrinsics) -> (Self, LLVMCache) {
+    pub fn new(
+        module: Module,
+        _intrinsics: Intrinsics,
+        _stackmaps: &StackmapRegistry,
+    ) -> (Self, LLVMCache) {
         Target::initialize_x86(&InitializationConfig {
             asm_parser: true,
             asm_printer: true,
@@ -285,10 +296,24 @@ impl LLVMBackend {
 
         let buffer = Arc::new(Buffer::LlvmMemory(memory_buffer));
 
+        let raw_stackmap = unsafe {
+            ::std::slice::from_raw_parts(
+                llvm_backend_get_stack_map_ptr(module),
+                llvm_backend_get_stack_map_size(module),
+            )
+        };
+        if raw_stackmap.len() > 0 {
+            let map = stackmap::StackMap::parse(raw_stackmap);
+            eprintln!("{:?}", map);
+        } else {
+            eprintln!("WARNING: No stack map");
+        }
+
         (
             Self {
                 module,
                 buffer: Arc::clone(&buffer),
+                msm: None,
             },
             LLVMCache { buffer },
         )
@@ -318,6 +343,7 @@ impl LLVMBackend {
             Self {
                 module,
                 buffer: Arc::clone(&buffer),
+                msm: None,
             },
             LLVMCache { buffer },
         ))
@@ -372,6 +398,19 @@ impl RunnableModule for LLVMBackend {
         Some(unsafe { Wasm::from_raw_parts(trampoline, invoke_trampoline, None) })
     }
 
+    fn get_code(&self) -> Option<&[u8]> {
+        Some(unsafe {
+            ::std::slice::from_raw_parts(
+                llvm_backend_get_code_ptr(self.module),
+                llvm_backend_get_code_size(self.module),
+            )
+        })
+    }
+
+    fn get_module_state_map(&self) -> Option<ModuleStateMap> {
+        self.msm.clone()
+    }
+
     unsafe fn do_early_trap(&self, data: Box<dyn Any>) -> ! {
         throw_any(Box::leak(data))
     }

lib/llvm-backend/src/code.rs

@@ -27,6 +27,7 @@ use wasmparser::{BinaryReaderError, MemoryImmediate, Operator, Type as WpType};
 use crate::backend::LLVMBackend;
 use crate::intrinsics::{CtxType, GlobalCache, Intrinsics, MemoryCache};
 use crate::read_info::{blocktype_to_type, type_to_type};
+use crate::stackmap::{StackmapEntry, StackmapEntryKind, StackmapRegistry};
 use crate::state::{ControlFrame, IfElseState, State};
 use crate::trampolines::generate_trampolines;
 
@@ -342,26 +343,6 @@ pub struct CodegenError {
     pub message: String,
 }
 
-#[derive(Default, Debug, Clone)]
-pub struct StackmapRegistry {
-    entries: Vec<StackmapEntry>,
-}
-
-#[derive(Debug, Clone)]
-pub struct StackmapEntry {
-    kind: StackmapEntryKind,
-    local_function_id: usize,
-    local_count: usize,
-    stack_count: usize,
-}
-
-#[derive(Debug, Clone, Copy)]
-pub enum StackmapEntryKind {
-    Loop,
-    Call,
-    Trappable,
-}
-
 pub struct LLVMModuleCodeGenerator {
     context: Option<Context>,
     builder: Option<Builder>,
@@ -2653,7 +2634,10 @@ impl ModuleCodeGenerator<LLVMFunctionCodeGenerator, LLVMBackend, CodegenError>
 
         // self.module.print_to_stderr();
 
-        let (backend, cache_gen) = LLVMBackend::new(self.module, self.intrinsics.take().unwrap());
+        let stackmaps = self.stackmaps.borrow();
+
+        let (backend, cache_gen) =
+            LLVMBackend::new(self.module, self.intrinsics.take().unwrap(), &*stackmaps);
         Ok((backend, Box::new(cache_gen)))
     }
 

lib/llvm-backend/src/lib.rs

@@ -6,6 +6,7 @@ mod code;
 mod intrinsics;
 mod platform;
 mod read_info;
+mod stackmap;
 mod state;
 mod trampolines;
 

lib/llvm-backend/src/stackmap.rs

@@ -0,0 +1,208 @@
+// https://llvm.org/docs/StackMaps.html#stackmap-section
+
+use byteorder::{LittleEndian, ReadBytesExt};
+use std::io::{self, Cursor};
+
+#[derive(Default, Debug, Clone)]
+pub struct StackmapRegistry {
+    pub entries: Vec<StackmapEntry>,
+}
+
+#[derive(Debug, Clone)]
+pub struct StackmapEntry {
+    pub kind: StackmapEntryKind,
+    pub local_function_id: usize,
+    pub local_count: usize,
+    pub stack_count: usize,
+}
+
+#[derive(Debug, Clone, Copy)]
+pub enum StackmapEntryKind {
+    Loop,
+    Call,
+    Trappable,
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct StackMap {
+    pub version: u8,
+    pub stk_size_records: Vec<StkSizeRecord>,
+    pub constants: Vec<Constant>,
+    pub stk_map_records: Vec<StkMapRecord>,
+}
+
+#[derive(Copy, Clone, Debug, Default)]
+pub struct StkSizeRecord {
+    pub function_address: u64,
+    pub stack_size: u64,
+    pub record_count: u64,
+}
+
+#[derive(Copy, Clone, Debug, Default)]
+pub struct Constant {
+    pub large_constant: u64,
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct StkMapRecord {
+    pub patchpoint_id: u64,
+    pub instruction_offset: u32,
+    pub locations: Vec<Location>,
+    pub live_outs: Vec<LiveOut>,
+}
+
+#[derive(Copy, Clone, Debug)]
+pub struct Location {
+    pub ty: LocationType,
+    pub location_size: u16,
+    pub dwarf_regnum: u16,
+    pub offset_or_small_constant: i32,
+}
+
+#[derive(Copy, Clone, Debug, Default)]
+pub struct LiveOut {
+    pub dwarf_regnum: u16,
+    pub size_in_bytes: u8,
+}
+
+#[derive(Copy, Clone, Debug)]
+pub enum LocationType {
+    Register,
+    Direct,
+    Indirect,
+    Constant,
+    ConstantIndex,
+}
+
+impl StackMap {
+    pub fn parse(raw: &[u8]) -> io::Result<StackMap> {
+        let mut reader = Cursor::new(raw);
+        let mut map = StackMap::default();
+
+        let version = reader.read_u8()?;
+        if version != 3 {
+            return Err(io::Error::new(io::ErrorKind::Other, "version is not 3"));
+        }
+        map.version = version;
+        if reader.read_u8()? != 0 {
+            return Err(io::Error::new(
+                io::ErrorKind::Other,
+                "reserved field is not zero (1)",
+            ));
+        }
+        if reader.read_u16::<LittleEndian>()? != 0 {
+            return Err(io::Error::new(
+                io::ErrorKind::Other,
+                "reserved field is not zero (2)",
+            ));
+        }
+        let num_functions = reader.read_u32::<LittleEndian>()?;
+        let num_constants = reader.read_u32::<LittleEndian>()?;
+        let num_records = reader.read_u32::<LittleEndian>()?;
+        for _ in 0..num_functions {
+            let mut record = StkSizeRecord::default();
+            record.function_address = reader.read_u64::<LittleEndian>()?;
+            record.stack_size = reader.read_u64::<LittleEndian>()?;
+            record.record_count = reader.read_u64::<LittleEndian>()?;
+            map.stk_size_records.push(record);
+        }
+        for _ in 0..num_constants {
+            map.constants.push(Constant {
+                large_constant: reader.read_u64::<LittleEndian>()?,
+            });
+        }
+        for _ in 0..num_records {
+            let mut record = StkMapRecord::default();
+
+            record.patchpoint_id = reader.read_u64::<LittleEndian>()?;
+            record.instruction_offset = reader.read_u32::<LittleEndian>()?;
+            if reader.read_u16::<LittleEndian>()? != 0 {
+                return Err(io::Error::new(
+                    io::ErrorKind::Other,
+                    "reserved field is not zero (3)",
+                ));
+            }
+            let num_locations = reader.read_u16::<LittleEndian>()?;
+            for _ in 0..num_locations {
+                let ty = reader.read_u8()?;
+
+                let mut location = Location {
+                    ty: match ty {
+                        1 => LocationType::Register,
+                        2 => LocationType::Direct,
+                        3 => LocationType::Indirect,
+                        4 => LocationType::Constant,
+                        5 => LocationType::ConstantIndex,
+                        _ => {
+                            return Err(io::Error::new(
+                                io::ErrorKind::Other,
+                                "unknown location type",
+                            ))
+                        }
+                    },
+                    location_size: 0,
+                    dwarf_regnum: 0,
+                    offset_or_small_constant: 0,
+                };
+
+                if reader.read_u8()? != 0 {
+                    return Err(io::Error::new(
+                        io::ErrorKind::Other,
+                        "reserved field is not zero (4)",
+                    ));
+                }
+                location.location_size = reader.read_u16::<LittleEndian>()?;
+                location.dwarf_regnum = reader.read_u16::<LittleEndian>()?;
+                if reader.read_u16::<LittleEndian>()? != 0 {
+                    return Err(io::Error::new(
+                        io::ErrorKind::Other,
+                        "reserved field is not zero (5)",
+                    ));
+                }
+                location.offset_or_small_constant = reader.read_i32::<LittleEndian>()?;
+
+                record.locations.push(location);
+            }
+            if reader.position() % 8 != 0 {
+                if reader.read_u32::<LittleEndian>()? != 0 {
+                    return Err(io::Error::new(
+                        io::ErrorKind::Other,
+                        "reserved field is not zero (6)",
+                    ));
+                }
+            }
+            if reader.read_u16::<LittleEndian>()? != 0 {
+                return Err(io::Error::new(
+                    io::ErrorKind::Other,
+                    "reserved field is not zero (7)",
+                ));
+            }
+            let num_live_outs = reader.read_u16::<LittleEndian>()?;
+            for _ in 0..num_live_outs {
+                let mut liveout = LiveOut::default();
+
+                liveout.dwarf_regnum = reader.read_u16::<LittleEndian>()?;
+                if reader.read_u8()? != 0 {
+                    return Err(io::Error::new(
+                        io::ErrorKind::Other,
+                        "reserved field is not zero (8)",
+                    ));
+                }
+                liveout.size_in_bytes = reader.read_u8()?;
+
+                record.live_outs.push(liveout);
+            }
+            if reader.position() % 8 != 0 {
+                if reader.read_u32::<LittleEndian>()? != 0 {
+                    return Err(io::Error::new(
+                        io::ErrorKind::Other,
+                        "reserved field is not zero (9)",
+                    ));
+                }
+            }
+
+            map.stk_map_records.push(record);
+        }
+        Ok(map)
+    }
+}