wasmtime_environ/
stack_map.rs

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

wasmtime_environ/
stack_map.rs
