1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152

use crate::bitset::BitSet;
use alloc::vec::Vec;

type Num = u32;
const NUM_BITS: usize = core::mem::size_of::<Num>() * 8;

/// Stack maps record which words in a stack frame contain live GC references at
/// a given instruction pointer.
///
/// Logically, a set of stack maps for a function record a table of the form:
///
/// ```text
/// +---------------------+-------------------------------------------+
/// | Instruction Pointer | SP-Relative Offsets of Live GC References |
/// +---------------------+-------------------------------------------+
/// | 0x12345678          | 2, 6, 12                                  |
/// | 0x1234abcd          | 2, 6                                      |
/// | ...                 | ...                                       |
/// +---------------------+-------------------------------------------+
/// ```
///
/// Where "instruction pointer" is an instruction pointer within the function,
/// and "offsets of live GC references" contains the offsets (in units of words)
/// from the frame's stack pointer where live GC references are stored on the
/// stack. Instruction pointers within the function that do not have an entry in
/// this table are not GC safepoints.
///
/// Because
///
/// * offsets of live GC references are relative from the stack pointer, and
/// * stack frames grow down from higher addresses to lower addresses,
///
/// to get a pointer to a live reference at offset `x` within a stack frame, you
/// add `x` from the frame's stack pointer.
///
/// For example, to calculate the pointer to the live GC reference inside "frame
/// 1" below, you would do `frame_1_sp + x`:
///
/// ```text
///           Stack
///         +-------------------+
///         | Frame 0           |
///         |                   |
///    |    |                   |
///    |    +-------------------+ <--- Frame 0's SP
///    |    | Frame 1           |
///  Grows  |                   |
///  down   |                   |
///    |    | Live GC reference | --+--
///    |    |                   |   |
///    |    |                   |   |
///    V    |                   |   x = offset of live GC reference
///         |                   |   |
///         |                   |   |
///         +-------------------+ --+--  <--- Frame 1's SP
///         | Frame 2           |
///         | ...               |
/// ```
///
/// An individual `StackMap` is associated with just one instruction pointer
/// within the function, contains the size of the stack frame, and represents
/// the stack frame as a bitmap. There is one bit per word in the stack frame,
/// and if the bit is set, then the word contains a live GC reference.
///
/// Note that a caller's `OutgoingArg` stack slots and callee's `IncomingArg`
/// stack slots overlap, so we must choose which function's stack maps record
/// live GC references in these slots. We record the `IncomingArg`s in the
/// callee's stack map.
#[derive(Clone, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "enable-serde", derive(serde::Deserialize, serde::Serialize))]
pub struct StackMap {
    bitmap: Vec<BitSet<Num>>,
    mapped_words: u32,
}

impl StackMap {
    /// Create a vec of Bitsets from a slice of bools.
    pub fn from_slice(vec: &[bool]) -> Self {
        let len = vec.len();
        let num_word = len / NUM_BITS + (len % NUM_BITS != 0) as usize;
        let mut bitmap = Vec::with_capacity(num_word);

        for segment in vec.chunks(NUM_BITS) {
            let mut curr_word = 0;
            for (i, set) in segment.iter().enumerate() {
                if *set {
                    curr_word |= 1 << i;
                }
            }
            bitmap.push(BitSet(curr_word));
        }
        Self {
            mapped_words: len as u32,
            bitmap,
        }
    }

    /// Returns a specified bit.
    pub fn get_bit(&self, bit_index: usize) -> bool {
        assert!(bit_index < NUM_BITS * self.bitmap.len());
        let word_index = bit_index / NUM_BITS;
        let word_offset = (bit_index % NUM_BITS) as u8;
        self.bitmap[word_index].contains(word_offset)
    }

    /// Returns the raw bitmap that represents this stack map.
    pub fn as_slice(&self) -> &[BitSet<u32>] {
        &self.bitmap
    }

    /// Returns the number of words represented by this stack map.
    pub fn mapped_words(&self) -> u32 {
        self.mapped_words
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn stack_maps() {
        let vec: Vec<bool> = Vec::new();
        assert!(StackMap::from_slice(&vec).bitmap.is_empty());

        let mut vec: [bool; NUM_BITS] = Default::default();
        let set_true_idx = [5, 7, 24, 31];

        for &idx in &set_true_idx {
            vec[idx] = true;
        }

        let mut vec = vec.to_vec();
        assert_eq!(
            vec![BitSet::<Num>(2164261024)],
            StackMap::from_slice(&vec).bitmap
        );

        vec.push(false);
        vec.push(true);
        let res = StackMap::from_slice(&vec);
        assert_eq!(
            vec![BitSet::<Num>(2164261024), BitSet::<Num>(2)],
            res.bitmap
        );

        assert!(res.get_bit(5));
        assert!(res.get_bit(31));
        assert!(res.get_bit(33));
        assert!(!res.get_bit(1));
    }
}