1 use crate::bitset::BitSet;
2 use crate::ir;
3 use crate::isa::TargetIsa;
4 use alloc::vec::Vec;
5 
6 type Num = u32;
7 const NUM_BITS: usize = core::mem::size_of::<Num>() * 8;
8 
9 /// A stack map is a bitmap with one bit per machine word on the stack. Stack
10 /// maps are created at `safepoint` instructions and record all live reference
11 /// values that are on the stack. All slot kinds, except `OutgoingArg` are
12 /// captured in a stack map. The `OutgoingArg`'s will be captured in the callee
13 /// function as `IncomingArg`'s.
14 ///
15 /// The first value in the bitmap is of the lowest addressed slot on the stack.
16 /// As all stacks in Isa's supported by Cranelift grow down, this means that
17 /// first value is of the top of the stack and values proceed down the stack.
18 #[derive(Clone, Debug, PartialEq, Eq)]
19 #[cfg_attr(feature = "enable-serde", derive(serde::Deserialize, serde::Serialize))]
20 pub struct Stackmap {
21     bitmap: Vec<BitSet<Num>>,
22     mapped_words: u32,
23 }
24 
25 impl Stackmap {
26     /// Create a stackmap based on where references are located on a function's stack.
from_values( args: &[ir::entities::Value], func: &ir::Function, isa: &dyn TargetIsa, ) -> Self27     pub fn from_values(
28         args: &[ir::entities::Value],
29         func: &ir::Function,
30         isa: &dyn TargetIsa,
31     ) -> Self {
32         let loc = &func.locations;
33         let mut live_ref_in_stack_slot = crate::HashSet::new();
34         // References can be in registers, and live registers values are pushed onto the stack before calls and traps.
35         // TODO: Implement register maps. If a register containing a reference is spilled and reused after a safepoint,
36         // it could contain a stale reference value if the garbage collector relocated the value.
37         for val in args {
38             if let Some(value_loc) = loc.get(*val) {
39                 match *value_loc {
40                     ir::ValueLoc::Stack(stack_slot) => {
41                         live_ref_in_stack_slot.insert(stack_slot);
42                     }
43                     _ => {}
44                 }
45             }
46         }
47 
48         let stack = &func.stack_slots;
49         let info = func.stack_slots.layout_info.unwrap();
50 
51         // Refer to the doc comment for `Stackmap` above to understand the
52         // bitmap representation used here.
53         let map_size = (info.frame_size + info.inbound_args_size) as usize;
54         let word_size = isa.pointer_bytes() as usize;
55         let num_words = map_size / word_size;
56 
57         let mut vec = alloc::vec::Vec::with_capacity(num_words);
58         vec.resize(num_words, false);
59 
60         for (ss, ssd) in stack.iter() {
61             if !live_ref_in_stack_slot.contains(&ss)
62                 || ssd.kind == ir::stackslot::StackSlotKind::OutgoingArg
63             {
64                 continue;
65             }
66 
67             debug_assert!(ssd.size as usize == word_size);
68             let bytes_from_bottom = info.frame_size as i32 + ssd.offset.unwrap();
69             let words_from_bottom = (bytes_from_bottom as usize) / word_size;
70             vec[words_from_bottom] = true;
71         }
72 
73         Self::from_slice(&vec)
74     }
75 
76     /// Create a vec of Bitsets from a slice of bools.
from_slice(vec: &[bool]) -> Self77     pub fn from_slice(vec: &[bool]) -> Self {
78         let len = vec.len();
79         let num_word = len / NUM_BITS + (len % NUM_BITS != 0) as usize;
80         let mut bitmap = Vec::with_capacity(num_word);
81 
82         for segment in vec.chunks(NUM_BITS) {
83             let mut curr_word = 0;
84             for (i, set) in segment.iter().enumerate() {
85                 if *set {
86                     curr_word |= 1 << i;
87                 }
88             }
89             bitmap.push(BitSet(curr_word));
90         }
91         Self {
92             mapped_words: len as u32,
93             bitmap,
94         }
95     }
96 
97     /// Returns a specified bit.
get_bit(&self, bit_index: usize) -> bool98     pub fn get_bit(&self, bit_index: usize) -> bool {
99         assert!(bit_index < NUM_BITS * self.bitmap.len());
100         let word_index = bit_index / NUM_BITS;
101         let word_offset = (bit_index % NUM_BITS) as u8;
102         self.bitmap[word_index].contains(word_offset)
103     }
104 
105     /// Returns the raw bitmap that represents this stack map.
as_slice(&self) -> &[BitSet<u32>]106     pub fn as_slice(&self) -> &[BitSet<u32>] {
107         &self.bitmap
108     }
109 
110     /// Returns the number of words represented by this stack map.
mapped_words(&self) -> u32111     pub fn mapped_words(&self) -> u32 {
112         self.mapped_words
113     }
114 }
115 
116 #[cfg(test)]
117 mod tests {
118     use super::*;
119 
120     #[test]
stackmaps()121     fn stackmaps() {
122         let vec: Vec<bool> = Vec::new();
123         assert!(Stackmap::from_slice(&vec).bitmap.is_empty());
124 
125         let mut vec: [bool; NUM_BITS] = Default::default();
126         let set_true_idx = [5, 7, 24, 31];
127 
128         for &idx in &set_true_idx {
129             vec[idx] = true;
130         }
131 
132         let mut vec = vec.to_vec();
133         assert_eq!(
134             vec![BitSet::<Num>(2164261024)],
135             Stackmap::from_slice(&vec).bitmap
136         );
137 
138         vec.push(false);
139         vec.push(true);
140         let res = Stackmap::from_slice(&vec);
141         assert_eq!(
142             vec![BitSet::<Num>(2164261024), BitSet::<Num>(2)],
143             res.bitmap
144         );
145 
146         assert!(res.get_bit(5));
147         assert!(res.get_bit(31));
148         assert!(res.get_bit(33));
149         assert!(!res.get_bit(1));
150     }
151 }
152