wasmtime_environ/
stack_map.rs

1use cranelift_bitset::ScalarBitSet;
2use object::{Bytes, LittleEndian, U32Bytes};
3
4struct StackMapSection<'a> {
5    pcs: &'a [U32Bytes<LittleEndian>],
6    pointers_to_stack_map: &'a [U32Bytes<LittleEndian>],
7    stack_map_data: &'a [U32Bytes<LittleEndian>],
8}
9
10impl<'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.
66pub struct StackMap<'a> {
67    frame_size: u32,
68    data: &'a [U32Bytes<LittleEndian>],
69}
70
71impl<'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}
wasmtime_environ/stack_map.rs

wasmtime_environ/
stack_map.rs
