Skip to main content

data_stream/
collections.rs

1use crate::{FromStream, ToStream, from_stream, numbers::EndianSettings, to_stream};
2
3use std::{
4    collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet, LinkedList, VecDeque},
5    hash::Hash,
6    io::{Error, ErrorKind, Read, Result, Write},
7};
8
9fn checked_size<T: TryFrom<usize>>(size: usize) -> Result<T> {
10    T::try_from(size).map_err(|_| {
11        Error::new(
12            ErrorKind::InvalidInput,
13            "collection size exceeds the configured size encoding",
14        )
15    })
16}
17
18/// The encoding for sizes for collections.
19pub enum Size {
20    /// Encodes sizes as `usize`.
21    Usize,
22    /// Encodes sizes as `u8`.
23    U8,
24    /// Encodes sizes as `u16`.
25    U16,
26    /// Encodes sizes as `u32`.
27    U32,
28    /// Encodes sizes as `u64`.
29    U64,
30    /// Encodes sizes as `u128`.
31    U128,
32}
33
34/// The size settings required for reading and writing collections.
35pub trait SizeSettings: EndianSettings {
36    /// The size used by these size settings.
37    const SIZE: Size;
38
39    /// Writes some size using specific settings.
40    fn size_to_stream<W: Write>(size: usize, stream: &mut W) -> Result<()> {
41        use Size::*;
42        match Self::SIZE {
43            U8 => to_stream::<Self, _, _>(&checked_size::<u8>(size)?, stream),
44            U16 => to_stream::<Self, _, _>(&checked_size::<u16>(size)?, stream),
45            U32 => to_stream::<Self, _, _>(&checked_size::<u32>(size)?, stream),
46            U64 => to_stream::<Self, _, _>(&checked_size::<u64>(size)?, stream),
47            U128 => to_stream::<Self, _, _>(&checked_size::<u128>(size)?, stream),
48            Usize => to_stream::<Self, _, _>(&size, stream),
49        }
50    }
51
52    /// Reads some size using specific settings.
53    fn size_from_stream<R: Read>(stream: &mut R) -> Result<usize> {
54        use Size::*;
55        Ok(match Self::SIZE {
56            U8 => from_stream::<Self, u8, _>(stream)? as usize,
57            U16 => from_stream::<Self, u16, _>(stream)? as usize,
58            U32 => from_stream::<Self, u32, _>(stream)? as usize,
59            U64 => from_stream::<Self, u64, _>(stream)? as usize,
60            U128 => from_stream::<Self, u128, _>(stream)? as usize,
61            Usize => from_stream::<Self, usize, _>(stream)?,
62        })
63    }
64}
65
66macro_rules! impl_list_collection {
67    ($t: ident) => {
68        impl<T: ToStream<S>, S: SizeSettings> ToStream<S> for $t<T> {
69            fn to_stream<W: Write>(&self, stream: &mut W) -> Result<()> {
70                S::size_to_stream(self.len(), stream)?;
71
72                for element in self {
73                    element.to_stream(stream)?
74                }
75
76                Ok(())
77            }
78        }
79
80        impl<T: FromStream<S>, S: SizeSettings> FromStream<S> for $t<T> {
81            fn from_stream<R: Read>(stream: &mut R) -> Result<Self> {
82                let size = S::size_from_stream(stream)?;
83
84                (0..size)
85                    .map(|_| from_stream(stream))
86                    .collect::<Result<_>>()
87            }
88        }
89    };
90    ($t: ident, $($rest: ident),+) => {
91        impl_list_collection!($t);
92        impl_list_collection!($($rest),+);
93    };
94}
95
96impl_list_collection!(Vec, VecDeque, LinkedList);
97
98macro_rules! impl_set_collection {
99    ($t: ident, $bound: tt $(+ $others: tt )*) => {
100        impl<T: ToStream<S> + $bound $(+ $others)*, S: SizeSettings> ToStream<S> for $t<T> {
101            fn to_stream<W: Write>(&self, stream: &mut W) -> Result<()> {
102                S::size_to_stream(self.len(), stream)?;
103
104                for element in self {
105                    element.to_stream(stream)?
106                }
107
108                Ok(())
109            }
110        }
111
112        impl<T: FromStream<S> + $bound $(+ $others)*, S: SizeSettings> FromStream<S> for $t<T> {
113            fn from_stream<R: Read>(stream: &mut R) -> Result<Self> {
114                let size = S::size_from_stream(stream)?;
115
116                (0..size)
117                    .map(|_| from_stream(stream))
118                    .collect::<Result<_>>()
119            }
120        }
121    };
122}
123
124impl_set_collection!(HashSet, Eq + Hash);
125impl_set_collection!(BTreeSet, Ord);
126impl_set_collection!(BinaryHeap, Ord);
127
128macro_rules! impl_map_collection {
129    ($t: ident, $bound: tt $(+ $others: tt )*) => {
130        impl<K: ToStream<S> + $bound $(+ $others)*, V: ToStream<S>, S: SizeSettings> ToStream<S> for $t<K, V> {
131            fn to_stream<W: Write>(&self, stream: &mut W) -> Result<()> {
132                S::size_to_stream(self.len(), stream)?;
133
134                for (key, value) in self {
135                    key.to_stream(stream)?;
136                    value.to_stream(stream)?;
137                }
138
139                Ok(())
140            }
141        }
142
143        impl<K: FromStream<S> + $bound $(+ $others)*, V: FromStream<S>, S: SizeSettings> FromStream<S> for $t<K, V> {
144            fn from_stream<R: Read>(stream: &mut R) -> Result<Self> {
145                let size = S::size_from_stream(stream)?;
146
147                (0..size)
148                    .map(|_| Ok((from_stream(stream)?, from_stream(stream)?)))
149                    .collect::<Result<_>>()
150            }
151        }
152    };
153}
154
155impl_map_collection!(HashMap, Eq + Hash);
156impl_map_collection!(BTreeMap, Ord);