Merge pull request #2811 from peterhuene/improve-store-registration

Refactor store frame information.

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

@@ -195,11 +195,19 @@ pub struct TypeTables {
 
 /// Container for data needed for an Instance function to exist.
 pub struct ModuleCode {
+    range: (usize, usize),
     code_memory: CodeMemory,
     #[allow(dead_code)]
     dbg_jit_registration: Option<GdbJitImageRegistration>,
 }
 
+impl ModuleCode {
+    /// Gets the [begin, end) range of the module's code.
+    pub fn range(&self) -> (usize, usize) {
+        self.range
+    }
+}
+
 /// A compiled wasm module, ready to be instantiated.
 pub struct CompiledModule {
     artifacts: CompilationArtifacts,
@@ -259,10 +267,13 @@ impl CompiledModule {
         };
 
         let finished_functions = FinishedFunctions(finished_functions);
+        let start = code_range.0 as usize;
+        let end = start + code_range.1;
 
         Ok(Arc::new(Self {
             artifacts,
             code: Arc::new(ModuleCode {
+                range: (start, end),
                 code_memory,
                 dbg_jit_registration,
             }),
@@ -312,25 +323,53 @@ impl CompiledModule {
         )
     }
 
-    /// Iterates over all functions in this module, returning information about
-    /// how to decode traps which happen in the function.
-    pub fn trap_information(
-        &self,
-    ) -> impl Iterator<
-        Item = (
-            DefinedFuncIndex,
-            *mut [VMFunctionBody],
-            &[TrapInformation],
-            &FunctionAddressMap,
-        ),
-    > {
-        self.finished_functions()
-            .iter()
-            .zip(self.artifacts.funcs.values())
-            .map(|((i, alloc), func)| (i, *alloc, func.traps.as_slice(), &func.address_map))
+    /// Lookups a defined function by a program counter value.
+    ///
+    /// 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();
+
+        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
+            }
+        };
+
+        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))
     }
 
-    /// Returns all ranges convered by JIT code.
+    /// Gets the function information for a given function index.
+    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.
     pub fn jit_code_ranges<'a>(&'a self) -> impl Iterator<Item = (usize, usize)> + 'a {
         self.code.code_memory.published_ranges()
     }
@@ -454,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();