commonware_codec/types/
btree_set.rs

1//! Codec implementations for no_std compatible set types.
2//!
3//! For portability and consistency between architectures,
4//! the size of the set must fit within a [u32].
5
6extern crate alloc;
7
8use crate::{
9    codec::{EncodeSize, Read, Write},
10    error::Error,
11    types::read_ordered_set,
12    RangeCfg,
13};
14use alloc::collections::BTreeSet;
15use bytes::{Buf, BufMut};
16
17const BTREESET_TYPE: &str = "BTreeSet";
18
19// ---------- BTreeSet ----------
20
21impl<K: Ord + Eq + Write> Write for BTreeSet<K> {
22    fn write(&self, buf: &mut impl BufMut) {
23        self.len().write(buf);
24
25        // Items are already sorted in BTreeSet, so we can iterate directly
26        for item in self {
27            item.write(buf);
28        }
29    }
30}
31
32impl<K: Ord + Eq + EncodeSize> EncodeSize for BTreeSet<K> {
33    fn encode_size(&self) -> usize {
34        let mut size = self.len().encode_size();
35        for item in self {
36            size += item.encode_size();
37        }
38        size
39    }
40}
41
42impl<K: Read + Clone + Ord + Eq> Read for BTreeSet<K> {
43    type Cfg = (RangeCfg<usize>, K::Cfg);
44
45    fn read_cfg(buf: &mut impl Buf, (range, cfg): &Self::Cfg) -> Result<Self, Error> {
46        // Read and validate the length prefix
47        let len = usize::read_cfg(buf, range)?;
48        let mut set = Self::new();
49
50        // Read items in ascending order
51        read_ordered_set(buf, len, cfg, |item| set.insert(item), BTREESET_TYPE)?;
52
53        Ok(set)
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60    use crate::{
61        codec::{Decode, Encode},
62        FixedSize,
63    };
64    use bytes::{Bytes, BytesMut};
65    use core::fmt::Debug;
66
67    // Generic round trip test function for BTreeSet
68    fn round_trip_btree<K>(set: &BTreeSet<K>, range_cfg: RangeCfg<usize>, item_cfg: K::Cfg)
69    where
70        K: Write + EncodeSize + Read + Clone + Ord + Eq + Debug + PartialEq,
71        BTreeSet<K>: Read<Cfg = (RangeCfg<usize>, K::Cfg)>
72            + Decode<Cfg = (RangeCfg<usize>, K::Cfg)>
73            + Debug
74            + PartialEq
75            + Write
76            + EncodeSize,
77    {
78        let encoded = set.encode();
79        assert_eq!(set.encode_size(), encoded.len());
80        let config_tuple = (range_cfg, item_cfg);
81        let decoded = BTreeSet::<K>::decode_cfg(encoded, &config_tuple).expect("decode_cfg failed");
82        assert_eq!(set, &decoded);
83    }
84
85    #[test]
86    fn test_empty_btreeset() {
87        let set = BTreeSet::<u32>::new();
88        round_trip_btree(&set, (..).into(), ());
89        assert_eq!(set.encode_size(), 1); // varint 0
90        let encoded = set.encode();
91        assert_eq!(encoded, Bytes::from_static(&[0]));
92    }
93
94    #[test]
95    fn test_simple_btreeset_u32() {
96        let mut set = BTreeSet::new();
97        set.insert(1u32);
98        set.insert(5u32);
99        set.insert(2u32);
100        round_trip_btree(&set, (..).into(), ());
101        assert_eq!(set.encode_size(), 1 + 3 * u32::SIZE);
102    }
103
104    #[test]
105    fn test_large_btreeset() {
106        // Fixed-size items
107        let set: BTreeSet<_> = (0..1000u16).collect();
108        round_trip_btree(&set, (1000..=1000).into(), ());
109
110        // Variable-size items
111        let set: BTreeSet<_> = (0..1000usize).collect();
112        round_trip_btree(&set, (1000..=1000).into(), (..=1000).into());
113    }
114
115    #[test]
116    fn test_btreeset_decode_length_limit_exceeded() {
117        let mut set = BTreeSet::new();
118        set.insert(1u32);
119        set.insert(5u32);
120        let encoded = set.encode();
121
122        let config_tuple = ((0..=1).into(), ());
123        let result = BTreeSet::<u32>::decode_cfg(encoded, &config_tuple);
124        assert!(matches!(result, Err(Error::InvalidLength(2))));
125    }
126
127    #[test]
128    fn test_btreeset_decode_invalid_item_order() {
129        let mut encoded = BytesMut::new();
130        2usize.write(&mut encoded); // Set length = 2
131        5u32.write(&mut encoded); // Item 5
132        2u32.write(&mut encoded); // Item 2 (out of order)
133
134        let config_tuple = ((..).into(), ());
135        let result = BTreeSet::<u32>::decode_cfg(encoded, &config_tuple);
136        assert!(matches!(
137            result,
138            Err(Error::Invalid("BTreeSet", "Items must ascend"))
139        ));
140    }
141
142    #[test]
143    fn test_btreeset_decode_duplicate_item() {
144        let mut encoded = BytesMut::new();
145        2usize.write(&mut encoded); // Set length = 2
146        1u32.write(&mut encoded); // Item 1
147        1u32.write(&mut encoded); // Duplicate Item 1
148
149        let config_tuple = ((..).into(), ());
150        let result = BTreeSet::<u32>::decode_cfg(encoded, &config_tuple);
151        assert!(matches!(
152            result,
153            Err(Error::Invalid("BTreeSet", "Duplicate item"))
154        ));
155    }
156
157    #[test]
158    fn test_btreeset_deterministic_encoding() {
159        let mut set1 = BTreeSet::new();
160        (0..1000u32).for_each(|i| {
161            set1.insert(i);
162        });
163
164        let mut set2 = BTreeSet::new();
165        (0..1000u32).rev().for_each(|i| {
166            set2.insert(i);
167        });
168
169        assert_eq!(set1.encode(), set2.encode());
170    }
171
172    #[cfg(feature = "arbitrary")]
173    mod conformance {
174        use super::*;
175        use crate::conformance::CodecConformance;
176
177        commonware_conformance::conformance_tests! {
178            CodecConformance<BTreeSet<u8>>,
179        }
180    }
181}