1use smallvec::SmallVec;
4
5use core::{fmt::Debug, iter, ops::Range};
6
7#[derive(Clone, Debug, PartialEq)]
11pub(crate) struct RangedStates<I, T> {
12 ranges: SmallVec<[(Range<I>, T); 1]>,
15}
16
17impl<I: Copy + Ord, T: Copy + PartialEq> RangedStates<I, T> {
18 pub fn from_range(range: Range<I>, value: T) -> Self {
19 Self {
20 ranges: iter::once((range, value)).collect(),
21 }
22 }
23
24 #[cfg(test)]
26 pub fn from_slice(values: &[(Range<I>, T)]) -> Self {
27 Self {
28 ranges: values.iter().cloned().collect(),
29 }
30 }
31
32 pub fn iter(&self) -> impl Iterator<Item = &(Range<I>, T)> + Clone {
33 self.ranges.iter()
34 }
35
36 pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut (Range<I>, T)> {
37 self.ranges.iter_mut()
38 }
39
40 #[cfg(test)]
43 fn check_sanity(&self) {
44 for a in self.ranges.iter() {
45 assert!(a.0.start < a.0.end);
46 }
47 for (a, b) in self.ranges.iter().zip(self.ranges[1..].iter()) {
48 assert!(a.0.end <= b.0.start);
49 }
50 }
51
52 pub fn coalesce(&mut self) {
54 let mut num_removed = 0;
55 let mut iter = self.ranges.iter_mut();
56 let mut cur = match iter.next() {
57 Some(elem) => elem,
58 None => return,
59 };
60 for next in iter {
61 if cur.0.end == next.0.start && cur.1 == next.1 {
62 num_removed += 1;
63 cur.0.end = next.0.end;
64 next.0.end = next.0.start;
65 } else {
66 cur = next;
67 }
68 }
69 if num_removed != 0 {
70 self.ranges.retain(|pair| pair.0.start != pair.0.end);
71 }
72 }
73
74 pub fn iter_filter<'a>(
75 &'a self,
76 range: &'a Range<I>,
77 ) -> impl Iterator<Item = (Range<I>, &'a T)> + 'a {
78 self.ranges
79 .iter()
80 .filter(move |&(inner, ..)| inner.end > range.start && inner.start < range.end)
81 .map(move |(inner, v)| {
82 let new_range = inner.start.max(range.start)..inner.end.min(range.end);
83
84 (new_range, v)
85 })
86 }
87
88 pub fn isolate(&mut self, index: &Range<I>, default: T) -> &mut [(Range<I>, T)] {
93 let mut start_pos = match self.ranges.iter().position(|pair| pair.0.end > index.start) {
98 Some(pos) => pos,
99 None => {
100 let pos = self.ranges.len();
101 self.ranges.push((index.clone(), default));
102 return &mut self.ranges[pos..];
103 }
104 };
105
106 {
107 let (range, value) = self.ranges[start_pos].clone();
108 if range.start < index.start {
109 self.ranges[start_pos].0.start = index.start;
110 self.ranges
111 .insert(start_pos, (range.start..index.start, value));
112 start_pos += 1;
113 }
114 }
115 let mut pos = start_pos;
116 let mut range_pos = index.start;
117 loop {
118 let (range, value) = self.ranges[pos].clone();
119 if range.start >= index.end {
120 self.ranges.insert(pos, (range_pos..index.end, default));
121 pos += 1;
122 break;
123 }
124 if range.start > range_pos {
125 self.ranges.insert(pos, (range_pos..range.start, default));
126 pos += 1;
127 range_pos = range.start;
128 }
129 if range.end >= index.end {
130 if range.end != index.end {
131 self.ranges[pos].0.start = index.end;
132 self.ranges.insert(pos, (range_pos..index.end, value));
133 }
134 pos += 1;
135 break;
136 }
137 pos += 1;
138 range_pos = range.end;
139 if pos == self.ranges.len() {
140 self.ranges.push((range_pos..index.end, default));
141 pos += 1;
142 break;
143 }
144 }
145
146 &mut self.ranges[start_pos..pos]
147 }
148
149 #[cfg(test)]
151 pub fn sanely_isolated(&self, index: Range<I>, default: T) -> alloc::vec::Vec<(Range<I>, T)> {
152 let mut clone = self.clone();
153 let result = clone.isolate(&index, default).to_vec();
154 clone.check_sanity();
155 result
156 }
157}
158
159#[cfg(test)]
160mod test {
161 use super::RangedStates;
163
164 #[test]
165 fn sane_good() {
166 let rs = RangedStates::from_slice(&[(1..4, 9u8), (4..5, 9)]);
167 rs.check_sanity();
168 }
169
170 #[test]
171 #[should_panic]
172 fn sane_empty() {
173 let rs = RangedStates::from_slice(&[(1..4, 9u8), (5..5, 9)]);
174 rs.check_sanity();
175 }
176
177 #[test]
178 #[should_panic]
179 fn sane_intersect() {
180 let rs = RangedStates::from_slice(&[(1..4, 9u8), (3..5, 9)]);
181 rs.check_sanity();
182 }
183
184 #[test]
185 fn coalesce() {
186 let mut rs = RangedStates::from_slice(&[(1..4, 9u8), (4..5, 9), (5..7, 1), (8..9, 1)]);
187 rs.coalesce();
188 rs.check_sanity();
189 assert_eq!(rs.ranges.as_slice(), &[(1..5, 9), (5..7, 1), (8..9, 1),]);
190 }
191
192 #[test]
193 fn isolate() {
194 let rs = RangedStates::from_slice(&[(1..4, 9u8), (4..5, 9), (5..7, 1), (8..9, 1)]);
195 assert_eq!(&rs.sanely_isolated(4..5, 0), &[(4..5, 9u8),]);
196 assert_eq!(
197 &rs.sanely_isolated(0..6, 0),
198 &[(0..1, 0), (1..4, 9u8), (4..5, 9), (5..6, 1),]
199 );
200 assert_eq!(&rs.sanely_isolated(8..10, 1), &[(8..9, 1), (9..10, 1),]);
201 assert_eq!(
202 &rs.sanely_isolated(6..9, 0),
203 &[(6..7, 1), (7..8, 0), (8..9, 1),]
204 );
205 }
206}