commonware_utils/
set.rs

1//! Ordered collections that guarantee sorted, deduplicated keys.
2
3#[cfg(not(feature = "std"))]
4use alloc::vec::Vec;
5use bytes::{Buf, BufMut};
6use commonware_codec::{EncodeSize, RangeCfg, Read, Write};
7use core::{
8    fmt,
9    ops::{Deref, Index, Range},
10};
11
12#[cfg(not(feature = "std"))]
13type VecIntoIter<T> = alloc::vec::IntoIter<T>;
14#[cfg(feature = "std")]
15type VecIntoIter<T> = std::vec::IntoIter<T>;
16
17/// An ordered, deduplicated slice of items.
18///
19/// After construction, the contained [`Vec<T>`] is sealed and cannot be modified. To unseal the
20/// inner [`Vec<T>`], use the [`Into<Vec<T>>`] impl.
21#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
22pub struct Ordered<T>(Vec<T>);
23
24impl<T> Ordered<T> {
25    /// Returns the size of the ordered collection.
26    pub fn len(&self) -> usize {
27        self.0.len()
28    }
29
30    /// Returns `true` if the collection is empty.
31    pub fn is_empty(&self) -> bool {
32        self.0.is_empty()
33    }
34
35    /// Returns an item by index, if it exists.
36    pub fn get(&self, index: usize) -> Option<&T> {
37        self.0.get(index)
38    }
39
40    /// Returns the position of a given item in the collection, if it exists.
41    pub fn position(&self, item: &T) -> Option<usize>
42    where
43        T: Ord,
44    {
45        self.0.binary_search(item).ok()
46    }
47
48    /// Returns an iterator over the items in the collection.
49    pub fn iter(&self) -> core::slice::Iter<'_, T> {
50        self.into_iter()
51    }
52}
53
54impl<T: Write> Write for Ordered<T> {
55    fn write(&self, buf: &mut impl BufMut) {
56        self.0.write(buf);
57    }
58}
59
60impl<T: EncodeSize> EncodeSize for Ordered<T> {
61    fn encode_size(&self) -> usize {
62        self.0.encode_size()
63    }
64}
65
66impl<T: Read> Read for Ordered<T> {
67    type Cfg = (RangeCfg<usize>, T::Cfg);
68
69    fn read_cfg(buf: &mut impl Buf, cfg: &Self::Cfg) -> Result<Self, commonware_codec::Error> {
70        Ok(Self(Vec::<T>::read_cfg(buf, cfg)?))
71    }
72}
73
74impl<T: Ord> FromIterator<T> for Ordered<T> {
75    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
76        let items: Vec<_> = iter.into_iter().collect();
77        items.into()
78    }
79}
80
81impl<T: Ord> From<Vec<T>> for Ordered<T> {
82    fn from(mut items: Vec<T>) -> Self {
83        items.sort();
84        items.dedup();
85        Self(items)
86    }
87}
88
89impl<T: Ord + Clone> From<&[T]> for Ordered<T> {
90    fn from(items: &[T]) -> Self {
91        items.iter().cloned().collect()
92    }
93}
94
95impl<T: Ord, const N: usize> From<[T; N]> for Ordered<T> {
96    fn from(items: [T; N]) -> Self {
97        items.into_iter().collect()
98    }
99}
100
101impl<T: Ord + Clone, const N: usize> From<&[T; N]> for Ordered<T> {
102    fn from(items: &[T; N]) -> Self {
103        items.as_slice().into()
104    }
105}
106
107impl<T> IntoIterator for Ordered<T> {
108    type Item = T;
109    type IntoIter = VecIntoIter<T>;
110
111    fn into_iter(self) -> Self::IntoIter {
112        self.0.into_iter()
113    }
114}
115
116impl<'a, T> IntoIterator for &'a Ordered<T> {
117    type Item = &'a T;
118    type IntoIter = core::slice::Iter<'a, T>;
119
120    fn into_iter(self) -> Self::IntoIter {
121        self.0.iter()
122    }
123}
124
125impl<T> Index<usize> for Ordered<T> {
126    type Output = T;
127
128    fn index(&self, index: usize) -> &Self::Output {
129        &self.0[index]
130    }
131}
132
133impl<T> Index<Range<usize>> for Ordered<T> {
134    type Output = [T];
135
136    fn index(&self, index: Range<usize>) -> &Self::Output {
137        &self.0[index]
138    }
139}
140
141impl<T> AsRef<[T]> for Ordered<T> {
142    fn as_ref(&self) -> &[T] {
143        &self.0
144    }
145}
146
147impl<T: fmt::Display> fmt::Display for Ordered<T> {
148    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
149        write!(f, "[")?;
150        for (i, item) in self.0.iter().enumerate() {
151            if i > 0 {
152                write!(f, ", ")?;
153            }
154            write!(f, "{item}")?;
155        }
156        write!(f, "]")
157    }
158}
159
160impl<T: Ord> From<Ordered<T>> for Vec<T> {
161    fn from(set: Ordered<T>) -> Self {
162        set.0
163    }
164}
165
166/// An ordered, deduplicated slice of items each paired with some associated value.
167///
168/// Like [`Ordered`], the contained [`Vec<(K, V)>`] is sealed after construction and cannot be modified. To unseal the
169/// inner [`Vec<(K, V)>`], use the [`Into<Vec<(K, V)>>`] impl.
170///
171/// Consumers that only need the ordered keys can treat an [`OrderedAssociated`] as an
172/// [`Ordered`] through deref coercions.
173#[derive(Clone, PartialEq, Eq, Hash)]
174pub struct OrderedAssociated<K, V> {
175    keys: Ordered<K>,
176    values: Vec<V>,
177}
178
179impl<K, V> OrderedAssociated<K, V> {
180    /// Returns the number of entries in the map.
181    pub fn len(&self) -> usize {
182        self.keys.len()
183    }
184
185    /// Returns `true` if the map is empty.
186    pub fn is_empty(&self) -> bool {
187        self.keys.is_empty()
188    }
189
190    /// Returns a key by index, if it exists.
191    pub fn get(&self, index: usize) -> Option<&K> {
192        self.keys.get(index)
193    }
194
195    /// Returns the position of the provided key, if it exists.
196    pub fn position(&self, key: &K) -> Option<usize>
197    where
198        K: Ord,
199    {
200        self.keys.position(key)
201    }
202
203    /// Returns the ordered keys as an [`Ordered`] reference.
204    pub fn keys(&self) -> &Ordered<K> {
205        &self.keys
206    }
207
208    /// Consumes the map and returns the ordered keys.
209    pub fn into_keys(self) -> Ordered<K> {
210        self.keys
211    }
212
213    /// Returns the associated value at `index`, if it exists.
214    pub fn value(&self, index: usize) -> Option<&V> {
215        self.values.get(index)
216    }
217
218    /// Returns the associated value for `key`, if it exists.
219    pub fn get_value(&self, key: &K) -> Option<&V>
220    where
221        K: Ord,
222    {
223        self.position(key).and_then(|index| self.values.get(index))
224    }
225
226    /// Returns the associated values.
227    pub fn values(&self) -> &[V] {
228        &self.values
229    }
230
231    /// Returns a zipped iterator over keys and values.
232    pub fn iter_pairs(&self) -> impl Iterator<Item = (&K, &V)> {
233        self.keys.iter().zip(self.values.iter())
234    }
235
236    /// Returns an iterator over the ordered keys.
237    pub fn iter(&self) -> core::slice::Iter<'_, K> {
238        self.keys.iter()
239    }
240}
241
242impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for OrderedAssociated<K, V> {
243    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
244        f.debug_tuple("OrderedAssociated")
245            .field(&self.iter_pairs().collect::<Vec<_>>())
246            .finish()
247    }
248}
249
250impl<K, V> AsRef<[K]> for OrderedAssociated<K, V> {
251    fn as_ref(&self) -> &[K] {
252        self.keys.as_ref()
253    }
254}
255
256impl<K, V> AsRef<Ordered<K>> for OrderedAssociated<K, V> {
257    fn as_ref(&self) -> &Ordered<K> {
258        &self.keys
259    }
260}
261
262impl<K, V> Deref for OrderedAssociated<K, V> {
263    type Target = Ordered<K>;
264
265    fn deref(&self) -> &Self::Target {
266        &self.keys
267    }
268}
269
270impl<K: Ord, V> FromIterator<(K, V)> for OrderedAssociated<K, V> {
271    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
272        let mut items: Vec<(K, V)> = iter.into_iter().collect();
273        items.sort_by(|(lk, _), (rk, _)| lk.cmp(rk));
274        items.dedup_by(|l, r| l.0 == r.0);
275
276        let mut keys = Vec::with_capacity(items.len());
277        let mut values = Vec::with_capacity(items.len());
278        for (key, value) in items {
279            keys.push(key);
280            values.push(value);
281        }
282
283        Self {
284            keys: Ordered(keys),
285            values,
286        }
287    }
288}
289
290impl<K: Ord + Clone, V: Clone> From<&[(K, V)]> for OrderedAssociated<K, V> {
291    fn from(items: &[(K, V)]) -> Self {
292        items.iter().cloned().collect()
293    }
294}
295
296impl<K: Ord, V> From<Vec<(K, V)>> for OrderedAssociated<K, V> {
297    fn from(items: Vec<(K, V)>) -> Self {
298        items.into_iter().collect()
299    }
300}
301
302impl<K: Ord, V, const N: usize> From<[(K, V); N]> for OrderedAssociated<K, V> {
303    fn from(items: [(K, V); N]) -> Self {
304        items.into_iter().collect()
305    }
306}
307
308impl<K: Ord + Clone, V: Clone, const N: usize> From<&[(K, V); N]> for OrderedAssociated<K, V> {
309    fn from(items: &[(K, V); N]) -> Self {
310        items.as_slice().into()
311    }
312}
313
314impl<K: Ord, V> From<OrderedAssociated<K, V>> for Vec<(K, V)> {
315    fn from(wrapped: OrderedAssociated<K, V>) -> Self {
316        wrapped.into_iter().collect()
317    }
318}
319
320impl<K: Write, V: Write> Write for OrderedAssociated<K, V> {
321    fn write(&self, buf: &mut impl BufMut) {
322        self.keys.write(buf);
323        self.values.write(buf);
324    }
325}
326
327impl<K: EncodeSize, V: EncodeSize> EncodeSize for OrderedAssociated<K, V> {
328    fn encode_size(&self) -> usize {
329        self.keys.encode_size() + self.values.encode_size()
330    }
331}
332
333impl<K: Read, V: Read> Read for OrderedAssociated<K, V> {
334    type Cfg = (RangeCfg<usize>, K::Cfg, V::Cfg);
335
336    fn read_cfg(buf: &mut impl Buf, cfg: &Self::Cfg) -> Result<Self, commonware_codec::Error> {
337        let (range_cfg, key_cfg, value_cfg) = cfg;
338        let keys = Ordered::<K>::read_cfg(buf, &(*range_cfg, key_cfg.clone()))?;
339        let values = Vec::<V>::read_cfg(buf, &(RangeCfg::exact(keys.len()), value_cfg.clone()))?;
340        Ok(Self { keys, values })
341    }
342}
343
344impl<K, V> IntoIterator for OrderedAssociated<K, V> {
345    type Item = (K, V);
346    type IntoIter = OrderedAssociatedIntoIter<K, V>;
347
348    fn into_iter(self) -> Self::IntoIter {
349        OrderedAssociatedIntoIter {
350            keys: self.keys.into_iter(),
351            values: self.values.into_iter(),
352        }
353    }
354}
355
356impl<'a, K, V> IntoIterator for &'a OrderedAssociated<K, V> {
357    type Item = (&'a K, &'a V);
358    type IntoIter = core::iter::Zip<core::slice::Iter<'a, K>, core::slice::Iter<'a, V>>;
359
360    fn into_iter(self) -> Self::IntoIter {
361        self.keys.iter().zip(self.values.iter())
362    }
363}
364
365/// Owned iterator over [`OrderedAssociated`].
366pub struct OrderedAssociatedIntoIter<K, V> {
367    keys: VecIntoIter<K>,
368    values: VecIntoIter<V>,
369}
370
371impl<K, V> Iterator for OrderedAssociatedIntoIter<K, V> {
372    type Item = (K, V);
373
374    fn next(&mut self) -> Option<Self::Item> {
375        let key = self.keys.next()?;
376        let value = self.values.next()?;
377        Some((key, value))
378    }
379
380    fn size_hint(&self) -> (usize, Option<usize>) {
381        self.keys.size_hint()
382    }
383}
384
385impl<K, V> ExactSizeIterator for OrderedAssociatedIntoIter<K, V> {}
386
387impl<K, V> DoubleEndedIterator for OrderedAssociatedIntoIter<K, V> {
388    fn next_back(&mut self) -> Option<Self::Item> {
389        let key = self.keys.next_back()?;
390        let value = self.values.next_back()?;
391        Some((key, value))
392    }
393}
394
395#[cfg(test)]
396mod test {
397    use super::*;
398
399    #[test]
400    fn test_sorted_unique_construct_unseal() {
401        const CASE: [u8; 12] = [1, 3, 2, 5, 4, 3, 1, 7, 9, 6, 8, 4];
402        const EXPECTED: [u8; 9] = [1, 2, 3, 4, 5, 6, 7, 8, 9];
403
404        let sorted: Ordered<_> = CASE.into_iter().collect();
405        assert_eq!(sorted.iter().copied().collect::<Vec<_>>(), EXPECTED);
406
407        let unsealed: Vec<_> = sorted.into();
408        assert_eq!(unsealed, EXPECTED);
409    }
410
411    #[test]
412    fn test_sorted_unique_codec_roundtrip() {
413        const CASE: [u8; 9] = [1, 2, 3, 4, 5, 6, 7, 8, 9];
414        let sorted: Ordered<_> = CASE.into_iter().collect();
415
416        let mut buf = Vec::with_capacity(sorted.encode_size());
417        sorted.write(&mut buf);
418        let decoded =
419            Ordered::<u8>::read_cfg(&mut buf.as_slice(), &(RangeCfg::from(0..=9), ())).unwrap();
420
421        assert_eq!(sorted, decoded);
422    }
423
424    #[test]
425    fn test_sorted_unique_display() {
426        const CASE: [u8; 9] = [1, 2, 3, 4, 5, 6, 7, 8, 9];
427
428        #[derive(Ord, PartialOrd, Eq, PartialEq)]
429        struct Example(u8);
430        impl fmt::Display for Example {
431            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
432                write!(f, "ex({})", self.0)
433            }
434        }
435        let examples = CASE.into_iter().map(Example).collect::<Vec<_>>();
436
437        let sorted: Ordered<_> = examples.into_iter().collect();
438        assert_eq!(
439            sorted.to_string(),
440            "[ex(1), ex(2), ex(3), ex(4), ex(5), ex(6), ex(7), ex(8), ex(9)]"
441        );
442    }
443
444    #[test]
445    fn test_ordered_from_slice() {
446        let items = [3u8, 1u8, 2u8, 2u8];
447        let ordered = Ordered::from(&items[..]);
448        assert_eq!(ordered.iter().copied().collect::<Vec<_>>(), vec![1, 2, 3]);
449    }
450
451    #[test]
452    fn test_ordered_from_iterator() {
453        let items = [3u8, 1u8, 2u8, 2u8];
454        let ordered = items.iter().copied().collect::<Ordered<_>>();
455        assert_eq!(ordered.iter().copied().collect::<Vec<_>>(), vec![1, 2, 3]);
456    }
457
458    #[test]
459    fn test_ordered_map_dedup_and_access() {
460        let items = vec![(3u8, "c"), (1u8, "a"), (2u8, "b"), (1u8, "duplicate")];
461
462        let map: OrderedAssociated<_, _> = items.into_iter().collect();
463
464        assert_eq!(map.len(), 3);
465        assert_eq!(map.iter().copied().collect::<Vec<_>>(), vec![1, 2, 3]);
466        assert_eq!(map.get_value(&1), Some(&"a"));
467        assert_eq!(map.get_value(&4), None);
468        assert_eq!(map.value(1), Some(&"b"));
469    }
470
471    #[test]
472    fn test_ordered_wrapped_from_slice() {
473        let pairs = [(3u8, "c"), (1u8, "a"), (2u8, "b")];
474        let wrapped = OrderedAssociated::from(&pairs[..]);
475
476        assert_eq!(wrapped.iter().copied().collect::<Vec<_>>(), vec![1, 2, 3]);
477        assert_eq!(wrapped.get_value(&2), Some(&"b"));
478    }
479
480    #[test]
481    fn test_ordered_wrapped_from_iterator() {
482        let pairs = [(3u8, "c"), (1u8, "a"), (2u8, "b")];
483        let wrapped = pairs
484            .iter()
485            .map(|(k, v)| (*k, *v))
486            .collect::<OrderedAssociated<_, _>>();
487
488        assert_eq!(wrapped.iter().copied().collect::<Vec<_>>(), vec![1, 2, 3]);
489        assert_eq!(wrapped.get_value(&1), Some(&"a"));
490    }
491
492    #[test]
493    fn test_ordered_map_deref_to_ordered() {
494        fn sum(set: &Ordered<u8>) -> u32 {
495            set.iter().map(|v| *v as u32).sum()
496        }
497
498        let map: OrderedAssociated<_, _> = vec![(2u8, "b"), (1u8, "a")].into_iter().collect();
499        assert_eq!(sum(&map), 3);
500    }
501
502    #[test]
503    fn test_ordered_map_from_ordered() {
504        let ordered: Ordered<_> = vec![(3u8, 'a'), (1u8, 'b'), (2u8, 'c')]
505            .into_iter()
506            .collect();
507        let wrapped: OrderedAssociated<_, _> = ordered.clone().into_iter().collect();
508
509        assert_eq!(
510            ordered.iter().map(|(k, _)| *k).collect::<Vec<_>>(),
511            wrapped.keys().iter().copied().collect::<Vec<_>>(),
512        );
513    }
514
515    #[test]
516    fn test_ordered_map_into_keys() {
517        let map: OrderedAssociated<_, _> = vec![(3u8, "c"), (1u8, "a"), (2u8, "b")]
518            .into_iter()
519            .collect();
520        let keys = map.into_keys();
521        assert_eq!(keys.iter().copied().collect::<Vec<_>>(), vec![1, 2, 3]);
522    }
523}