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