messagepack_core/encode/
map.rs

1//! Map encoders.
2
3use core::{cell::RefCell, marker::PhantomData, ops::Deref};
4
5use super::{Encode, Error, Result};
6use crate::{formats::Format, io::IoWrite};
7
8/// A key-value encoder that writes a single `key, value` pair.
9pub trait KVEncode<W>
10where
11    W: IoWrite,
12{
13    /// Encode this key‑value pair to the writer and return the number of bytes written.
14    fn encode(&self, writer: &mut W) -> Result<usize, W::Error>;
15}
16
17impl<W: IoWrite, KV: KVEncode<W>> KVEncode<W> for &KV {
18    fn encode(&self, writer: &mut W) -> Result<usize, W::Error> {
19        KV::encode(self, writer)
20    }
21}
22
23impl<W: IoWrite, K: Encode<W>, V: Encode<W>> KVEncode<W> for (K, V) {
24    fn encode(&self, writer: &mut W) -> Result<usize, W::Error> {
25        let (k, v) = self;
26        let k_len = k.encode(writer)?;
27        let v_len = v.encode(writer)?;
28        Ok(k_len + v_len)
29    }
30}
31
32/// Encode only the map header for a map of a given length.
33pub struct MapFormatEncoder(pub usize);
34impl MapFormatEncoder {
35    /// Construct from the number of pairs contained in the map.
36    pub fn new(size: usize) -> Self {
37        Self(size)
38    }
39}
40
41impl<W: IoWrite> Encode<W> for MapFormatEncoder {
42    fn encode(&self, writer: &mut W) -> Result<usize, W::Error> {
43        match self.0 {
44            0x00..=0xf => {
45                let cast = self.0 as u8;
46                writer.write(&[Format::FixMap(cast).as_byte()])?;
47
48                Ok(1)
49            }
50            0x10..=0xffff => {
51                let cast = (self.0 as u16).to_be_bytes();
52                writer.write(&[Format::Map16.as_byte(), cast[0], cast[1]])?;
53
54                Ok(3)
55            }
56            0x10000..=0xffffffff => {
57                let cast = (self.0 as u32).to_be_bytes();
58                writer.write(&[Format::Map32.as_byte(), cast[0], cast[1], cast[2], cast[3]])?;
59
60                Ok(5)
61            }
62            _ => Err(Error::InvalidFormat),
63        }
64    }
65}
66
67/// Encode a stream of key-value pairs from an iterator.
68pub struct MapDataEncoder<I, J, KV> {
69    data: RefCell<J>,
70    _phantom: PhantomData<(I, J, KV)>,
71}
72
73impl<I, KV> MapDataEncoder<I, I::IntoIter, KV>
74where
75    I: IntoIterator<Item = KV>,
76{
77    /// Construct from any iterable of key-value pairs.
78    pub fn new(data: I) -> Self {
79        Self {
80            data: RefCell::new(data.into_iter()),
81            _phantom: Default::default(),
82        }
83    }
84}
85
86impl<W, I, J, KV> Encode<W> for MapDataEncoder<I, J, KV>
87where
88    W: IoWrite,
89    J: Iterator<Item = KV>,
90    KV: KVEncode<W>,
91{
92    fn encode(&self, writer: &mut W) -> Result<usize, W::Error> {
93        let map_len = self
94            .data
95            .borrow_mut()
96            .by_ref()
97            .map(|kv| kv.encode(writer))
98            .try_fold(0, |acc, v| v.map(|n| acc + n))?;
99        Ok(map_len)
100    }
101}
102
103/// Encode a slice of key-value pairs.
104pub struct MapSliceEncoder<'data, KV> {
105    data: &'data [KV],
106    _phantom: PhantomData<KV>,
107}
108
109impl<'data, KV> MapSliceEncoder<'data, KV> {
110    /// Construct from a slice of key-value pairs.
111    pub fn new(data: &'data [KV]) -> Self {
112        Self {
113            data,
114            _phantom: Default::default(),
115        }
116    }
117}
118
119impl<'data, KV> Deref for MapSliceEncoder<'data, KV> {
120    type Target = &'data [KV];
121    fn deref(&self) -> &Self::Target {
122        &self.data
123    }
124}
125
126impl<W, KV> Encode<W> for MapSliceEncoder<'_, KV>
127where
128    W: IoWrite,
129    KV: KVEncode<W>,
130{
131    fn encode(&self, writer: &mut W) -> Result<usize, W::Error> {
132        let self_len = self.data.len();
133        let format_len = MapFormatEncoder::new(self_len).encode(writer)?;
134        let map_len = MapDataEncoder::new(self.data.iter()).encode(writer)?;
135
136        Ok(format_len + map_len)
137    }
138}
139
140/// Encode a map from an owned iterator, writing items lazily.
141pub struct MapEncoder<W, I, J, KV> {
142    map: RefCell<J>,
143    _phantom: PhantomData<(W, I, J, KV)>,
144}
145
146impl<W, I, KV> MapEncoder<W, I, I::IntoIter, KV>
147where
148    W: IoWrite,
149    I: IntoIterator<Item = KV>,
150    KV: KVEncode<W>,
151{
152    /// Construct from any iterable of key-value pairs.
153    pub fn new(map: I) -> Self {
154        Self {
155            map: RefCell::new(map.into_iter()),
156            _phantom: Default::default(),
157        }
158    }
159}
160
161impl<W, I, J, KV> Encode<W> for MapEncoder<W, I, J, KV>
162where
163    W: IoWrite,
164    J: Iterator<Item = KV> + ExactSizeIterator,
165    KV: KVEncode<W>,
166{
167    fn encode(&self, writer: &mut W) -> Result<usize, W::Error> {
168        let self_len = self.map.borrow().len();
169        let format_len = MapFormatEncoder::new(self_len).encode(writer)?;
170        let map_len = MapDataEncoder::new(self.map.borrow_mut().by_ref()).encode(writer)?;
171
172        Ok(format_len + map_len)
173    }
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179    use crate::encode::int::EncodeMinimizeInt;
180    use rstest::rstest;
181
182    #[rstest]
183    #[case([("123", EncodeMinimizeInt(123)), ("456", EncodeMinimizeInt(456))], [0x82, 0xa3, 0x31, 0x32, 0x33, 0x7b, 0xa3, 0x34, 0x35, 0x36, 0xcd, 0x01, 0xc8])]
184    fn encode_slice_fix_array<K, V, Map, E>(#[case] value: Map, #[case] expected: E)
185    where
186        K: Encode<Vec<u8>>,
187        V: Encode<Vec<u8>>,
188        Map: AsRef<[(K, V)]>,
189        E: AsRef<[u8]> + Sized,
190    {
191        let expected = expected.as_ref();
192        let encoder = MapSliceEncoder::new(value.as_ref());
193
194        let mut buf = vec![];
195        let n = encoder.encode(&mut buf).unwrap();
196        assert_eq!(buf, expected);
197        assert_eq!(n, expected.len());
198    }
199
200    #[rstest]
201    #[case([("123", EncodeMinimizeInt(123)), ("456", EncodeMinimizeInt(456))], [0x82, 0xa3, 0x31, 0x32, 0x33, 0x7b, 0xa3, 0x34, 0x35, 0x36, 0xcd, 0x01, 0xc8])]
202    fn encode_iter_fix_array<I, KV, E>(#[case] value: I, #[case] expected: E)
203    where
204        I: IntoIterator<Item = KV>,
205        I::IntoIter: ExactSizeIterator,
206        KV: KVEncode<Vec<u8>>,
207        E: AsRef<[u8]> + Sized,
208    {
209        let expected = expected.as_ref();
210
211        let encoder = MapEncoder::new(value.into_iter());
212        let mut buf = vec![];
213        let n = encoder.encode(&mut buf).unwrap();
214        assert_eq!(buf, expected);
215        assert_eq!(n, expected.len());
216    }
217}
messagepack_core/encode/map.rs

messagepack_core/encode/
map.rs
