Skip to main content

cranpose_core/
snapshot_id_set.rs

1use std::fmt;
2
3/// The identifier of a snapshot. Ids grow in the order that snapshots are created.
4pub type SnapshotId = usize;
5
6const BITS_PER_SET: usize = 64;
7const SNAPSHOT_ID_SIZE: usize = 64;
8
9/// A compact set of snapshot ids, as Jetpack Compose's `SnapshotIdSet`. Ids close to the set's
10/// lower bound take one bit each.
11#[derive(Clone, PartialEq, Eq)]
12pub struct SnapshotIdSet {
13    upper_set: u64,
14    lower_set: u64,
15    lower_bound: SnapshotId,
16    below_bound: Option<Box<[SnapshotId]>>,
17}
18
19impl SnapshotIdSet {
20    /// Empty snapshot ID set.
21    pub const EMPTY: SnapshotIdSet = SnapshotIdSet {
22        upper_set: 0,
23        lower_set: 0,
24        lower_bound: 0,
25        below_bound: None,
26    };
27
28    /// Create a new empty snapshot ID set.
29    pub fn new() -> Self {
30        Self::EMPTY
31    }
32
33    /// Check if a snapshot ID is in the set.
34    pub fn get(&self, id: SnapshotId) -> bool {
35        let offset = id.wrapping_sub(self.lower_bound);
36
37        if offset < BITS_PER_SET {
38            let mask = 1u64 << offset;
39            (self.lower_set & mask) != 0
40        } else if offset < BITS_PER_SET * 2 {
41            let mask = 1u64 << (offset - BITS_PER_SET);
42            (self.upper_set & mask) != 0
43        } else if id > self.lower_bound {
44            false
45        } else {
46            self.below_bound
47                .as_ref()
48                .is_some_and(|arr| arr.binary_search(&id).is_ok())
49        }
50    }
51
52    /// Add a snapshot ID to the set (returns a new set if modified).
53    pub fn set(&self, id: SnapshotId) -> Self {
54        if id < self.lower_bound {
55            if let Some(ref arr) = self.below_bound {
56                match arr.binary_search(&id) {
57                    Ok(_) => {
58                        return self.clone();
59                    }
60                    Err(insert_pos) => {
61                        let mut new_arr = Vec::with_capacity(arr.len() + 1);
62                        new_arr.extend_from_slice(&arr[..insert_pos]);
63                        new_arr.push(id);
64                        new_arr.extend_from_slice(&arr[insert_pos..]);
65                        return Self {
66                            upper_set: self.upper_set,
67                            lower_set: self.lower_set,
68                            lower_bound: self.lower_bound,
69                            below_bound: Some(new_arr.into_boxed_slice()),
70                        };
71                    }
72                }
73            } else {
74                return Self {
75                    upper_set: self.upper_set,
76                    lower_set: self.lower_set,
77                    lower_bound: self.lower_bound,
78                    below_bound: Some(vec![id].into_boxed_slice()),
79                };
80            }
81        }
82
83        let offset = id - self.lower_bound;
84
85        if offset < BITS_PER_SET {
86            let mask = 1u64 << offset;
87            if (self.lower_set & mask) == 0 {
88                return Self {
89                    upper_set: self.upper_set,
90                    lower_set: self.lower_set | mask,
91                    lower_bound: self.lower_bound,
92                    below_bound: self.below_bound.clone(),
93                };
94            }
95        } else if offset < BITS_PER_SET * 2 {
96            let mask = 1u64 << (offset - BITS_PER_SET);
97            if (self.upper_set & mask) == 0 {
98                return Self {
99                    upper_set: self.upper_set | mask,
100                    lower_set: self.lower_set,
101                    lower_bound: self.lower_bound,
102                    below_bound: self.below_bound.clone(),
103                };
104            }
105        } else if offset >= BITS_PER_SET * 2 && !self.get(id) {
106            return self.shift_and_set(id);
107        }
108
109        self.clone()
110    }
111
112    /// Remove a snapshot ID from the set (returns a new set if modified).
113    pub fn clear(&self, id: SnapshotId) -> Self {
114        let offset = id.wrapping_sub(self.lower_bound);
115
116        if offset < BITS_PER_SET {
117            let mask = 1u64 << offset;
118            if (self.lower_set & mask) != 0 {
119                return Self {
120                    upper_set: self.upper_set,
121                    lower_set: self.lower_set & !mask,
122                    lower_bound: self.lower_bound,
123                    below_bound: self.below_bound.clone(),
124                };
125            }
126        } else if offset < BITS_PER_SET * 2 {
127            let mask = 1u64 << (offset - BITS_PER_SET);
128            if (self.upper_set & mask) != 0 {
129                return Self {
130                    upper_set: self.upper_set & !mask,
131                    lower_set: self.lower_set,
132                    lower_bound: self.lower_bound,
133                    below_bound: self.below_bound.clone(),
134                };
135            }
136        } else if id < self.lower_bound
137            && let Some(ref arr) = self.below_bound
138            && let Ok(pos) = arr.binary_search(&id)
139        {
140            let mut new_arr = Vec::with_capacity(arr.len() - 1);
141            new_arr.extend_from_slice(&arr[..pos]);
142            new_arr.extend_from_slice(&arr[pos + 1..]);
143            return Self {
144                upper_set: self.upper_set,
145                lower_set: self.lower_set,
146                lower_bound: self.lower_bound,
147                below_bound: if new_arr.is_empty() {
148                    None
149                } else {
150                    Some(new_arr.into_boxed_slice())
151                },
152            };
153        }
154
155        self.clone()
156    }
157
158    /// Remove all IDs in `other` from this set (a & ~b).
159    pub fn and_not(&self, other: &Self) -> Self {
160        if other.is_empty() {
161            return self.clone();
162        }
163        if self.is_empty() {
164            return Self::EMPTY;
165        }
166
167        if self.lower_bound == other.lower_bound && self.below_bound_equals(&other.below_bound) {
168            return Self {
169                upper_set: self.upper_set & !other.upper_set,
170                lower_set: self.lower_set & !other.lower_set,
171                lower_bound: self.lower_bound,
172                below_bound: self.below_bound.clone(),
173            };
174        }
175
176        let mut result = self.clone();
177        for id in other.iter() {
178            result = result.clear(id);
179        }
180        result
181    }
182
183    /// Union this set with another (a | b).
184    pub fn or(&self, other: &Self) -> Self {
185        if other.is_empty() {
186            return self.clone();
187        }
188        if self.is_empty() {
189            return other.clone();
190        }
191
192        if self.lower_bound == other.lower_bound && self.below_bound_equals(&other.below_bound) {
193            return Self {
194                upper_set: self.upper_set | other.upper_set,
195                lower_set: self.lower_set | other.lower_set,
196                lower_bound: self.lower_bound,
197                below_bound: self.below_bound.clone(),
198            };
199        }
200
201        let mut result = self.clone();
202        for id in other.iter() {
203            result = result.set(id);
204        }
205        result
206    }
207
208    /// Find the lowest snapshot ID in the set that is <= upper.
209    pub fn lowest(&self, upper: SnapshotId) -> SnapshotId {
210        if let Some(ref arr) = self.below_bound
211            && let Some(&lowest) = arr.first()
212            && lowest <= upper
213        {
214            return lowest;
215        }
216
217        if self.lower_set != 0 {
218            let lowest_in_lower = self.lower_bound + self.lower_set.trailing_zeros() as usize;
219            if lowest_in_lower <= upper {
220                return lowest_in_lower;
221            }
222        }
223
224        if self.upper_set != 0 {
225            let lowest_in_upper =
226                self.lower_bound + BITS_PER_SET + self.upper_set.trailing_zeros() as usize;
227            if lowest_in_upper <= upper {
228                return lowest_in_upper;
229            }
230        }
231
232        upper
233    }
234
235    /// Check if the set is empty.
236    pub fn is_empty(&self) -> bool {
237        self.lower_set == 0 && self.upper_set == 0 && self.below_bound.is_none()
238    }
239
240    /// Iterate over all snapshot IDs in the set.
241    pub fn iter(&self) -> SnapshotIdSetIter<'_> {
242        SnapshotIdSetIter::new(self)
243    }
244
245    /// Convert to a Vec of snapshot IDs (for testing/debugging).
246    pub fn to_list(&self) -> Vec<SnapshotId> {
247        self.iter().collect()
248    }
249
250    /// Add a contiguous range of IDs [from, until) to the set.
251    /// Mirrors AndroidX SnapshotIdSet.addRange semantics used by Snapshot.kt.
252    pub fn add_range(&self, from: SnapshotId, until: SnapshotId) -> Self {
253        if from >= until {
254            return self.clone();
255        }
256        let mut result = self.clone();
257        let mut id = from;
258        while id < until {
259            result = result.set(id);
260            id += 1;
261        }
262        result
263    }
264
265    fn below_bound_equals(&self, other: &Option<Box<[SnapshotId]>>) -> bool {
266        match (&self.below_bound, other) {
267            (None, None) => true,
268            (Some(a), Some(b)) => a == b,
269            _ => false,
270        }
271    }
272
273    fn shift_and_set(&self, id: SnapshotId) -> Self {
274        let target_lower_bound = (id / SNAPSHOT_ID_SIZE) * SNAPSHOT_ID_SIZE;
275
276        let mut new_upper_set = self.upper_set;
277        let mut new_lower_set = self.lower_set;
278        let mut new_lower_bound = self.lower_bound;
279        let mut new_below_bound: Vec<SnapshotId> = if let Some(ref arr) = self.below_bound {
280            arr.to_vec()
281        } else {
282            Vec::new()
283        };
284
285        while new_lower_bound < target_lower_bound {
286            if new_lower_set != 0 {
287                for bit_offset in 0..BITS_PER_SET {
288                    if (new_lower_set & (1u64 << bit_offset)) != 0 {
289                        let id_to_add = new_lower_bound + bit_offset;
290                        match new_below_bound.binary_search(&id_to_add) {
291                            Ok(_) => {}
292                            Err(pos) => new_below_bound.insert(pos, id_to_add),
293                        }
294                    }
295                }
296            }
297
298            if new_upper_set == 0 {
299                new_lower_bound = target_lower_bound;
300                new_lower_set = 0;
301                break;
302            }
303
304            new_lower_set = new_upper_set;
305            new_upper_set = 0;
306            new_lower_bound += BITS_PER_SET;
307        }
308
309        let result = Self {
310            upper_set: new_upper_set,
311            lower_set: new_lower_set,
312            lower_bound: new_lower_bound,
313            below_bound: if new_below_bound.is_empty() {
314                None
315            } else {
316                Some(new_below_bound.into_boxed_slice())
317            },
318        };
319
320        result.set(id)
321    }
322}
323
324impl Default for SnapshotIdSet {
325    fn default() -> Self {
326        Self::EMPTY
327    }
328}
329
330impl fmt::Debug for SnapshotIdSet {
331    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
332        write!(f, "SnapshotIdSet{{")?;
333        let ids: Vec<_> = self.iter().collect();
334        for (i, id) in ids.iter().enumerate() {
335            if i > 0 {
336                write!(f, ", ")?;
337            }
338            write!(f, "{id}")?;
339        }
340        write!(f, "}}")
341    }
342}
343
344/// Iterator over snapshot IDs in a set.
345pub struct SnapshotIdSetIter<'a> {
346    set: &'a SnapshotIdSet,
347    below_index: usize,
348    lower_set: u64,
349    upper_set: u64,
350    current_offset: usize,
351}
352
353impl<'a> SnapshotIdSetIter<'a> {
354    fn new(set: &'a SnapshotIdSet) -> Self {
355        Self {
356            set,
357            below_index: 0,
358            lower_set: set.lower_set,
359            upper_set: set.upper_set,
360            current_offset: 0,
361        }
362    }
363}
364
365impl Iterator for SnapshotIdSetIter<'_> {
366    type Item = SnapshotId;
367
368    fn next(&mut self) -> Option<Self::Item> {
369        if let Some(ref arr) = self.set.below_bound
370            && self.below_index < arr.len()
371        {
372            let id = arr[self.below_index];
373            self.below_index += 1;
374            return Some(id);
375        }
376
377        while self.current_offset < BITS_PER_SET {
378            if (self.lower_set & (1u64 << self.current_offset)) != 0 {
379                let id = self.set.lower_bound + self.current_offset;
380                self.current_offset += 1;
381                return Some(id);
382            }
383            self.current_offset += 1;
384        }
385
386        while self.current_offset < BITS_PER_SET * 2 {
387            let bit_offset = self.current_offset - BITS_PER_SET;
388            if (self.upper_set & (1u64 << bit_offset)) != 0 {
389                let id = self.set.lower_bound + self.current_offset;
390                self.current_offset += 1;
391                return Some(id);
392            }
393            self.current_offset += 1;
394        }
395
396        None
397    }
398}
399
400#[cfg(test)]
401#[path = "tests/snapshot_id_set_tests.rs"]
402mod tests;