1 use cranelift_bitset::ScalarBitSet;
2 use object::{Bytes, LittleEndian, U32Bytes};
3 
4 struct StackMapSection<'a> {
5     pcs: &'a [U32Bytes<LittleEndian>],
6     pointers_to_stack_map: &'a [U32Bytes<LittleEndian>],
7     stack_map_data: &'a [U32Bytes<LittleEndian>],
8 }
9 
10 impl<'a> StackMapSection<'a> {
11     fn parse(section: &'a [u8]) -> Option<StackMapSection<'a>> {
12         let mut section = Bytes(section);
13         // NB: this matches the encoding written by `append_to` in the
14         // `compile::stack_map` module.
15         let pc_count = section.read::<U32Bytes<LittleEndian>>().ok()?;
16         let pc_count = usize::try_from(pc_count.get(LittleEndian)).ok()?;
17         let (pcs, section) =
18             object::slice_from_bytes::<U32Bytes<LittleEndian>>(section.0, pc_count).ok()?;
19         let (pointers_to_stack_map, section) =
20             object::slice_from_bytes::<U32Bytes<LittleEndian>>(section, pc_count).ok()?;
21         let stack_map_data =
22             object::slice_from_all_bytes::<U32Bytes<LittleEndian>>(section).ok()?;
23         Some(StackMapSection {
24             pcs,
25             pointers_to_stack_map,
26             stack_map_data,
27         })
28     }
29 
30     fn lookup(&self, pc: u32) -> Option<StackMap<'a>> {
31         let pc_index = self
32             .pcs
33             .binary_search_by_key(&pc, |v| v.get(LittleEndian))
34             .ok()?;
35         self.get(pc_index)
36     }
37 
38     fn into_iter(self) -> impl Iterator<Item = (u32, StackMap<'a>)> + 'a {
39         self.pcs
40             .iter()
41             .enumerate()
42             .map(move |(i, pc)| (pc.get(LittleEndian), self.get(i).unwrap()))
43     }
44 
45     /// Returns the stack map corresponding to the `i`th pc.
46     fn get(&self, i: usize) -> Option<StackMap<'a>> {
47         let pointer_to_stack_map = self.pointers_to_stack_map[i].get(LittleEndian) as usize;
48         let data = self.stack_map_data.get(pointer_to_stack_map..)?;
49 
50         let (frame_size, data) = data.split_first()?;
51         let (count, data) = data.split_first()?;
52         let data = data.get(..count.get(LittleEndian) as usize)?;
53 
54         Some(StackMap {
55             frame_size: frame_size.get(LittleEndian),
56             data,
57         })
58     }
59 }
60 
61 /// A map for determining where live GC references live in a stack frame.
62 ///
63 /// Note that this is currently primarily documented as cranelift's
64 /// `binemit::StackMap`, so for detailed documentation about this please read
65 /// the docs over there.
66 pub struct StackMap<'a> {
67     frame_size: u32,
68     data: &'a [U32Bytes<LittleEndian>],
69 }
70 
71 impl<'a> StackMap<'a> {
72     /// Looks up a stack map for `pc` within the `section` provided.
73     ///
74     /// The `section` should be produced by `StackMapSection` in the
75     /// `compile::stack_map` module. The `pc` should be relative to the start
76     /// of the `.text` section in the final executable.
77     pub fn lookup(pc: u32, section: &'a [u8]) -> Option<StackMap<'a>> {
78         StackMapSection::parse(section)?.lookup(pc)
79     }
80 
81     /// Iterate over the stack maps contained in the given stack map section.
82     ///
83     /// This function takes a `section` as its first argument which must have
84     /// been created with `StackMapSection` builder. This is intended to be the
85     /// raw `ELF_WASMTIME_STACK_MAP` section from the compilation artifact.
86     ///
87     /// The yielded offsets are relative to the start of the text section for
88     /// this map's code object.
89     pub fn iter(section: &'a [u8]) -> Option<impl Iterator<Item = (u32, StackMap<'a>)> + 'a> {
90         Some(StackMapSection::parse(section)?.into_iter())
91     }
92 
93     /// Returns the byte size of this stack map's frame.
94     pub fn frame_size(&self) -> u32 {
95         self.frame_size
96     }
97 
98     /// Given a frame pointer, get the stack pointer.
99     ///
100     /// # Safety
101     ///
102     /// The `fp` must be the frame pointer at the code offset that this stack
103     /// map is associated with.
104     pub unsafe fn sp(&self, fp: *mut usize) -> *mut usize {
105         let frame_size = usize::try_from(self.frame_size).unwrap();
106         unsafe { fp.byte_sub(frame_size) }
107     }
108 
109     /// Given the stack pointer, get a reference to each live GC reference in
110     /// the stack frame.
111     ///
112     /// # Safety
113     ///
114     /// The `sp` must be the stack pointer at the code offset that this stack
115     /// map is associated with.
116     pub unsafe fn live_gc_refs(&self, sp: *mut usize) -> impl Iterator<Item = *mut u32> + '_ {
117         self.offsets().map(move |i| {
118             log::trace!("Live GC ref in frame at frame offset {i:#x}");
119             let i = usize::try_from(i).unwrap();
120             let ptr_to_gc_ref = unsafe { sp.byte_add(i) };
121 
122             // Assert that the pointer is inside this stack map's frame.
123             assert!({
124                 let delta = ptr_to_gc_ref as usize - sp as usize;
125                 let frame_size = usize::try_from(self.frame_size).unwrap();
126                 delta < frame_size
127             });
128 
129             ptr_to_gc_ref.cast::<u32>()
130         })
131     }
132 
133     /// Returns the offsets that this stack map registers GC references at.
134     pub fn offsets(&self) -> impl Iterator<Item = u32> + '_ {
135         // Here `self.data` is a bit set of offsets divided by 4, so iterate
136         // over all the bits in `self.data` and multiply their position by 4.
137         let bit_positions = self.data.iter().enumerate().flat_map(|(i, word)| {
138             ScalarBitSet(word.get(LittleEndian))
139                 .iter()
140                 .map(move |bit| (i as u32) * 32 + u32::from(bit))
141         });
142 
143         bit_positions.map(|pos| pos * 4)
144     }
145 }
146