Code review feedback.

* Remove `once-cell` dependency.
* Remove function address `BTreeMap` from `CompiledModule` in favor of binary
  searching finished functions directly.
* Use `with_capacity` when populating `CompiledModule` finished functions and
  trampolines.

Cargo.lock

@@ -3407,7 +3407,6 @@ dependencies = [
  "log",
  "more-asserts",
  "object",
- "once_cell",
  "rayon",
  "region",
  "serde",

cranelift/entity/src/primary.rs

@@ -148,6 +148,28 @@ where
     pub fn into_boxed_slice(self) -> BoxedSlice<K, V> {
         unsafe { BoxedSlice::<K, V>::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<K, K>
+    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<K, V> Default for PrimaryMap<K, V>

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"] }

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<Item = (FuncIndex, &'a mut [VMFunctionBody])> + '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<Item = (SignatureIndex, &'a mut [VMFunctionBody])> + 'a {

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<ModuleCode>,
     finished_functions: FinishedFunctions,
     trampolines: Vec<(SignatureIndex, VMTrampoline)>,
-    func_map: OnceCell<BTreeMap<usize, (usize, DefinedFuncIndex)>>,
 }
 
 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
+    /// Returns the defined function index, the start address, and the end address (exclusive).
-    /// defined function index.
+    pub fn func_by_pc(&self, pc: usize) -> Option<(DefinedFuncIndex, usize, usize)> {
-    ///
+        let functions = self.finished_functions();
-    /// The map is lazily-initialized, so it will be populated the first time this
-    /// method is called.
-    pub fn func_map(&self) -> &BTreeMap<usize, (usize, DefinedFuncIndex)> {
-        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)
-                };
 
-                // Finished functions cannot be empty
+        let index = match functions.binary_search_values_by_key(&pc, |body| unsafe {
-                assert!(start <= end);
+            debug_assert!(!(**body).is_empty());
-                assert!(functions.insert(end, (start, index)).is_none());
+            // 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(
+    pub fn func_info(&self, index: DefinedFuncIndex) -> (&FunctionAddressMap, &[TrapInformation]) {
-        &self,
-        index: DefinedFuncIndex,
-    ) -> Option<(&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.
+    // Populate the finished functions from the allocation
-    let mut finished_functions = PrimaryMap::new();
+    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
+    // Populate the trampolines from the allocation
-        .trampolines()
+    let mut trampolines = Vec::with_capacity(allocation.trampolines_len());
-        .map(|(i, fat_ptr)| {
+    for (i, fat_ptr) in allocation.trampolines() {
-            let fnptr = unsafe {
+        let fnptr =
-                std::mem::transmute::<*const VMFunctionBody, VMTrampoline>(fat_ptr.as_ptr())
+            unsafe { std::mem::transmute::<*const VMFunctionBody, VMTrampoline>(fat_ptr.as_ptr()) };
-            };
+        trampolines.push((i, fnptr));
-            (i, fnptr)
+    }
-        })
-        .collect();
 
     let code_range = allocation.code_range();
 

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<FrameInfo> {
         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()?;
+        let (index, start, _) = self.module.func_by_pc(pc)?;
-
+        Some((index, (pc - start) as u32))
-        if pc < *start || *end < pc {
-            return None;
-        }
-
-        Some((*index, (pc - *start) as u32))
     }
 
     fn instr_pos(offset: u32, addr_map: &FunctionAddressMap) -> Option<usize> {
@@ -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,