regalloc2/fastalloc/
vregset.rs

1use core::fmt;
2
3use crate::ion::data_structures::VRegIndex;
4use crate::VReg;
5use alloc::vec;
6use alloc::vec::Vec;
7
8#[derive(Clone)]
9struct VRegNode {
10    next: VRegIndex,
11    prev: VRegIndex,
12    vreg: VReg,
13}
14
15// Using a doubly linked list here for fast insertion,
16// removal and iteration.
17pub struct VRegSet {
18    items: Vec<VRegNode>,
19    head: VRegIndex,
20}
21
22impl VRegSet {
23    pub fn with_capacity(num_vregs: usize) -> Self {
24        Self {
25            items: vec![
26                VRegNode {
27                    prev: VRegIndex::new(num_vregs),
28                    next: VRegIndex::new(num_vregs),
29                    vreg: VReg::invalid()
30                };
31                num_vregs + 1
32            ],
33            head: VRegIndex::new(num_vregs),
34        }
35    }
36
37    pub fn insert(&mut self, vreg: VReg) {
38        debug_assert_eq!(self.items[vreg.vreg()].vreg, VReg::invalid());
39        let old_head_next = self.items[self.head.index()].next;
40        self.items[vreg.vreg()] = VRegNode {
41            next: old_head_next,
42            prev: self.head,
43            vreg,
44        };
45        self.items[self.head.index()].next = VRegIndex::new(vreg.vreg());
46        self.items[old_head_next.index()].prev = VRegIndex::new(vreg.vreg());
47    }
48
49    pub fn remove(&mut self, vreg_num: usize) {
50        let prev = self.items[vreg_num].prev;
51        let next = self.items[vreg_num].next;
52        self.items[prev.index()].next = next;
53        self.items[next.index()].prev = prev;
54        self.items[vreg_num].vreg = VReg::invalid();
55    }
56
57    pub fn is_empty(&self) -> bool {
58        self.items[self.head.index()].next == self.head
59    }
60
61    pub fn iter(&self) -> VRegSetIter<'_> {
62        VRegSetIter {
63            curr_item: self.items[self.head.index()].next,
64            head: self.head,
65            items: &self.items,
66        }
67    }
68}
69
70pub struct VRegSetIter<'a> {
71    curr_item: VRegIndex,
72    head: VRegIndex,
73    items: &'a [VRegNode],
74}
75
76impl<'a> Iterator for VRegSetIter<'a> {
77    type Item = VReg;
78
79    fn next(&mut self) -> Option<Self::Item> {
80        if self.curr_item != self.head {
81            let item = self.items[self.curr_item.index()].clone();
82            self.curr_item = item.next;
83            Some(item.vreg)
84        } else {
85            None
86        }
87    }
88}
89
90impl fmt::Debug for VRegSet {
91    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
92        write!(f, "{{ ")?;
93        for vreg in self.iter() {
94            write!(f, "{vreg} ")?;
95        }
96        write!(f, "}}")
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103    use crate::RegClass;
104    use RegClass::*;
105    const VREG: fn(usize, RegClass) -> VReg = VReg::new;
106
107    #[test]
108    fn operations() {
109        let mut set = VRegSet::with_capacity(3090);
110        assert!(set.is_empty());
111        set.insert(VREG(10, Int));
112        set.insert(VREG(2000, Int));
113        set.insert(VREG(11, Vector));
114        set.insert(VREG(199, Float));
115        set.insert(VREG(23, Int));
116        let mut iter = set.iter();
117        assert_eq!(iter.next(), Some(VREG(23, Int)));
118        assert_eq!(iter.next(), Some(VREG(199, Float)));
119        assert_eq!(iter.next(), Some(VREG(11, Vector)));
120        assert_eq!(iter.next(), Some(VREG(2000, Int)));
121        assert_eq!(iter.next(), Some(VREG(10, Int)));
122
123        set.remove(23);
124        set.remove(11);
125        set.insert(VREG(73, Vector));
126        let mut iter = set.iter();
127        assert_eq!(iter.next(), Some(VREG(73, Vector)));
128        assert_eq!(iter.next(), Some(VREG(199, Float)));
129        assert_eq!(iter.next(), Some(VREG(2000, Int)));
130        assert_eq!(iter.next(), Some(VREG(10, Int)));
131        assert!(!set.is_empty());
132    }
133
134    #[test]
135    fn empty() {
136        let mut set = VRegSet::with_capacity(2000);
137        assert!(set.is_empty());
138        set.insert(VREG(100, Int));
139        assert!(!set.is_empty());
140        set.remove(100);
141        assert!(set.is_empty());
142    }
143}