celox_backend_common/regalloc/
stack_color.rs

1//! Opcode-free stack-slot coloring over exact sparse live intervals.
2
3use std::cmp::Reverse;
4use std::collections::{BTreeMap, BTreeSet, BinaryHeap};
5use std::fmt;
6
7use super::LiveInterval;
8
9/// Failure while assigning target-owned spill homes to reusable frame slots.
10#[derive(Debug, Clone, PartialEq, Eq)]
11pub enum StackColorError<V> {
12    DuplicateValue(V),
13    EmptyInterval(V),
14    SlotCountOverflow,
15}
16
17impl<V: fmt::Debug> fmt::Display for StackColorError<V> {
18    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
19        match self {
20            Self::DuplicateValue(value) => {
21                write!(formatter, "stack interval for {value:?} was colored twice")
22            }
23            Self::EmptyInterval(value) => {
24                write!(
25                    formatter,
26                    "stack interval for {value:?} has no live segments"
27                )
28            }
29            Self::SlotCountOverflow => formatter.write_str("stack-slot count exceeds u32"),
30        }
31    }
32}
33
34impl<V: fmt::Debug> std::error::Error for StackColorError<V> {}
35
36/// Deterministic stack-slot assignment for target-owned spill homes.
37///
38/// Backends retain responsibility for choosing spill values and translating a
39/// slot number into their frame layout. This helper conservatively projects
40/// each exact sparse interval to a linear block-order envelope, then reuses
41/// slots with a sweep. Envelope overlap may miss a legal reuse across a sparse
42/// gap, but cannot make interfering homes alias. The algorithm is independent
43/// of target MIR and opcode semantics and runs in O(n log n).
44#[derive(Debug, Clone, PartialEq, Eq)]
45pub struct StackSlotColoring<V> {
46    assignments: BTreeMap<V, u32>,
47    slot_count: u32,
48}
49
50impl<V> StackSlotColoring<V>
51where
52    V: Ord,
53{
54    pub fn get(&self, value: &V) -> Option<u32> {
55        self.assignments.get(value).copied()
56    }
57
58    pub fn slot_count(&self) -> u32 {
59        self.slot_count
60    }
61}
62
63/// Color a complete target spill batch using conservative linear envelopes.
64pub fn color_stack_slots<'a, V, I>(intervals: I) -> Result<StackSlotColoring<V>, StackColorError<V>>
65where
66    V: 'a + Copy + Ord,
67    I: IntoIterator<Item = &'a LiveInterval<V>>,
68{
69    let mut seen = BTreeSet::new();
70    let mut ordered = Vec::new();
71    for interval in intervals {
72        if !seen.insert(interval.value) {
73            return Err(StackColorError::DuplicateValue(interval.value));
74        }
75        let Some(first) = interval.segments.first() else {
76            return Err(StackColorError::EmptyInterval(interval.value));
77        };
78        let last = interval
79            .segments
80            .last()
81            .expect("a nonempty interval has a last segment");
82        ordered.push((
83            (first.block, first.start),
84            (last.block, last.end),
85            interval.value,
86        ));
87    }
88    ordered.sort_unstable();
89
90    let mut active = BinaryHeap::<Reverse<((usize, u64), u32, V)>>::new();
91    let mut available = BinaryHeap::<Reverse<u32>>::new();
92    let mut next_slot = 0_u32;
93    let mut assignments = BTreeMap::new();
94    for (start, end, value) in ordered {
95        while active
96            .peek()
97            .is_some_and(|Reverse((active_end, _, _))| *active_end <= start)
98        {
99            let Reverse((_, slot, _)) = active.pop().expect("peeked active interval exists");
100            available.push(Reverse(slot));
101        }
102        let slot = if let Some(Reverse(slot)) = available.pop() {
103            slot
104        } else {
105            let slot = next_slot;
106            next_slot = next_slot
107                .checked_add(1)
108                .ok_or(StackColorError::SlotCountOverflow)?;
109            slot
110        };
111        assignments.insert(value, slot);
112        active.push(Reverse((end, slot, value)));
113    }
114    Ok(StackSlotColoring {
115        assignments,
116        slot_count: next_slot,
117    })
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123    use crate::regalloc::LiveSegment;
124
125    fn interval(value: u32, segments: &[(usize, u64, u64)]) -> LiveInterval<u32> {
126        LiveInterval {
127            value,
128            segments: segments
129                .iter()
130                .map(|&(block, start, end)| LiveSegment { block, start, end })
131                .collect(),
132        }
133    }
134
135    #[test]
136    fn reuses_a_slot_for_disjoint_sparse_intervals() {
137        let intervals = [
138            interval(0, &[(0, 0, 4)]),
139            interval(1, &[(0, 4, 8)]),
140            interval(2, &[(1, 0, 8)]),
141        ];
142        let coloring = color_stack_slots(&intervals).unwrap();
143
144        assert_eq!(coloring.get(&0), Some(0));
145        assert_eq!(coloring.get(&1), Some(0));
146        assert_eq!(coloring.get(&2), Some(0));
147        assert_eq!(coloring.slot_count(), 1);
148    }
149
150    #[test]
151    fn separates_any_pair_with_an_overlapping_segment() {
152        let intervals = [
153            interval(0, &[(0, 0, 2), (2, 0, 8)]),
154            interval(1, &[(1, 0, 2), (2, 7, 9)]),
155        ];
156        let coloring = color_stack_slots(&intervals).unwrap();
157
158        assert_eq!(coloring.get(&0), Some(0));
159        assert_eq!(coloring.get(&1), Some(1));
160    }
161
162    #[test]
163    fn rejects_duplicate_values() {
164        let value = interval(7, &[(0, 0, 1)]);
165
166        assert_eq!(
167            color_stack_slots([&value, &value]),
168            Err(StackColorError::DuplicateValue(7))
169        );
170    }
171}
celox_backend_common/regalloc/stack_color.rs

celox_backend_common/regalloc/
stack_color.rs
