diff --git a/Cargo.lock b/Cargo.lock index e7faf37210..1fb9e09a6b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3407,7 +3407,6 @@ dependencies = [ "log", "more-asserts", "object", - "once_cell", "rayon", "region", "serde", diff --git a/cranelift/entity/src/primary.rs b/cranelift/entity/src/primary.rs index 9f43f088ce..f35c6f44a6 100644 --- a/cranelift/entity/src/primary.rs +++ b/cranelift/entity/src/primary.rs @@ -148,6 +148,28 @@ where pub fn into_boxed_slice(self) -> BoxedSlice { unsafe { BoxedSlice::::from_raw(Box::<[V]>::into_raw(self.elems.into_boxed_slice())) } } + + /// Performs a binary search on the values with a key extraction function. + /// + /// Assumes that the values are sorted by the key extracted by the function. + /// + /// If the value is found then `Ok(K)` is returned, containing the entity key + /// of the matching value. + /// + /// If there are multiple matches, then any one of the matches could be returned. + /// + /// If the value is not found then Err(K) is returned, containing the entity key + /// where a matching element could be inserted while maintaining sorted order. + pub fn binary_search_values_by_key<'a, B, F>(&'a self, b: &B, f: F) -> Result + where + F: FnMut(&'a V) -> B, + B: Ord, + { + self.elems + .binary_search_by_key(b, f) + .map(|i| K::new(i)) + .map_err(|i| K::new(i)) + } } impl Default for PrimaryMap diff --git a/crates/jit/Cargo.toml b/crates/jit/Cargo.toml index f653df5942..8d7ef7a04b 100644 --- a/crates/jit/Cargo.toml +++ b/crates/jit/Cargo.toml @@ -37,7 +37,6 @@ gimli = { version = "0.23.0", default-features = false, features = ["write"] } object = { version = "0.23.0", default-features = false, features = ["write"] } serde = { version = "1.0.94", features = ["derive"] } addr2line = { version = "0.14", default-features = false } -once_cell = "1.7.2" [target.'cfg(target_os = "windows")'.dependencies] winapi = { version = "0.3.8", features = ["winnt", "impl-default"] } diff --git a/crates/jit/src/code_memory.rs b/crates/jit/src/code_memory.rs index 9c703f6440..3362ed8b57 100644 --- a/crates/jit/src/code_memory.rs +++ b/crates/jit/src/code_memory.rs @@ -61,6 +61,15 @@ impl<'a> CodeMemoryObjectAllocation<'a> { pub fn code_range(self) -> &'a mut [u8] { self.buf } + + pub fn funcs_len(&self) -> usize { + self.funcs.len() + } + + pub fn trampolines_len(&self) -> usize { + self.trampolines.len() + } + pub fn funcs(&'a self) -> impl Iterator + 'a { let buf = self.buf as *const _ as *mut [u8]; self.funcs.iter().map(move |(i, (start, len))| { @@ -69,6 +78,7 @@ impl<'a> CodeMemoryObjectAllocation<'a> { }) }) } + pub fn trampolines( &'a self, ) -> impl Iterator + 'a { diff --git a/crates/jit/src/instantiate.rs b/crates/jit/src/instantiate.rs index 6a6007c255..c6e6f0ef05 100644 --- a/crates/jit/src/instantiate.rs +++ b/crates/jit/src/instantiate.rs @@ -8,11 +8,9 @@ use crate::compiler::{Compilation, Compiler}; use crate::link::link_module; use crate::object::ObjectUnwindInfo; use object::File as ObjectFile; -use once_cell::sync::OnceCell; #[cfg(feature = "parallel-compilation")] use rayon::prelude::*; use serde::{Deserialize, Serialize}; -use std::collections::BTreeMap; use std::ops::Range; use std::sync::Arc; use thiserror::Error; @@ -216,7 +214,6 @@ pub struct CompiledModule { code: Arc, finished_functions: FinishedFunctions, trampolines: Vec<(SignatureIndex, VMTrampoline)>, - func_map: OnceCell>, } impl CompiledModule { @@ -282,7 +279,6 @@ impl CompiledModule { }), finished_functions, trampolines, - func_map: Default::default(), })) } @@ -327,42 +323,50 @@ impl CompiledModule { ) } - /// Gets the function map of the compiled module. + /// Lookups a defined function by a program counter value. /// - /// The map is from ending address (inclusive) to a tuple of starting address and - /// defined function index. - /// - /// The map is lazily-initialized, so it will be populated the first time this - /// method is called. - pub fn func_map(&self) -> &BTreeMap { - self.func_map.get_or_init(|| { - let mut functions = BTreeMap::new(); - for (index, allocated) in self.finished_functions().iter() { - let (start, end) = unsafe { - let ptr = (**allocated).as_ptr(); - let len = (**allocated).len(); - // First and last byte of the function text. - (ptr as usize, ptr as usize + len - 1) - }; + /// Returns the defined function index, the start address, and the end address (exclusive). + pub fn func_by_pc(&self, pc: usize) -> Option<(DefinedFuncIndex, usize, usize)> { + let functions = self.finished_functions(); - // Finished functions cannot be empty - assert!(start <= end); - assert!(functions.insert(end, (start, index)).is_none()); + let index = match functions.binary_search_values_by_key(&pc, |body| unsafe { + debug_assert!(!(**body).is_empty()); + // Return the inclusive "end" of the function + (**body).as_ptr() as usize + (**body).len() - 1 + }) { + Ok(k) => { + // Exact match, pc is at the end of this function + k } + Err(k) => { + // Not an exact match, k is where `pc` would be "inserted" + // Since we key based on the end, function `k` might contain `pc`, + // so we'll validate on the range check below + k + } + }; - functions - }) + let body = functions.get(index)?; + let (start, end) = unsafe { + let ptr = (**body).as_ptr(); + let len = (**body).len(); + (ptr as usize, ptr as usize + len) + }; + + if pc < start || end < pc { + return None; + } + + Some((index, start, end)) } /// Gets the function information for a given function index. - pub fn func_info( - &self, - index: DefinedFuncIndex, - ) -> Option<(&FunctionAddressMap, &[TrapInformation])> { + pub fn func_info(&self, index: DefinedFuncIndex) -> (&FunctionAddressMap, &[TrapInformation]) { self.artifacts .funcs .get(index) .map(|f| (&f.address_map, f.traps.as_ref())) + .expect("defined function should be present") } /// Returns all ranges covered by JIT code. @@ -489,25 +493,33 @@ fn build_code_memory( let allocation = code_memory.allocate_for_object(&obj, unwind_info)?; - // Second, create a PrimaryMap from result vector of pointers. - let mut finished_functions = PrimaryMap::new(); + // Populate the finished functions from the allocation + let mut finished_functions = PrimaryMap::with_capacity(allocation.funcs_len()); for (i, fat_ptr) in allocation.funcs() { + let start = fat_ptr.as_ptr() as usize; let fat_ptr: *mut [VMFunctionBody] = fat_ptr; + // Assert that the function bodies are pushed in sort order + // This property is relied upon to search for functions by PC values + assert!( + start + > finished_functions + .last() + .map(|f: &*mut [VMFunctionBody]| unsafe { (**f).as_ptr() as usize }) + .unwrap_or(0) + ); assert_eq!( Some(finished_functions.push(fat_ptr)), module.defined_func_index(i) ); } - let trampolines = allocation - .trampolines() - .map(|(i, fat_ptr)| { - let fnptr = unsafe { - std::mem::transmute::<*const VMFunctionBody, VMTrampoline>(fat_ptr.as_ptr()) - }; - (i, fnptr) - }) - .collect(); + // Populate the trampolines from the allocation + let mut trampolines = Vec::with_capacity(allocation.trampolines_len()); + for (i, fat_ptr) in allocation.trampolines() { + let fnptr = + unsafe { std::mem::transmute::<*const VMFunctionBody, VMTrampoline>(fat_ptr.as_ptr()) }; + trampolines.push((i, fnptr)); + } let code_range = allocation.code_range(); diff --git a/crates/wasmtime/src/frame_info.rs b/crates/wasmtime/src/frame_info.rs index f9334461ac..c3abc1e99c 100644 --- a/crates/wasmtime/src/frame_info.rs +++ b/crates/wasmtime/src/frame_info.rs @@ -132,7 +132,7 @@ impl ModuleFrameInfo { /// if no information can be found. pub fn lookup_frame_info(&self, pc: usize) -> Option { let (index, offset) = self.func(pc)?; - let (addr_map, _) = self.module.func_info(index)?; + let (addr_map, _) = self.module.func_info(index); let pos = Self::instr_pos(offset, addr_map); // In debug mode for now assert that we found a mapping for `pc` within @@ -200,7 +200,7 @@ impl ModuleFrameInfo { /// Fetches trap information about a program counter in a backtrace. pub fn lookup_trap_info(&self, pc: usize) -> Option<&TrapInformation> { let (index, offset) = self.func(pc)?; - let (_, traps) = self.module.func_info(index)?; + let (_, traps) = self.module.func_info(index); let idx = traps .binary_search_by_key(&offset, |info| info.code_offset) .ok()?; @@ -208,13 +208,8 @@ impl ModuleFrameInfo { } fn func(&self, pc: usize) -> Option<(DefinedFuncIndex, u32)> { - let (end, (start, index)) = self.module.func_map().range(pc..).next()?; - - if pc < *start || *end < pc { - return None; - } - - Some((*index, (pc - *start) as u32)) + let (index, start, _) = self.module.func_by_pc(pc)?; + Some((index, (pc - start) as u32)) } fn instr_pos(offset: u32, addr_map: &FunctionAddressMap) -> Option { @@ -286,7 +281,7 @@ impl GlobalFrameInfo { match info.module.func(pc) { Some((index, offset)) => { - let (addr_map, _) = info.module.module.func_info(index).unwrap(); + let (addr_map, _) = info.module.module.func_info(index); ModuleFrameInfo::instr_pos(offset, addr_map).is_some() } None => false,