usize_set/
vec.rs

1//! Index set backed by a [`Vec`].
2
3use alloc::vec::Vec;
4#[cfg(feature = "serialize-borsh")]
5use alloc::{format, string::ToString};
6#[cfg(feature = "serialize-borsh")]
7use borsh::{BorshDeserialize, BorshSchema, BorshSerialize};
8#[cfg(feature = "serialize-serde")]
9use serde::{Deserialize, Serialize};
10
11use super::calculate_map_and_set_indices;
12use super::macros::*;
13use super::storage;
14use super::IndexSet;
15
16#[cfg(feature = "serialize-serde")]
17mod serde_deserialize {
18    use alloc::vec::Vec;
19
20    use serde::{Deserialize, Deserializer};
21
22    /// Deserialize a [`VecIndexSet`] from serde data.
23    pub fn from<'de, D, S>(deserializer: D) -> Result<Vec<(usize, S)>, D::Error>
24    where
25        D: Deserializer<'de>,
26        S: Deserialize<'de>,
27    {
28        let bit_sets: Vec<(usize, S)> = Deserialize::deserialize(deserializer)?;
29        for window in bit_sets.windows(2) {
30            let &[(a, _), (b, _)] = window else {
31                unreachable!()
32            };
33            if a > b {
34                return Err(serde::de::Error::custom(
35                    "VecIndexSet should have been sorted",
36                ));
37            }
38        }
39        Ok(bit_sets)
40    }
41}
42
43#[cfg(feature = "serialize-borsh")]
44mod borsh_deserialize {
45    use super::*;
46
47    /// Deserialize a [`VecIndexSet`] from borsh data.
48    pub fn from<R, S>(reader: &mut R) -> Result<Vec<(usize, S)>, borsh::io::Error>
49    where
50        R: borsh::io::Read,
51        S: borsh::de::BorshDeserialize,
52    {
53        let bit_sets: Vec<(usize, S)> = borsh::BorshDeserialize::deserialize_reader(reader)?;
54        for window in bit_sets.windows(2) {
55            let &[(a, _), (b, _)] = window else {
56                unreachable!()
57            };
58            if a > b {
59                return Err(borsh::io::Error::new(
60                    borsh::io::ErrorKind::Other,
61                    "VecIndexSet should have been sorted",
62                ));
63            }
64        }
65        Ok(bit_sets)
66    }
67}
68
69/// Index set backed by a [`Vec`].
70#[derive(Default, Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
71#[cfg_attr(
72    feature = "serialize-borsh",
73    derive(BorshSerialize, BorshDeserialize, BorshSchema)
74)]
75#[cfg_attr(feature = "serialize-serde", derive(Serialize, Deserialize))]
76#[repr(transparent)]
77pub struct VecIndexSet<S = u64> {
78    /// Pairs of indices to bit vectors, containing the actual boolean
79    /// values to be asserted.
80    ///
81    /// If the bit `B` is set, at the bit vector with index `S`, then
82    /// the index `S::WIDTH * S + B` is in the set.
83    #[cfg_attr(
84        feature = "serialize-borsh",
85        borsh(deserialize_with = "borsh_deserialize::from")
86    )]
87    #[cfg_attr(
88        feature = "serialize-serde",
89        serde(deserialize_with = "serde_deserialize::from")
90    )]
91    #[cfg_attr(
92        feature = "serialize-serde",
93        serde(bound(deserialize = "S: Deserialize<'de>"))
94    )]
95    bit_sets: Vec<(usize, S)>,
96}
97
98impl<S> VecIndexSet<S> {
99    /// Create a new [`VecIndexSet`].
100    pub const fn new() -> Self {
101        Self {
102            bit_sets: Vec::new(),
103        }
104    }
105
106    /// Create a new [`VecIndexSet`] with the given capacity.
107    #[inline]
108    pub fn with_capacity(capacity: usize) -> Self {
109        Self {
110            bit_sets: Vec::with_capacity(capacity),
111        }
112    }
113}
114
115impl<S: storage::Storage> VecIndexSet<S> {
116    /// Lookup the bit set at `map_index`, or initialize it
117    /// with zero, if it doesn't exist.
118    #[inline]
119    fn lookup_or_zero(&mut self, map_index: usize) -> &mut S {
120        let pair_index = self.lookup_or_initialize_pair(map_index);
121        let (_, set) = &mut self.bit_sets[pair_index];
122        set
123    }
124
125    /// Lookup the vec index of the bit set at `map_index`, or initialize it
126    /// with zero, if it doesn't exist, returning the initialized index.
127    #[inline]
128    fn lookup_or_initialize_pair(&mut self, map_index: usize) -> usize {
129        self.lookup_pair(map_index)
130            .unwrap_or_else(|insert_at_index| {
131                self.bit_sets.insert(insert_at_index, (map_index, S::ZERO));
132                insert_at_index
133            })
134    }
135
136    /// Lookup the vec index of the bit set at `map_index`.
137    #[inline]
138    fn lookup_pair(&self, map_index: usize) -> Result<usize, usize> {
139        self.bit_sets.binary_search_by_key(&map_index, |&(i, _)| i)
140    }
141}
142
143impl<S: storage::Storage> IndexSet for VecIndexSet<S> {
144    #[inline]
145    fn len(&self) -> usize {
146        self.bit_sets
147            .iter()
148            .map(|(_, set)| set.num_of_high_bits())
149            .sum::<usize>()
150    }
151
152    #[inline]
153    fn is_empty(&self) -> bool {
154        self.bit_sets.is_empty()
155    }
156
157    fn insert(&mut self, index: usize) {
158        let (map_index, bit_set_index) = calculate_map_and_set_indices::<S>(index);
159        let set = self.lookup_or_zero(map_index);
160        *set |= S::from_usize(1 << bit_set_index);
161    }
162
163    fn remove(&mut self, index: usize) {
164        let (map_index, bit_set_index) = calculate_map_and_set_indices::<S>(index);
165        let maybe_remove_index = self.lookup_pair(map_index).ok().and_then(|pair_index| {
166            let (_, set) = &mut self.bit_sets[pair_index];
167            *set &= !S::from_usize(1 << bit_set_index);
168            if *set == S::ZERO {
169                Some(pair_index)
170            } else {
171                None
172            }
173        });
174        if let Some(pair_index) = maybe_remove_index {
175            self.bit_sets.remove(pair_index);
176        }
177    }
178
179    fn contains(&self, index: usize) -> bool {
180        let (map_index, bit_set_index) = calculate_map_and_set_indices::<S>(index);
181        self.lookup_pair(map_index)
182            .map(|pair_index| {
183                let &(_, set) = &self.bit_sets[pair_index];
184                set & S::from_usize(1 << bit_set_index) != S::ZERO
185            })
186            .unwrap_or(false)
187    }
188
189    #[inline]
190    fn iter(&self) -> impl Iterator<Item = usize> + '_ {
191        self.bit_sets.iter().flat_map(|&(map_index, set)| {
192            (0..S::WIDTH).filter_map(move |bit_set_index| {
193                let is_bit_set = (set & S::from_usize(1 << bit_set_index)) != S::ZERO;
194                if is_bit_set {
195                    Some(map_index * S::WIDTH + bit_set_index)
196                } else {
197                    None
198                }
199            })
200        })
201    }
202
203    #[inline]
204    fn union(&mut self, other: &VecIndexSet<S>) {
205        // naive implementation
206        for &(map_index, other_set) in other.bit_sets.iter() {
207            let set = self.lookup_or_zero(map_index);
208            *set |= other_set;
209        }
210    }
211
212    #[inline]
213    fn reserve(&mut self, size: usize) {
214        self.bit_sets.reserve(size);
215    }
216}
217
218index_set_impl_from!(crate::vec::VecIndexSet);
219index_set_impl_from_iterator!(crate::vec::VecIndexSet);
220index_set_impl_extend!(crate::vec::VecIndexSet);
221index_set_tests!(crate::vec::VecIndexSet);
usize_set/vec.rs

usize_set/
vec.rs
