tidepool_codegen/
stack_map.rs

1use std::collections::BTreeMap;
2
3/// Information about GC roots at a single safepoint.
4#[derive(Debug, Clone)]
5pub struct StackMapInfo {
6    /// Size of the frame in bytes (span from user_stack_maps tuple).
7    pub frame_size: u32,
8    /// SP-relative offsets of heap pointer slots.
9    /// root_addr = SP + offset at the safepoint.
10    pub offsets: Vec<u32>,
11}
12
13pub type RawStackMapEntry = (cranelift_codegen::ir::types::Type, u32);
14pub type RawStackMap = (u32, u32, Vec<RawStackMapEntry>);
15
16/// Maps absolute return addresses to stack map info.
17///
18/// Key = function_base_ptr + code_offset
19/// (i.e., the return address, which is what the frame walker sees as caller_pc).
20/// Cranelift's `code_offset` for user stack maps already points to the
21/// instruction AFTER the call (the return point).
22#[derive(Debug, Default)]
23pub struct StackMapRegistry {
24    entries: BTreeMap<usize, StackMapInfo>,
25    /// Known JIT function address ranges (start, end).
26    ranges: Vec<(usize, usize)>,
27}
28
29impl StackMapRegistry {
30    pub fn new() -> Self {
31        Self::default()
32    }
33
34    /// Register stack map entries from a compiled function.
35    ///
36    /// `base_ptr` is the start address of the compiled function in memory.
37    /// `size` is the total size of the function in bytes.
38    /// `raw_entries` come from `CompiledCode.buffer.user_stack_maps()`:
39    ///   each tuple is (code_offset, frame_size, UserStackMap).
40    ///
41    /// We key by `base_ptr + code_offset` as the return address. Cranelift's
42    /// `code_offset` for user stack maps points to the instruction AFTER the call
43    /// (the return point), so `base_ptr + code_offset` IS the absolute return address.
44    pub fn register(&mut self, base_ptr: usize, size: u32, raw_entries: &[RawStackMap]) {
45        self.ranges.push((base_ptr, base_ptr + size as usize));
46        
47        for (code_offset, frame_size, slot_entries) in raw_entries {
48            let return_addr = base_ptr + *code_offset as usize;
49            let offsets: Vec<u32> = slot_entries.iter().map(|(_, offset)| *offset).collect();
50            self.entries.insert(return_addr, StackMapInfo {
51                frame_size: *frame_size,
52                offsets,
53            });
54        }
55    }
56
57    /// Look up stack map info by return address (PC value from frame walker).
58    pub fn lookup(&self, return_addr: usize) -> Option<&StackMapInfo> {
59        self.entries.get(&return_addr)
60    }
61
62    /// Number of registered safepoints.
63    pub fn len(&self) -> usize {
64        self.entries.len()
65    }
66
67    /// Whether the registry is empty.
68    pub fn is_empty(&self) -> bool {
69        self.entries.is_empty()
70    }
71
72    /// Check if an address falls within the known JIT code region.
73    /// Used by the frame walker to determine when to stop walking.
74    pub fn contains_address(&self, addr: usize) -> bool {
75        self.ranges.iter().any(|(start, end)| addr >= *start && addr < *end)
76    }
77}
tidepool_codegen/stack_map.rs

tidepool_codegen/
stack_map.rs
