tidepool_codegen/gc/
frame_walker.rs

1use crate::stack_map::StackMapRegistry;
2
3/// A collected GC root: the address on the stack where a heap pointer lives.
4#[derive(Debug, Clone, Copy)]
5pub struct StackRoot {
6    /// Address on the stack containing the heap pointer.
7    pub stack_slot_addr: *mut u64,
8    /// Current value of the heap pointer.
9    pub heap_ptr: *mut u8,
10}
11
12/// Walk JIT frames starting from the given RBP, collecting all GC roots.
13///
14/// # Safety
15/// - `start_rbp` must be a valid frame pointer from within a JIT call chain.
16/// - `stack_maps` must contain entries for all JIT functions in the call chain.
17#[cfg(target_arch = "x86_64")]
18pub unsafe fn walk_frames(
19    start_rbp: usize,
20    stack_maps: &StackMapRegistry,
21    start_rsp: usize,
22) -> Vec<StackRoot> {
23    let mut roots = Vec::new();
24    let mut rbp = start_rbp;
25    
26    // First, try to find the first JIT return address by searching up from RSP.
27    // This handles cases where gc_trigger doesn't have a frame pointer.
28    let mut current_return_addr = None;
29    let mut search_ptr = start_rsp;
30    // We search up to RBP + 16 (where the JIT return addr would be if gc_trigger has no frame).
31    while search_ptr < start_rbp + 16 {
32        let val = *(search_ptr as *const usize);
33        if stack_maps.contains_address(val) {
34            current_return_addr = Some(val);
35            // If we found a JIT return address, the RBP for this frame is the current RBP
36            // if gc_trigger has no frame, or it's the saved RBP if it does.
37            // Actually, if we found a JIT return address at search_ptr, 
38            // then search_ptr + 8 is the SP at the safepoint!
39            break;
40        }
41        search_ptr += 8;
42    }
43
44    loop {
45        if rbp == 0 {
46            break;
47        }
48
49        // Determine return address for this frame
50        let return_addr = if let Some(addr) = current_return_addr.take() {
51            addr
52        } else {
53            *((rbp + 8) as *const usize)
54        };
55        
56        // Check if this return address is in JIT code
57        if !stack_maps.contains_address(return_addr) {
58            if !roots.is_empty() {
59                // We were in JIT territory and now we left it. Stop.
60                break;
61            } else {
62                // We haven't hit JIT territory yet. Skip this frame.
63                let next_rbp = *(rbp as *const usize);
64                if next_rbp == 0 || next_rbp == rbp || next_rbp < rbp {
65                    break;
66                }
67                rbp = next_rbp;
68                continue;
69            }
70        }
71
72        // Look up stack map for this return address
73        if let Some(info) = stack_maps.lookup(return_addr) {
74            // Compute SP at safepoint.
75            // Cranelift stack map offsets are SP-relative at the safepoint.
76            // The return address we found is the one pushed by the 'call' in JIT code.
77            // The SP just before that 'call' was (addr_of_return_addr + 8).
78            
79            // We need to find where this return_addr was on the stack.
80            //
81            // If this is the first frame we found, and it was found via the initial RSP search
82            // (proxied by `search_ptr < start_rbp + 16`), we use that search address.
83            // Otherwise, for any subsequent JIT frames, the return address is at [rbp + 8].
84            let addr_of_return_addr = if roots.is_empty() && search_ptr < start_rbp + 16 {
85                search_ptr
86            } else {
87                rbp + 8
88            };
89            
90            let sp_at_safepoint = addr_of_return_addr + 8;
91
92            for &offset in &info.offsets {
93                let root_addr = (sp_at_safepoint + offset as usize) as *mut u64;
94                let heap_ptr = *root_addr as *mut u8;
95                roots.push(StackRoot {
96                    stack_slot_addr: root_addr,
97                    heap_ptr,
98                });
99            }
100        }
101
102        // Walk to next frame: *(rbp) is the saved caller RBP
103        let next_rbp = *(rbp as *const usize);
104        
105        // Basic sanity checks to prevent infinite loops or jumping to null
106        if next_rbp == 0 || next_rbp == rbp || next_rbp < rbp {
107            break;
108        }
109        rbp = next_rbp;
110    }
111
112    roots
113}
114
115/// Rewrite forwarding pointers in stack slots after GC.
116///
117/// For each root, if the heap object has been moved (forwarding pointer),
118/// update the stack slot to point to the new location.
119///
120/// # Safety
121/// All roots must still be valid stack addresses.
122pub unsafe fn rewrite_roots(roots: &[StackRoot], forwarding_map: &dyn Fn(*mut u8) -> *mut u8) {
123    for root in roots {
124        let new_ptr = forwarding_map(root.heap_ptr);
125        if new_ptr != root.heap_ptr {
126            *root.stack_slot_addr = new_ptr as u64;
127        }
128    }
129}
tidepool_codegen/gc/frame_walker.rs

tidepool_codegen/gc/
frame_walker.rs
