Skip to main content

serde_multipart/
part.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9use std::ops::Deref;
10
11use bincode::Options;
12use bytes::Bytes;
13use bytes::BytesMut;
14use bytes::buf::Reader as BufReader;
15use bytes::buf::Writer as BufWriter;
16use serde::Deserialize;
17use serde::Serialize;
18use serde::de::DeserializeOwned;
19use typeuri::Named;
20
21use crate::UnsafeBufCellRef;
22use crate::de;
23use crate::ser;
24
25/// Part represents a single part of a multipart message.
26///
27/// A part is backed by byte fragments, which permit zero-copy shared ownership
28/// of the underlying buffers. It may also carry a typehash for framework-owned,
29/// bincode-serialized values.
30#[derive(Clone, Debug, PartialEq, Eq, Default)]
31pub struct Part {
32    typehash: Option<u64>,
33    fragments: Vec<Bytes>,
34}
35
36impl Part {
37    /// Consumes the part, returning its underlying byte fragments.
38    pub fn into_fragments(self) -> Vec<Bytes> {
39        self.fragments
40    }
41
42    /// Consumes the part, concatenating fragments if necessary into a single byte buffer.
43    pub fn into_bytes(self) -> Bytes {
44        match self.fragments.len() {
45            0 => Bytes::new(),
46            1 => self.fragments.into_iter().next().unwrap(),
47            _ => {
48                let total_len: usize = self.fragments.iter().map(|p| p.len()).sum();
49                let mut result = BytesMut::with_capacity(total_len);
50                for fragment in self.fragments {
51                    result.extend_from_slice(&fragment);
52                }
53                result.freeze()
54            }
55        }
56    }
57
58    /// Get bytes as a reference, concatenating fragments if necessary.
59    pub fn to_bytes(&self) -> Bytes {
60        match self.fragments.len() {
61            0 => Bytes::new(),
62            1 => self.fragments.first().unwrap().clone(),
63            _ => {
64                let total_len: usize = self.fragments.iter().map(|p| p.len()).sum();
65                let mut result = BytesMut::with_capacity(total_len);
66                for fragment in &self.fragments {
67                    result.extend_from_slice(fragment);
68                }
69                result.freeze()
70            }
71        }
72    }
73
74    /// Returns the total length in bytes.
75    pub fn len(&self) -> usize {
76        self.fragments
77            .iter()
78            .try_fold(0usize, |len, fragment| len.checked_add(fragment.len()))
79            .expect("part length exceeds usize")
80    }
81
82    /// Returns the number of fragments
83    pub fn num_fragments(&self) -> usize {
84        self.fragments.len()
85    }
86
87    /// Returns whether the part is empty.
88    pub fn is_empty(&self) -> bool {
89        self.fragments.iter().all(|b| b.is_empty())
90    }
91
92    /// Returns the typehash carried by this part, if any.
93    pub fn typehash(&self) -> Option<u64> {
94        self.typehash
95    }
96
97    /// Returns whether this part contains a serialized `T`-typed value.
98    pub fn is<T: Named>(&self) -> bool {
99        self.typehash == Some(T::typehash())
100    }
101
102    /// Serialize a value into a typed part.
103    pub fn serialize<T: Serialize + Named>(value: &T) -> crate::Result<Self> {
104        Self::serialize_as::<T, T>(value)
105    }
106
107    /// Serialize a value into a part carrying `T`'s typehash.
108    pub fn serialize_as<T, U>(value: &U) -> crate::Result<Self>
109    where
110        T: Named,
111        U: Serialize,
112    {
113        Ok(Self {
114            typehash: Some(T::typehash()),
115            fragments: vec![Bytes::from(crate::options().serialize(value)?)],
116        })
117    }
118
119    /// Deserialize this part into the provided type `T`.
120    pub fn deserialized<T: DeserializeOwned + Named>(&self) -> crate::Result<T> {
121        self.deserialized_as::<T, T>()
122    }
123
124    /// Deserialize this part as `U` after checking that it carries `T`'s typehash.
125    pub fn deserialized_as<T, U>(&self) -> crate::Result<U>
126    where
127        T: Named,
128        U: DeserializeOwned,
129    {
130        if !self.is::<T>() {
131            return Err(crate::Error::TypeMismatch {
132                expected: T::typename(),
133                actual: self
134                    .typehash
135                    .map_or_else(|| "unknown".to_string(), |typehash| typehash.to_string()),
136            });
137        }
138        self.deserialized_unchecked()
139    }
140
141    /// Deserialize this part without checking its typehash.
142    pub fn deserialized_unchecked<T: DeserializeOwned>(&self) -> crate::Result<T> {
143        Ok(crate::options().deserialize(&self.to_bytes())?)
144    }
145
146    /// Returns a part from byte fragments.
147    pub fn from_fragments(fragments: Vec<Bytes>) -> Self {
148        Self {
149            typehash: None,
150            fragments,
151        }
152    }
153
154    /// Returns a part from an optional typehash and byte fragments.
155    pub(crate) fn from_typehash_and_fragments(
156        typehash: Option<u64>,
157        fragments: Vec<Bytes>,
158    ) -> Self {
159        Self {
160            typehash,
161            fragments,
162        }
163    }
164}
165
166impl<T: Into<Bytes>> From<T> for Part {
167    fn from(bytes: T) -> Self {
168        Self {
169            typehash: None,
170            fragments: vec![bytes.into()],
171        }
172    }
173}
174
175impl Deref for Part {
176    type Target = Vec<Bytes>;
177
178    fn deref(&self) -> &Self::Target {
179        &self.fragments
180    }
181}
182
183impl Serialize for Part {
184    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
185        <Part as PartSerializer<S>>::serialize(self, s)
186    }
187}
188
189impl<'de> Deserialize<'de> for Part {
190    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
191        <Part as PartDeserializer<'de, D>>::deserialize(d)
192    }
193}
194
195/// PartSerializer is the trait that selects serialization strategy based on the
196/// the serializer's type.
197pub trait PartSerializer<S: serde::Serializer> {
198    fn serialize(this: &Part, s: S) -> Result<S::Ok, S::Error>;
199}
200
201/// By default, we use the underlying byte serializer, which copies the underlying bytes
202/// into the serialization buffer.
203impl<S: serde::Serializer> PartSerializer<S> for Part {
204    default fn serialize(this: &Part, s: S) -> Result<S::Ok, S::Error> {
205        (&this.typehash, &this.fragments).serialize(s)
206    }
207}
208
209/// The options type used by the underlying bincode codec. We capture this here to make sure
210/// we consistently use the type, which is required to correctly specialize the multipart codec.
211pub(crate) type BincodeOptionsType = bincode::config::WithOtherTrailing<
212    bincode::config::WithOtherIntEncoding<bincode::DefaultOptions, bincode::config::FixintEncoding>,
213    bincode::config::AllowTrailing,
214>;
215
216/// The serializer type used by the underlying bincode codec. We capture this here to make sure
217/// we consistently use the type, which is required to correctly specialize the multipart codec.
218pub(crate) type BincodeSerializer =
219    ser::bincode::Serializer<BufWriter<UnsafeBufCellRef>, BincodeOptionsType>;
220
221/// Specialized implementaiton for our multipart serializer.
222impl<'a> PartSerializer<&'a mut BincodeSerializer> for Part {
223    fn serialize(this: &Part, s: &'a mut BincodeSerializer) -> Result<(), bincode::Error> {
224        s.serialize_part(this);
225        Ok(())
226    }
227}
228
229/// PartDeserializer is the trait that selects serialization strategy based on the
230/// the deserializer's type.
231trait PartDeserializer<'de, S: serde::Deserializer<'de>>: Sized {
232    fn deserialize(this: S) -> Result<Self, S::Error>;
233}
234
235/// By default, we use the underlying byte deserializer, which copies the serialized bytes
236/// into the value directly.
237impl<'de, D: serde::Deserializer<'de>> PartDeserializer<'de, D> for Part {
238    default fn deserialize(deserializer: D) -> Result<Self, D::Error> {
239        let (typehash, fragments) = <(Option<u64>, Vec<Bytes>)>::deserialize(deserializer)?;
240        Ok(Self {
241            typehash,
242            fragments,
243        })
244    }
245}
246
247/// The deserializer type used by the underlying bincode codec. We capture this here to make sure
248/// we consistently use the type, which is required to correctly specialize the multipart codec.
249pub(crate) type BincodeDeserializer =
250    de::bincode::Deserializer<bincode::de::read::IoReader<BufReader<Bytes>>, BincodeOptionsType>;
251
252/// Specialized implementation for our multipart deserializer.
253impl<'a> PartDeserializer<'_, &'a mut BincodeDeserializer> for Part {
254    fn deserialize(deserializer: &'a mut BincodeDeserializer) -> Result<Self, bincode::Error> {
255        deserializer.deserialize_part()
256    }
257}