ostd/util/
id_set.rs

1// SPDX-License-Identifier: MPL-2.0
2
3//! A fixed-size set of unique IDs.
4//!
5//! This module introduces two abstract data collection types.
6//! The first one is [`IdSet<I>`],
7//! a set that contains at most `I::cardinality()` items,
8//! each of which represents a unique ID of type `I`.
9//! The second one is [`AtomicIdSet<I>`],
10//! the atomic version of `IdSet<I>`.
11//!
12//! The type parameter `I` implements the [`Id`] trait,
13//! which abstracts any ID type whose instances
14//! can be 1:1 mapped to the integers from 0 to `Id::cardinality()` (exclusive).
15//! The ID type is required to implement `Into<u32>` and `TryFrom<u32>`.
16//!
17//! One use case of `IdSet<I>` inside OSTD is
18//! to implement a set of CPU IDs.
19//! [`crate::cpu::CpuSet`] has been simply defined
20//! as a typa alias for `IdSet<CpuId>`.
21
22use core::{
23    fmt::Debug,
24    marker::PhantomData,
25    ops::{Bound, Range, RangeFrom, RangeFull, RangeTo, RangeToInclusive},
26    sync::atomic::{AtomicU64, Ordering},
27};
28
29use bitvec::{order::Lsb0, view::BitView};
30use smallvec::SmallVec;
31
32use crate::const_assert;
33
34/// A trait to abstract an ID type.
35///
36/// # Safety
37///
38/// There must be a 1:1 mapping between this ID type and
39/// the integers from 0 to `Self::cardinality()` (exclusive).
40/// This implies that the implementation must ensure that
41/// if one invokes `Into::<u32>::into()` for an `Id` value,
42/// then the returned integer always falls within `0..Id::cardinality()`.
43/// Furthermore, the implementation must ensure that
44/// for any `id_value` of type `MyId: Id`,
45/// the following assertion always succeed
46///
47/// ```rust
48/// assert!(id_value == MyId::new(id_value.into()));
49/// ```
50///
51/// There are also constraints on the implementation of `self::cardinality`.
52/// For one thing, the cardinality of the ID type must not change, i.e.,
53/// different calls to `self::cardinality` return the same value.
54/// In addition, the value of a cardinality must be greater than zero.
55pub unsafe trait Id: Copy + Clone + Debug + Eq + Into<u32> + PartialEq {
56    /// Creates an ID instance given a raw ID number.
57    ///
58    /// # Panics
59    ///
60    /// The given number must be less than `Self::cardinality()`.
61    /// Otherwise, this method would panic.
62    fn new(raw_id: u32) -> Self {
63        assert!(raw_id < Self::cardinality());
64        // SAFETY: The raw ID is a valid one.
65        unsafe { Self::new_unchecked(raw_id) }
66    }
67
68    /// Creates an ID instance given a raw ID number.
69    ///
70    /// # Safety
71    ///
72    /// The given number must be less than `Self::cardinality()`.
73    unsafe fn new_unchecked(raw_id: u32) -> Self;
74
75    /// The number of unique IDs representable by this type.
76    fn cardinality() -> u32;
77
78    /// Returns an [`usize`] from the [`Id`]'s corresponding [`u32`].
79    fn as_usize(self) -> usize {
80        Into::<u32>::into(self) as usize
81    }
82}
83
84/// A set of IDs.
85///
86/// # Examples
87///
88/// Assume that you have a type named `MyId`,
89/// which represents a set of IDs from 0 to 10 (exclusive).
90///
91/// ```
92/// #[derive(Copy, Clone, Debug, Eq, PartialEq)]
93/// pub struct MyId(u32);
94///
95/// // SAFETY: `MyId` maintains the 1:1 mapping invariant for 0..10.
96/// unsafe impl Id for MyId {
97///     fn new_unchecked(raw_id: u32) -> Self { Self(raw_id) }
98///     fn cardinality() -> u32 { 10 } // Fixed cardinality for this example
99/// }
100/// impl From<MyId> for u32 { fn from(id: MyId) -> u32 { id.0 } }
101/// ```
102///
103/// Now you can use `IdSet<MyId>` as a container for `MyID`s.
104///
105/// ```
106/// let mut my_id_set: IdSet<MyId> = IdSet::new_empty();
107///
108/// let id0 = MyId::new(0);
109/// my_id_set.add(id0);
110/// assert!(my_id_set.contains(id0));
111/// assert_eq!(my_id_set.count(), 1);
112///
113/// let id5 = MyId::new(5);
114/// my_id_set.add(id5);
115/// assert!(my_id_set.contains(id5));
116/// assert_eq!(my_id_set.count(), 2);
117///
118/// my_id_set.remove(id0);
119/// assert!(!my_id_set.contains(id0));
120/// assert_eq!(my_id_set.count(), 1);
121///
122/// let full_set = IdSet::<MyId>::new_full();
123/// assert_eq!(full_set.count(), 10);
124/// assert!(full_set.contains(MyId::new(9)));
125/// ```
126#[derive(Clone, Debug, Eq, PartialEq)]
127pub struct IdSet<I> {
128    bits: SmallVec<[InnerPart; NR_PARTS_NO_ALLOC]>,
129    phantom: PhantomData<I>,
130}
131
132type InnerPart = u64;
133
134const BITS_PER_PART: usize = InnerPart::BITS as usize;
135const NR_PARTS_NO_ALLOC: usize = 2;
136
137fn part_idx<I: Id>(id: I) -> usize {
138    (id.into() as usize) / BITS_PER_PART
139}
140
141fn bit_idx<I: Id>(id: I) -> usize {
142    (id.into() as usize) % BITS_PER_PART
143}
144
145fn parts_for_ids<I: Id>() -> usize {
146    (I::cardinality() as usize).div_ceil(BITS_PER_PART)
147}
148
149impl<I: Id> IdSet<I> {
150    /// Creates a new `IdSet` with all IDs in the system.
151    pub fn new_full() -> Self {
152        let mut bits = Self::with_bit_pattern(!0);
153        Self::clear_invalid_id_bits(&mut bits);
154        Self {
155            bits,
156            phantom: PhantomData,
157        }
158    }
159
160    /// Creates a new `IdSet` with no IDs in the system.
161    pub fn new_empty() -> Self {
162        let bits = Self::with_bit_pattern(0);
163        Self {
164            bits,
165            phantom: PhantomData,
166        }
167    }
168
169    /// Creates a new bitmap with each of its parts filled with the given bits.
170    ///
171    /// Depending on the bit pattern and the number of IDs,
172    /// the resulting bitmap may end with some invalid bits.
173    fn with_bit_pattern(part_bits: InnerPart) -> SmallVec<[InnerPart; NR_PARTS_NO_ALLOC]> {
174        let num_parts = parts_for_ids::<I>();
175        let mut bits = SmallVec::with_capacity(num_parts);
176        bits.resize(num_parts, part_bits);
177        bits
178    }
179
180    fn clear_invalid_id_bits(bits: &mut SmallVec<[InnerPart; NR_PARTS_NO_ALLOC]>) {
181        let num_ids = I::cardinality() as usize;
182        if !num_ids.is_multiple_of(BITS_PER_PART) {
183            let num_parts = parts_for_ids::<I>();
184            bits[num_parts - 1] &= (1 << (num_ids % BITS_PER_PART)) - 1;
185        }
186    }
187
188    /// Adds an ID to the set.
189    pub fn add(&mut self, id: I) {
190        let part_idx = part_idx(id);
191        let bit_idx = bit_idx(id);
192        if part_idx >= self.bits.len() {
193            self.bits.resize(part_idx + 1, 0);
194        }
195        self.bits[part_idx] |= 1 << bit_idx;
196    }
197
198    /// Removes an ID from the set.
199    pub fn remove(&mut self, id: I) {
200        let part_idx = part_idx(id);
201        let bit_idx = bit_idx(id);
202        if part_idx < self.bits.len() {
203            self.bits[part_idx] &= !(1 << bit_idx);
204        }
205    }
206
207    /// Returns true if the set contains the specified ID.
208    pub fn contains(&self, id: I) -> bool {
209        let part_idx = part_idx(id);
210        let bit_idx = bit_idx(id);
211        part_idx < self.bits.len() && (self.bits[part_idx] & (1 << bit_idx)) != 0
212    }
213
214    /// Returns the number of IDs in the set.
215    pub fn count(&self) -> usize {
216        self.bits
217            .iter()
218            .map(|part| part.count_ones() as usize)
219            .sum()
220    }
221
222    /// Returns true if the set is empty.
223    pub fn is_empty(&self) -> bool {
224        self.bits.iter().all(|part| *part == 0)
225    }
226
227    /// Returns true if the set is full.
228    pub fn is_full(&self) -> bool {
229        let num_ids = I::cardinality() as usize;
230        self.bits.iter().enumerate().all(|(idx, part)| {
231            if idx == self.bits.len() - 1 && !num_ids.is_multiple_of(BITS_PER_PART) {
232                *part == (1 << (num_ids % BITS_PER_PART)) - 1
233            } else {
234                *part == !0
235            }
236        })
237    }
238
239    /// Adds all IDs to the set.
240    pub fn add_all(&mut self) {
241        self.bits.fill(!0);
242        Self::clear_invalid_id_bits(&mut self.bits);
243    }
244
245    /// Removes all IDs from the set.
246    pub fn clear(&mut self) {
247        self.bits.fill(0);
248    }
249
250    /// Iterates over all IDs in the set.
251    ///
252    /// The iteration is guaranteed to be in ascending order.
253    #[inline]
254    pub fn iter(&self) -> impl Iterator<Item = I> + '_ {
255        self.iter_in(..)
256    }
257
258    /// Iterates over the IDs in the set within the specified range.
259    ///
260    /// The iteration is guaranteed to be in ascending order.
261    /// Only IDs that are both in the set and within the specified range will be returned.
262    pub fn iter_in<S: IdSetSlicer<I>>(&self, slicer: S) -> impl Iterator<Item = I> + '_ {
263        let (start, end) = slicer.to_range_bounds();
264
265        self.bits.view_bits::<Lsb0>()[start..end]
266            .iter_ones()
267            .map(move |offset| {
268                // SAFETY: `offset` is relative to the slice `[start..end]`,
269                // therefore `start + offset` is the absolute index of the bit.
270                // Since `offset` only iterates over relative positions of bit 1s, the
271                // resulting absolute index must refer to an active bit in `self.bits`.
272                unsafe { I::new_unchecked((start + offset) as u32) }
273            })
274    }
275}
276
277/// A trait that unifies all types that slice a portion of [`IdSet`].
278pub trait IdSetSlicer<I: Id> {
279    /// Converts the index type to inclusive start and exclusive end bounds.
280    ///
281    /// Returns `(start, end)` where:
282    /// - `start`: inclusive lower bound
283    /// - `end`: exclusive upper bound
284    fn to_range_bounds(self) -> (usize, usize);
285}
286
287// In the following implementations of `IdSetSlicer`, the `Id` values are upcast
288// from `u32` to `usize`. So adding one is guaranteed to *not* overflow.
289impl<I: Id> IdSetSlicer<I> for RangeTo<I> {
290    fn to_range_bounds(self) -> (usize, usize) {
291        (0, self.end.as_usize())
292    }
293}
294impl<I: Id> IdSetSlicer<I> for RangeFrom<I> {
295    fn to_range_bounds(self) -> (usize, usize) {
296        (self.start.as_usize(), I::cardinality() as usize)
297    }
298}
299impl<I: Id> IdSetSlicer<I> for Range<I> {
300    fn to_range_bounds(self) -> (usize, usize) {
301        (self.start.as_usize(), self.end.as_usize())
302    }
303}
304impl<I: Id> IdSetSlicer<I> for RangeFull {
305    fn to_range_bounds(self) -> (usize, usize) {
306        (0, I::cardinality() as usize)
307    }
308}
309impl<I: Id> IdSetSlicer<I> for RangeToInclusive<I> {
310    fn to_range_bounds(self) -> (usize, usize) {
311        (0, self.end.as_usize() + 1)
312    }
313}
314impl<I: Id> IdSetSlicer<I> for (Bound<I>, Bound<I>) {
315    fn to_range_bounds(self) -> (usize, usize) {
316        let (start_bound, end_bound) = self;
317        let start = match start_bound {
318            Bound::Included(id) => id.as_usize(),
319            Bound::Excluded(id) => id.as_usize() + 1,
320            Bound::Unbounded => 0,
321        };
322        let end = match end_bound {
323            Bound::Included(id) => id.as_usize() + 1,
324            Bound::Excluded(id) => id.as_usize(),
325            Bound::Unbounded => I::cardinality() as usize,
326        };
327        (start, end)
328    }
329}
330
331impl<I: Id> From<I> for IdSet<I> {
332    fn from(id: I) -> Self {
333        let mut set = Self::new_empty();
334        set.add(id);
335        set
336    }
337}
338
339impl<I: Id> Default for IdSet<I> {
340    fn default() -> Self {
341        Self::new_empty()
342    }
343}
344
345/// A set of IDs that may be accessed concurrently.
346///
347/// `AtomicIdSet` is backed by an array of `AtomicU64`.
348/// This allows individual ID bits to be added, removed, or tested atomically.
349/// But operations that span multiple `AtomicU64` values are non-atomic.
350#[derive(Debug)]
351pub struct AtomicIdSet<I> {
352    bits: SmallVec<[AtomicInnerPart; NR_PARTS_NO_ALLOC]>,
353    phantom: PhantomData<I>,
354}
355
356type AtomicInnerPart = AtomicU64;
357const_assert!(size_of::<AtomicInnerPart>() == size_of::<InnerPart>());
358
359impl<I: Id> AtomicIdSet<I> {
360    /// Creates a new `AtomicIdSet` from an `IdSet`.
361    pub fn new(value: IdSet<I>) -> Self {
362        let bits = value.bits.into_iter().map(AtomicU64::new).collect();
363        Self {
364            bits,
365            phantom: PhantomData,
366        }
367    }
368
369    /// Loads the value of the set with the given ordering.
370    ///
371    /// This operation is not atomic. When racing with a [`Self::store`]
372    /// operation, this load may return a set that contains a portion of the
373    /// new value and a portion of the old value. Load on each specific
374    /// word is atomic, and follows the specified ordering.
375    ///
376    /// Note that load with [`Ordering::Release`] is a valid operation, which
377    /// is different from the normal atomic operations. When coupled with
378    /// [`Ordering::Release`], it actually performs `fetch_or(0, Release)`.
379    pub fn load(&self, ordering: Ordering) -> IdSet<I> {
380        let bits = self
381            .bits
382            .iter()
383            .map(|part| match ordering {
384                Ordering::Release => part.fetch_or(0, ordering),
385                _ => part.load(ordering),
386            })
387            .collect();
388        IdSet {
389            bits,
390            phantom: PhantomData,
391        }
392    }
393
394    /// Stores a new value to the set with the given ordering.
395    ///
396    /// This operation is not atomic. When racing with a [`Self::load`]
397    /// operation, that load may return a set that contains a portion of the
398    /// new value and a portion of the old value. Load on each specific
399    /// word is atomic, and follows the specified ordering.
400    pub fn store(&self, value: &IdSet<I>, ordering: Ordering) {
401        for (part, new_part) in self.bits.iter().zip(value.bits.iter()) {
402            part.store(*new_part, ordering);
403        }
404    }
405
406    /// Atomically adds an ID with the given ordering.
407    pub fn add(&self, id: I, ordering: Ordering) {
408        let part_idx = part_idx(id);
409        let bit_idx = bit_idx(id);
410        if part_idx < self.bits.len() {
411            self.bits[part_idx].fetch_or(1 << bit_idx, ordering);
412        }
413    }
414
415    /// Atomically removes an ID with the given ordering.
416    pub fn remove(&self, id: I, ordering: Ordering) {
417        let part_idx = part_idx(id);
418        let bit_idx = bit_idx(id);
419        if part_idx < self.bits.len() {
420            self.bits[part_idx].fetch_and(!(1 << bit_idx), ordering);
421        }
422    }
423
424    /// Atomically checks if the set contains the specified ID.
425    pub fn contains(&self, id: I, ordering: Ordering) -> bool {
426        let part_idx = part_idx(id);
427        let bit_idx = bit_idx(id);
428        part_idx < self.bits.len() && (self.bits[part_idx].load(ordering) & (1 << bit_idx)) != 0
429    }
430}
431
432#[cfg(ktest)]
433mod id_set_tests {
434    use alloc::vec;
435
436    use super::*;
437    use crate::prelude::*;
438
439    /// A mock ID type for testing `IdSet`.
440    /// `C` is the cardinality of this ID type.
441    #[derive(Copy, Clone, Debug, Eq, PartialEq)]
442    struct MockId<const C: u32>(u32);
443
444    unsafe impl<const C: u32> Id for MockId<C> {
445        unsafe fn new_unchecked(raw_id: u32) -> Self {
446            MockId(raw_id)
447        }
448
449        fn cardinality() -> u32 {
450            C
451        }
452    }
453
454    impl<const C: u32> From<MockId<C>> for u32 {
455        fn from(id: MockId<C>) -> u32 {
456            id.0
457        }
458    }
459
460    // Test cases for `IdSet` with various cardinalities
461
462    #[ktest]
463    fn id_set_empty() {
464        type TestId = MockId<10>; // Cardinality 10
465        let set: IdSet<TestId> = IdSet::new_empty();
466        assert!(set.is_empty());
467        assert_eq!(set.count(), 0);
468        for i in 0..10 {
469            assert!(!set.contains(TestId::new(i)));
470        }
471    }
472
473    #[ktest]
474    fn id_set_full() {
475        type TestId = MockId<10>; // Cardinality 10
476        let set: IdSet<TestId> = IdSet::new_full();
477        assert!(!set.is_empty());
478        assert_eq!(set.count(), 10);
479        for i in 0..10 {
480            assert!(set.contains(TestId::new(i)));
481        }
482    }
483
484    #[ktest]
485    fn id_set_add_remove() {
486        type TestId = MockId<64>; // Cardinality 64 (one InnerPart)
487        let mut set: IdSet<TestId> = IdSet::new_empty();
488
489        assert!(set.is_empty());
490
491        set.add(TestId::new(0));
492        assert!(set.contains(TestId::new(0)));
493        assert_eq!(set.count(), 1);
494        assert!(!set.is_empty());
495
496        set.add(TestId::new(63));
497        assert!(set.contains(TestId::new(63)));
498        assert_eq!(set.count(), 2);
499
500        set.add(TestId::new(32));
501        assert!(set.contains(TestId::new(32)));
502        assert_eq!(set.count(), 3);
503
504        set.remove(TestId::new(0));
505        assert!(!set.contains(TestId::new(0)));
506        assert_eq!(set.count(), 2);
507
508        set.remove(TestId::new(63));
509        assert!(!set.contains(TestId::new(63)));
510        assert_eq!(set.count(), 1);
511
512        set.remove(TestId::new(32));
513        assert!(!set.contains(TestId::new(32)));
514        assert_eq!(set.count(), 0);
515        assert!(set.is_empty());
516
517        // Removing a non-existent ID should not panic or change the set
518        set.remove(TestId::new(1));
519        assert!(set.is_empty());
520    }
521
522    #[ktest]
523    fn id_set_add_remove_multi_part() {
524        type TestId = MockId<128>; // Cardinality 128 (two InnerParts)
525        let mut set: IdSet<TestId> = IdSet::new_empty();
526
527        set.add(TestId::new(0));
528        set.add(TestId::new(63)); // Last bit of first part
529        set.add(TestId::new(64)); // First bit of second part
530        set.add(TestId::new(127)); // Last bit of second part
531
532        assert_eq!(set.count(), 4);
533        assert!(set.contains(TestId::new(0)));
534        assert!(set.contains(TestId::new(63)));
535        assert!(set.contains(TestId::new(64)));
536        assert!(set.contains(TestId::new(127)));
537
538        set.remove(TestId::new(63));
539        assert!(!set.contains(TestId::new(63)));
540        assert_eq!(set.count(), 3);
541
542        set.remove(TestId::new(64));
543        assert!(!set.contains(TestId::new(64)));
544        assert_eq!(set.count(), 2);
545    }
546
547    #[ktest]
548    fn id_set_add_all_clear() {
549        type TestId = MockId<70>; // Cardinality 70 (spans two parts, second part not full)
550        let mut set: IdSet<TestId> = IdSet::new_empty();
551
552        set.add_all();
553        assert_eq!(set.count(), 70);
554        assert!(set.is_full());
555        for i in 0..70 {
556            assert!(set.contains(TestId::new(i)));
557        }
558
559        set.clear();
560        assert!(set.is_empty());
561        assert_eq!(set.count(), 0);
562        for i in 0..70 {
563            assert!(!set.contains(TestId::new(i)));
564        }
565    }
566
567    #[ktest]
568    fn id_set_iter() {
569        type TestId = MockId<5>; // Cardinality 5
570        let mut set: IdSet<TestId> = IdSet::new_empty();
571
572        set.add(TestId::new(2));
573        set.add(TestId::new(0));
574        set.add(TestId::new(4));
575
576        let collected_ids: Vec<TestId> = set.iter().collect();
577        assert_eq!(
578            collected_ids,
579            vec![TestId::new(0), TestId::new(2), TestId::new(4)]
580        );
581
582        set.clear();
583        let collected_ids: Vec<TestId> = set.iter().collect();
584        assert!(collected_ids.is_empty());
585    }
586
587    #[ktest]
588    fn id_set_iter_full() {
589        type TestId = MockId<3>; // Cardinality 3
590        let set: IdSet<TestId> = IdSet::new_full();
591        let collected_ids: Vec<TestId> = set.iter().collect();
592        assert_eq!(
593            collected_ids,
594            vec![TestId::new(0), TestId::new(1), TestId::new(2)]
595        );
596    }
597
598    #[ktest]
599    fn id_set_iter_multi_part() {
600        type TestId = MockId<100>; // Cardinality 100
601        let mut set: IdSet<TestId> = IdSet::new_empty();
602        set.add(TestId::new(1));
603        set.add(TestId::new(65));
604        set.add(TestId::new(99));
605        set.add(TestId::new(0));
606        set.add(TestId::new(63));
607
608        let collected_ids: Vec<TestId> = set.iter().collect();
609        assert_eq!(
610            collected_ids,
611            vec![
612                TestId::new(0),
613                TestId::new(1),
614                TestId::new(63),
615                TestId::new(65),
616                TestId::new(99)
617            ]
618        );
619    }
620
621    #[ktest]
622    fn id_set_from_id() {
623        type TestId = MockId<10>;
624        let id = TestId::new(5);
625        let set: IdSet<TestId> = id.into();
626        assert_eq!(set.count(), 1);
627        assert!(set.contains(id));
628        assert!(!set.contains(TestId::new(0)));
629    }
630
631    #[ktest]
632    fn id_set_cardinality_one() {
633        type TestId = MockId<1>; // Cardinality 1
634        let mut set: IdSet<TestId> = IdSet::new_empty();
635        assert!(set.is_empty());
636        assert_eq!(set.count(), 0);
637
638        set.add(TestId::new(0));
639        assert!(set.contains(TestId::new(0)));
640        assert_eq!(set.count(), 1);
641        assert!(set.is_full());
642
643        set.remove(TestId::new(0));
644        assert!(!set.contains(TestId::new(0)));
645        assert_eq!(set.count(), 0);
646        assert!(set.is_empty());
647
648        let full_set = IdSet::<TestId>::new_full();
649        assert!(full_set.contains(TestId::new(0)));
650        assert_eq!(full_set.count(), 1);
651    }
652
653    #[ktest]
654    fn id_set_exact_part_boundary() {
655        type TestId = MockId<64>; // Cardinality exactly one full part
656        let mut set: IdSet<TestId> = IdSet::new_empty();
657
658        set.add(TestId::new(0));
659        set.add(TestId::new(63));
660        assert_eq!(set.count(), 2);
661
662        let full_set = IdSet::<TestId>::new_full();
663        assert!(full_set.is_full());
664        assert_eq!(full_set.count(), 64);
665        for i in 0..64 {
666            assert!(full_set.contains(TestId::new(i)));
667        }
668    }
669
670    #[ktest]
671    fn id_set_just_over_part_boundary() {
672        type TestId = MockId<65>; // Cardinality just over one part
673        let mut set: IdSet<TestId> = IdSet::new_empty();
674
675        set.add(TestId::new(0));
676        set.add(TestId::new(63)); // End of first part
677        set.add(TestId::new(64)); // Start of second part
678        assert_eq!(set.count(), 3);
679
680        let full_set = IdSet::<TestId>::new_full();
681        assert!(full_set.is_full());
682        assert_eq!(full_set.count(), 65);
683        for i in 0..65 {
684            assert!(full_set.contains(TestId::new(i)));
685        }
686    }
687
688    #[ktest]
689    fn id_set_is_full_with_less_than_full_last_part() {
690        type TestId = MockId<70>; // Cardinality 70 (64 in first part, 6 in second)
691        let mut set: IdSet<TestId> = IdSet::new_full();
692
693        assert!(set.is_full());
694        assert_eq!(set.count(), 70);
695
696        // Remove one ID from the last part
697        set.remove(TestId::new(69));
698        assert!(!set.is_full());
699        assert_eq!(set.count(), 69);
700
701        // Add it back
702        set.add(TestId::new(69));
703        assert!(set.is_full());
704        assert_eq!(set.count(), 70);
705    }
706
707    #[ktest]
708    fn id_set_default() {
709        type TestId = MockId<10>;
710        let set: IdSet<TestId> = Default::default();
711        assert!(set.is_empty());
712        assert_eq!(set.count(), 0);
713    }
714
715    #[ktest]
716    fn iter_in_range() {
717        type TestId = MockId<7>;
718        let mut set: IdSet<TestId> = IdSet::new_empty();
719        set.add(TestId::new(0));
720        set.add(TestId::new(1));
721        set.add(TestId::new(2));
722        set.add(TestId::new(5));
723        set.add(TestId::new(6));
724
725        let collected_ids: Vec<TestId> = set.iter_in(TestId::new(1)..TestId::new(5)).collect();
726        assert_eq!(collected_ids, vec![TestId::new(1), TestId::new(2)],);
727    }
728
729    #[ktest]
730    fn iter_in_range_to() {
731        type TestId = MockId<7>;
732        let mut set: IdSet<TestId> = IdSet::new_empty();
733        set.add(TestId::new(0));
734        set.add(TestId::new(1));
735        set.add(TestId::new(2));
736        set.add(TestId::new(5));
737        set.add(TestId::new(6));
738
739        let collected_ids: Vec<TestId> = set.iter_in(..TestId::new(5)).collect();
740        assert_eq!(
741            collected_ids,
742            vec![TestId::new(0), TestId::new(1), TestId::new(2)],
743        );
744    }
745
746    #[ktest]
747    fn iter_in_range_to_inclusive() {
748        type TestId = MockId<7>;
749        let mut set: IdSet<TestId> = IdSet::new_empty();
750        set.add(TestId::new(0));
751        set.add(TestId::new(1));
752        set.add(TestId::new(2));
753        set.add(TestId::new(5));
754        set.add(TestId::new(6));
755
756        let collected_ids: Vec<TestId> = set.iter_in(..=TestId::new(5)).collect();
757        assert_eq!(
758            collected_ids,
759            vec![
760                TestId::new(0),
761                TestId::new(1),
762                TestId::new(2),
763                TestId::new(5)
764            ],
765        );
766    }
767
768    #[ktest]
769    fn iter_in_range_from() {
770        type TestId = MockId<7>;
771        let mut set: IdSet<TestId> = IdSet::new_empty();
772        set.add(TestId::new(0));
773        set.add(TestId::new(1));
774        set.add(TestId::new(2));
775        set.add(TestId::new(5));
776        set.add(TestId::new(6));
777
778        let collected_ids: Vec<TestId> = set.iter_in(TestId::new(2)..).collect();
779        assert_eq!(
780            collected_ids,
781            vec![TestId::new(2), TestId::new(5), TestId::new(6)],
782        );
783    }
784
785    #[ktest]
786    fn iter_in_range_full() {
787        type TestId = MockId<7>;
788        let mut set: IdSet<TestId> = IdSet::new_empty();
789        set.add(TestId::new(0));
790        set.add(TestId::new(1));
791        set.add(TestId::new(2));
792        set.add(TestId::new(5));
793        set.add(TestId::new(6));
794
795        let collected_ids: Vec<TestId> = set.iter_in(..).collect();
796        assert_eq!(
797            collected_ids,
798            vec![
799                TestId::new(0),
800                TestId::new(1),
801                TestId::new(2),
802                TestId::new(5),
803                TestId::new(6)
804            ],
805        );
806    }
807
808    #[ktest]
809    fn iter_in_bound_tuple_inclusive_exclusive() {
810        type TestId = MockId<7>;
811        let mut set: IdSet<TestId> = IdSet::new_empty();
812        set.add(TestId::new(0));
813        set.add(TestId::new(1));
814        set.add(TestId::new(2));
815        set.add(TestId::new(5));
816        set.add(TestId::new(6));
817
818        let collected_ids: Vec<TestId> = set
819            .iter_in((
820                Bound::Included(TestId::new(1)),
821                Bound::Excluded(TestId::new(5)),
822            ))
823            .collect();
824        assert_eq!(collected_ids, vec![TestId::new(1), TestId::new(2)],);
825    }
826
827    #[ktest]
828    fn iter_in_bound_tuple_exclusive_inclusive() {
829        type TestId = MockId<7>;
830        let mut set: IdSet<TestId> = IdSet::new_empty();
831        set.add(TestId::new(0));
832        set.add(TestId::new(1));
833        set.add(TestId::new(2));
834        set.add(TestId::new(5));
835        set.add(TestId::new(6));
836
837        let collected_ids: Vec<TestId> = set
838            .iter_in((
839                Bound::Excluded(TestId::new(1)),
840                Bound::Included(TestId::new(5)),
841            ))
842            .collect();
843        assert_eq!(collected_ids, vec![TestId::new(2), TestId::new(5)],);
844    }
845
846    #[ktest]
847    fn iter_in_unbounded_bounds() {
848        type TestId = MockId<7>;
849        let mut set: IdSet<TestId> = IdSet::new_empty();
850        set.add(TestId::new(0));
851        set.add(TestId::new(1));
852        set.add(TestId::new(2));
853        set.add(TestId::new(5));
854        set.add(TestId::new(6));
855
856        let collected_ids: Vec<TestId> = set
857            .iter_in((Bound::Unbounded::<TestId>, Bound::Unbounded::<TestId>))
858            .collect();
859        assert_eq!(
860            collected_ids,
861            vec![
862                TestId::new(0),
863                TestId::new(1),
864                TestId::new(2),
865                TestId::new(5),
866                TestId::new(6)
867            ],
868        );
869    }
870
871    #[ktest]
872    fn iter_in_half_unbounded() {
873        type TestId = MockId<7>;
874        let mut set: IdSet<TestId> = IdSet::new_empty();
875        set.add(TestId::new(0));
876        set.add(TestId::new(1));
877        set.add(TestId::new(2));
878        set.add(TestId::new(5));
879        set.add(TestId::new(6));
880
881        let collected_ids: Vec<TestId> = set
882            .iter_in((Bound::Included(TestId::new(2)), Bound::Unbounded::<TestId>))
883            .collect();
884        assert_eq!(
885            collected_ids,
886            vec![TestId::new(2), TestId::new(5), TestId::new(6)],
887        );
888
889        let collected_ids: Vec<TestId> = set
890            .iter_in((Bound::Unbounded::<TestId>, Bound::Included(TestId::new(2))))
891            .collect();
892        assert_eq!(
893            collected_ids,
894            vec![TestId::new(0), TestId::new(1), TestId::new(2)],
895        );
896    }
897
898    #[ktest]
899    fn iter_in_range_starts_after_last() {
900        type TestId = MockId<7>;
901        let mut set: IdSet<TestId> = IdSet::new_empty();
902        set.add(TestId::new(0));
903        set.add(TestId::new(1));
904        set.add(TestId::new(2));
905
906        let collected_ids: Vec<TestId> = set.iter_in(TestId::new(3)..).collect();
907        assert_eq!(collected_ids, vec![],);
908    }
909
910    #[ktest]
911    fn iter_in_range_ends_after_last() {
912        type TestId = MockId<7>;
913        let mut set: IdSet<TestId> = IdSet::new_empty();
914        set.add(TestId::new(0));
915        set.add(TestId::new(1));
916        set.add(TestId::new(2));
917
918        let collected_ids: Vec<TestId> = set.iter_in(..TestId::new(3)).collect();
919        assert_eq!(
920            collected_ids,
921            vec![TestId::new(0), TestId::new(1), TestId::new(2)],
922        );
923    }
924
925    #[ktest]
926    fn iter_in_range_next_part() {
927        type TestId = MockId<{ InnerPart::BITS }>;
928        let last_id = TestId::new(InnerPart::BITS - 1);
929
930        let mut set: IdSet<TestId> = IdSet::new_empty();
931        set.add(last_id);
932
933        let collected_ids: Vec<TestId> = set
934            .iter_in((Bound::Excluded(last_id), Bound::Included(last_id)))
935            .collect();
936        assert_eq!(collected_ids, vec![],);
937    }
938}
ostd/util/id_set.rs

ostd/util/
id_set.rs
