roaring/treemap/
serialization.rs

1use super::RoaringTreemap;
2use crate::RoaringBitmap;
3use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
4use std::{io, mem::size_of};
5
6impl RoaringTreemap {
7    /// Return the size in bytes of the serialized output.
8    /// This is compatible with the official C/C++, Java and Go implementations.
9    ///
10    /// # Examples
11    ///
12    /// ```rust
13    /// use roaring::RoaringTreemap;
14    ///
15    /// let rb1: RoaringTreemap = (1..4).collect();
16    /// let mut bytes = Vec::with_capacity(rb1.serialized_size());
17    /// rb1.serialize_into(&mut bytes).unwrap();
18    /// let rb2 = RoaringTreemap::deserialize_from(&bytes[..]).unwrap();
19    ///
20    /// assert_eq!(rb1, rb2);
21    /// ```
22    pub fn serialized_size(&self) -> usize {
23        self.map
24            .values()
25            .fold(size_of::<u64>(), |acc, bitmap| acc + size_of::<u32>() + bitmap.serialized_size())
26    }
27
28    /// Serialize this bitmap.
29    /// This is compatible with the official C/C++, Java and Go implementations.
30    ///
31    /// # Examples
32    ///
33    /// ```rust
34    /// use roaring::RoaringTreemap;
35    ///
36    /// let rb1: RoaringTreemap = (1..4).collect();
37    /// let mut bytes = vec![];
38    /// rb1.serialize_into(&mut bytes).unwrap();
39    /// let rb2 = RoaringTreemap::deserialize_from(&bytes[..]).unwrap();
40    ///
41    /// assert_eq!(rb1, rb2);
42    /// ```
43    pub fn serialize_into<W: io::Write>(&self, mut writer: W) -> io::Result<()> {
44        writer.write_u64::<LittleEndian>(self.map.len() as u64)?;
45
46        for (key, bitmap) in &self.map {
47            writer.write_u32::<LittleEndian>(*key)?;
48            bitmap.serialize_into(&mut writer)?;
49        }
50
51        Ok(())
52    }
53
54    /// Deserialize a bitmap into memory.
55    ///
56    /// This is compatible with the official C/C++, Java and Go implementations.
57    /// This method checks that all of the internal values are valid.
58    ///
59    /// # Examples
60    ///
61    /// ```rust
62    /// use roaring::RoaringTreemap;
63    ///
64    /// let rb1: RoaringTreemap = (1..4).collect();
65    /// let mut bytes = vec![];
66    /// rb1.serialize_into(&mut bytes).unwrap();
67    /// let rb2 = RoaringTreemap::deserialize_from(&bytes[..]).unwrap();
68    ///
69    /// assert_eq!(rb1, rb2);
70    /// ```
71    pub fn deserialize_from<R: io::Read>(reader: R) -> io::Result<Self> {
72        RoaringTreemap::deserialize_from_impl(reader, |reader| {
73            RoaringBitmap::deserialize_from(reader)
74        })
75    }
76
77    /// Deserialize a bitmap into memory.
78    ///
79    /// This is compatible with the official C/C++, Java and Go implementations.
80    /// This method is memory safe but will not check if the data is a valid bitmap.
81    ///
82    /// # Examples
83    ///
84    /// ```rust
85    /// use roaring::RoaringTreemap;
86    ///
87    /// let rb1: RoaringTreemap = (1..4).collect();
88    /// let mut bytes = vec![];
89    /// rb1.serialize_into(&mut bytes).unwrap();
90    /// let rb2 = RoaringTreemap::deserialize_unchecked_from(&bytes[..]).unwrap();
91    ///
92    /// assert_eq!(rb1, rb2);
93    /// ```
94    pub fn deserialize_unchecked_from<R: io::Read>(reader: R) -> io::Result<Self> {
95        RoaringTreemap::deserialize_from_impl(reader, |reader| {
96            RoaringBitmap::deserialize_unchecked_from(reader)
97        })
98    }
99
100    fn deserialize_from_impl<R, F>(mut reader: R, mut deserialize_bitmap: F) -> io::Result<Self>
101    where
102        R: io::Read,
103        F: FnMut(&mut R) -> io::Result<RoaringBitmap>,
104    {
105        let size = reader.read_u64::<LittleEndian>()?;
106
107        let mut s = Self::new();
108
109        for _ in 0..size {
110            let key = reader.read_u32::<LittleEndian>()?;
111            let bitmap = deserialize_bitmap(&mut reader)?;
112
113            s.map.insert(key, bitmap);
114        }
115
116        Ok(s)
117    }
118}
119
120#[cfg(test)]
121mod test {
122    use crate::RoaringTreemap;
123    use proptest::prelude::*;
124
125    proptest! {
126        #[test]
127        fn test_serialization(
128            treemap in RoaringTreemap::arbitrary(),
129        ) {
130            let mut buffer = Vec::new();
131            treemap.serialize_into(&mut buffer).unwrap();
132            prop_assert_eq!(treemap, RoaringTreemap::deserialize_from(buffer.as_slice()).unwrap());
133        }
134    }
135}