boomerang_tinymap/
key_set.rs

1//! A set of keys
2//!
3//! [`KeySet`] is more efficient than a [`super::TinySecondaryMap<K, ()>`] in memory and compute since it uses a
4//! [`FixedBitSet`] to store the keys under the hood.
5use std::{marker::PhantomData, ops::Index};
6
7use fixedbitset::{FixedBitSet, Ones};
8
9use super::Key;
10
11/// A unique set of keys.
12#[derive(Clone, Debug)]
13pub struct KeySet<K: Key> {
14    data: FixedBitSet,
15    _k: PhantomData<K>,
16}
17
18impl<K: Key> KeySet<K> {
19    /// Construct a new, empty [`KeySet`].
20    pub fn with_capacity(capacity: usize) -> Self {
21        Self {
22            data: FixedBitSet::with_capacity(capacity),
23            _k: PhantomData,
24        }
25    }
26
27    /// Returns `true` if the set is empty.
28    pub fn is_empty(&self) -> bool {
29        self.data.is_clear()
30    }
31
32    /// Insert a key into the set.
33    #[inline]
34    pub fn insert(&mut self, key: K) {
35        self.data.set(key.index(), true);
36    }
37
38    /// Extend the set from an iterable.
39    #[inline]
40    pub fn extend(&mut self, keys: impl IntoIterator<Item = K>) {
41        self.data.extend(keys.into_iter().map(|key| key.index()));
42    }
43
44    /// Clear the set.
45    #[inline]
46    pub fn clear(&mut self) {
47        self.data.clear();
48    }
49
50    /// Returns an iterator over the `K` entries in the map.
51    pub fn iter(&self) -> Iter<'_, K> {
52        Iter {
53            inner: self.data.ones(),
54            count: self.data.count_ones(..),
55            _k: PhantomData,
56        }
57    }
58}
59
60impl<K: Key + std::fmt::Display> std::fmt::Display for KeySet<K> {
61    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
62        f.debug_list()
63            .entries(self.iter().map(|k| k.to_string()))
64            .finish()
65    }
66}
67
68pub struct Iter<'a, K: Key> {
69    inner: Ones<'a>,
70    count: usize,
71    _k: PhantomData<K>,
72}
73
74impl<'a, K: Key> Iterator for Iter<'a, K> {
75    type Item = K;
76
77    #[inline]
78    fn next(&mut self) -> Option<Self::Item> {
79        self.inner.next().map(|idx| K::from(idx))
80    }
81}
82
83impl<'a, K: Key> ExactSizeIterator for Iter<'a, K> {
84    fn len(&self) -> usize {
85        self.count
86    }
87}
88
89impl<K: Key> FromIterator<K> for KeySet<K> {
90    fn from_iter<T: IntoIterator<Item = K>>(iter: T) -> Self {
91        Self {
92            data: iter.into_iter().map(|k| k.index()).collect(),
93            _k: PhantomData,
94        }
95    }
96}
97
98impl<K: Key> Index<K> for KeySet<K> {
99    type Output = bool;
100
101    fn index(&self, key: K) -> &Self::Output {
102        // Note: bits outside the capcity are always disabled, thus this will never panic
103        &self.data[key.index()]
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110    use crate::DefaultKey;
111
112    #[test]
113    fn test_new() {
114        let set: KeySet<DefaultKey> = KeySet::with_capacity(10);
115        assert!(set.is_empty());
116    }
117
118    #[test]
119    fn test_insert_and_index() {
120        let mut set = KeySet::with_capacity(10);
121        let key1 = DefaultKey::from(0);
122        let key2 = DefaultKey::from(1);
123
124        set.insert(key1);
125        assert!(set[key1]);
126        assert!(!set[key2]);
127
128        set.insert(key2);
129        assert!(set[key1]);
130        assert!(set[key2]);
131    }
132
133    #[test]
134    fn test_iter() {
135        let mut set = KeySet::with_capacity(10);
136        set.insert(DefaultKey::from(0));
137        set.insert(DefaultKey::from(2));
138        set.insert(DefaultKey::from(4));
139
140        let keys: Vec<DefaultKey> = set.iter().collect();
141        assert_eq!(
142            keys,
143            vec![
144                DefaultKey::from(0),
145                DefaultKey::from(2),
146                DefaultKey::from(4)
147            ]
148        );
149    }
150
151    #[test]
152    fn test_from_iter() {
153        let set: KeySet<DefaultKey> = vec![DefaultKey::from(0), DefaultKey::from(1)]
154            .into_iter()
155            .collect();
156        assert_eq!(
157            set.iter().collect::<Vec<DefaultKey>>(),
158            vec![DefaultKey::from(0), DefaultKey::from(1)]
159        );
160    }
161
162    #[test]
163    fn test_empty_iter() {
164        let set: KeySet<DefaultKey> = KeySet::with_capacity(10);
165        assert!(set.iter().next().is_none());
166    }
167
168    #[test]
169    fn test_extend() {
170        let mut set = KeySet::with_capacity(10);
171        set.extend(vec![DefaultKey::from(0), DefaultKey::from(1)]);
172        assert_eq!(
173            set.iter().collect::<Vec<DefaultKey>>(),
174            vec![DefaultKey::from(0), DefaultKey::from(1)]
175        );
176    }
177
178    #[test]
179    fn test_exact_size_iter() {
180        let mut set = KeySet::with_capacity(10);
181        set.extend(vec![DefaultKey::from(0), DefaultKey::from(1)]);
182        assert_eq!(set.iter().len(), 2);
183    }
184}
boomerang_tinymap/key_set.rs

boomerang_tinymap/
key_set.rs
