rkyv/impls/alloc/collections/
btree_map.rs

1use core::ops::ControlFlow;
2
3use rancor::{Fallible, Source};
4
5use crate::{
6    alloc::collections::BTreeMap,
7    collections::btree_map::{ArchivedBTreeMap, BTreeMapResolver},
8    ser::{Allocator, Writer},
9    Archive, Deserialize, Place, Serialize,
10};
11
12impl<K: Archive + Ord, V: Archive> Archive for BTreeMap<K, V>
13where
14    K::Archived: Ord,
15{
16    type Archived = ArchivedBTreeMap<K::Archived, V::Archived>;
17    type Resolver = BTreeMapResolver;
18
19    fn resolve(&self, resolver: Self::Resolver, out: Place<Self::Archived>) {
20        Self::Archived::resolve_from_len(self.len(), resolver, out);
21    }
22}
23
24impl<K, V, S> Serialize<S> for BTreeMap<K, V>
25where
26    K: Serialize<S> + Ord,
27    K::Archived: Ord,
28    V: Serialize<S>,
29    S: Allocator + Fallible + Writer + ?Sized,
30    S::Error: Source,
31{
32    fn serialize(
33        &self,
34        serializer: &mut S,
35    ) -> Result<Self::Resolver, S::Error> {
36        Self::Archived::serialize_from_ordered_iter::<_, _, _, K, V, _>(
37            self.iter(),
38            serializer,
39        )
40    }
41}
42
43impl<K, V, D> Deserialize<BTreeMap<K, V>, D>
44    for ArchivedBTreeMap<K::Archived, V::Archived>
45where
46    K: Archive + Ord,
47    K::Archived: Deserialize<K, D> + Ord,
48    V: Archive,
49    V::Archived: Deserialize<V, D>,
50    D: Fallible + ?Sized,
51{
52    fn deserialize(
53        &self,
54        deserializer: &mut D,
55    ) -> Result<BTreeMap<K, V>, D::Error> {
56        let mut result = BTreeMap::new();
57        let r = self.visit(|ak, av| {
58            let k = match ak.deserialize(deserializer) {
59                Ok(k) => k,
60                Err(e) => return ControlFlow::Break(e),
61            };
62            let v = match av.deserialize(deserializer) {
63                Ok(v) => v,
64                Err(e) => return ControlFlow::Break(e),
65            };
66            result.insert(k, v);
67            ControlFlow::Continue(())
68        });
69        match r {
70            Some(e) => Err(e),
71            None => Ok(result),
72        }
73    }
74}
75
76impl<K, V, AK, AV> PartialEq<BTreeMap<K, V>> for ArchivedBTreeMap<AK, AV>
77where
78    AK: PartialEq<K>,
79    AV: PartialEq<V>,
80{
81    fn eq(&self, other: &BTreeMap<K, V>) -> bool {
82        if self.len() != other.len() {
83            false
84        } else {
85            let mut iter = other.iter();
86            self.visit(|ak, av| {
87                if let Some((k, v)) = iter.next() {
88                    if ak.eq(k) && av.eq(v) {
89                        return ControlFlow::Continue(());
90                    }
91                }
92                ControlFlow::Break(())
93            })
94            .is_none()
95        }
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use core::ops::ControlFlow;
102
103    use crate::{
104        alloc::{
105            collections::BTreeMap,
106            string::{String, ToString},
107            vec,
108            vec::Vec,
109        },
110        api::test::{roundtrip, to_archived},
111        collections::btree_map::ArchivedBTreeMap,
112        primitive::ArchivedI32,
113        seal::Seal,
114        Archive, Deserialize, Serialize,
115    };
116
117    #[test]
118    fn roundtrip_btree_map() {
119        let mut value = BTreeMap::new();
120        value.insert("foo".to_string(), 10);
121        value.insert("bar".to_string(), 20);
122        value.insert("baz".to_string(), 40);
123        value.insert("bat".to_string(), 80);
124
125        roundtrip(&value);
126    }
127
128    #[test]
129    fn roundtrip_empty_btree_map() {
130        roundtrip(&BTreeMap::<String, i32>::new());
131    }
132
133    #[test]
134    fn roundtrip_btree_map_zst() {
135        let mut value = BTreeMap::new();
136        value.insert(0, ());
137        value.insert(1, ());
138        roundtrip(&value);
139
140        let mut value = BTreeMap::new();
141        value.insert((), 10);
142        roundtrip(&value);
143
144        let mut value = BTreeMap::new();
145        value.insert((), ());
146        roundtrip(&value);
147    }
148
149    #[test]
150    fn roundtrip_btree_map_increasing_sizes() {
151        // These sizes are chosen based on a branching factor of 6.
152        // 0-5: Leaf root node, variably-filled
153        // 6: Inner root node and one leaf node with one element
154        // 17: Inner root node, two filled leaf nodes, and one with two elements
155        // 35: Two full levels
156        // 36: Two full levels and one additional element
157        // 112: Two full levels and a third level half-filled
158        // 215: Three full levels
159        const SIZES: &[usize] = &[0, 1, 2, 3, 4, 5, 6, 17, 35, 36, 112, 215];
160        for &size in SIZES {
161            let mut value = BTreeMap::new();
162            for i in 0..size {
163                value.insert(i.to_string(), i as i32);
164            }
165
166            roundtrip(&value);
167        }
168    }
169
170    // This test creates structures too big to fit in 16-bit offsets, and MIRI
171    // can't run it quickly enough.
172    #[cfg(not(any(feature = "pointer_width_16", miri)))]
173    #[test]
174    fn roundtrip_large_btree_map() {
175        const ENTRIES: usize = 100_000;
176
177        let mut value = BTreeMap::new();
178        for i in 0..ENTRIES {
179            value.insert(i.to_string(), i as i32);
180        }
181
182        roundtrip(&value);
183    }
184
185    #[test]
186    fn roundtrip_btree_map_with_struct_member() {
187        #[derive(
188            Archive, Serialize, Deserialize, Debug, Default, PartialEq,
189        )]
190        #[rkyv(crate, compare(PartialEq), derive(Debug))]
191        pub struct MyType {
192            pub some_list: BTreeMap<String, Vec<f32>>,
193            pub values: Vec<f32>,
194        }
195
196        let mut value = MyType::default();
197
198        value
199            .some_list
200            .entry("Asdf".to_string())
201            .and_modify(|e| e.push(1.0))
202            .or_insert_with(|| vec![2.0]);
203
204        roundtrip(&value);
205    }
206
207    #[test]
208    fn mutable_btree_map() {
209        let mut value = BTreeMap::new();
210        value.insert("foo".to_string(), 10);
211        value.insert("bar".to_string(), 20);
212        value.insert("baz".to_string(), 40);
213        value.insert("bat".to_string(), 80);
214
215        to_archived(&value, |mut archived| {
216            ArchivedBTreeMap::visit_seal(
217                archived.as_mut(),
218                |_, mut v: Seal<'_, ArchivedI32>| {
219                    *v = ArchivedI32::from_native(v.to_native() + 10);
220                    ControlFlow::<(), ()>::Continue(())
221                },
222            );
223            assert_eq!(archived.get("foo").map(|x| x.to_native()), Some(20));
224            assert_eq!(archived.get("bar").map(|x| x.to_native()), Some(30));
225            assert_eq!(archived.get("baz").map(|x| x.to_native()), Some(50));
226            assert_eq!(archived.get("bat").map(|x| x.to_native()), Some(90));
227
228            *ArchivedBTreeMap::get_seal(archived.as_mut(), "foo").unwrap() =
229                ArchivedI32::from_native(123);
230            *ArchivedBTreeMap::get_seal(archived.as_mut(), "bat").unwrap() =
231                ArchivedI32::from_native(456);
232
233            assert_eq!(archived.get("foo").map(|x| x.to_native()), Some(123));
234            assert_eq!(archived.get("bat").map(|x| x.to_native()), Some(456));
235        });
236    }
237
238    #[test]
239    fn btree_map_iter() {
240        let mut value = BTreeMap::<String, i32>::new();
241        value.insert("foo".to_string(), 10);
242        value.insert("bar".to_string(), 20);
243        value.insert("baz".to_string(), 40);
244        value.insert("bat".to_string(), 80);
245
246        to_archived(&value, |archived| {
247            let mut i =
248                archived.iter().map(|(k, v)| (k.as_str(), v.to_native()));
249            assert_eq!(i.next(), Some(("bar", 20)));
250            assert_eq!(i.next(), Some(("bat", 80)));
251            assert_eq!(i.next(), Some(("baz", 40)));
252            assert_eq!(i.next(), Some(("foo", 10)));
253            assert_eq!(i.next(), None);
254        });
255    }
256
257    #[test]
258    fn btree_map_mutable_iter() {
259        let mut value = BTreeMap::<String, i32>::new();
260        value.insert("foo".to_string(), 10);
261        value.insert("bar".to_string(), 20);
262        value.insert("baz".to_string(), 40);
263        value.insert("bat".to_string(), 80);
264
265        to_archived(&value, |archived| {
266            let mut i =
267                archived.iter().map(|(k, v)| (k.as_str(), v.to_native()));
268            assert_eq!(i.next(), Some(("bar", 20)));
269            assert_eq!(i.next(), Some(("bat", 80)));
270            assert_eq!(i.next(), Some(("baz", 40)));
271            assert_eq!(i.next(), Some(("foo", 10)));
272            assert_eq!(i.next(), None);
273        });
274    }
275
276    // MIRI can't run this test quickly enough
277    #[cfg(not(miri))]
278    #[test]
279    fn size_matches_iter_len() {
280        #[derive(Archive, Deserialize, Serialize)]
281        #[rkyv(crate)]
282        struct Container {
283            transforms: BTreeMap<String, String>,
284        }
285
286        impl Container {
287            pub fn fill(count: usize) -> Self {
288                let mut transforms = BTreeMap::new();
289                for i in 0..count {
290                    transforms.insert(i.to_string(), i.to_string());
291                }
292                Container { transforms }
293            }
294
295            pub fn check(&self, expected: usize) {
296                to_archived(self, |archived| {
297                    assert_eq!(archived.transforms.len(), expected);
298                    let mut count = 0;
299                    for (..) in archived.transforms.iter() {
300                        count += 1;
301                    }
302                    assert_eq!(count, expected);
303                });
304            }
305        }
306
307        for expected in 0..=200 {
308            let container = Container::fill(expected);
309            container.check(expected);
310        }
311    }
312}
rkyv/impls/alloc/collections/btree_map.rs

rkyv/impls/alloc/collections/
btree_map.rs
