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
18pub enum Size {
20 Usize,
22 U8,
24 U16,
26 U32,
28 U64,
30 U128,
32}
33
34pub trait SizeSettings: EndianSettings {
36 const SIZE: Size;
38
39 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 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);