relay_protocol/
size.rs

1use serde::de::value::Error;
2use serde::{ser, Serialize};
3use smallvec::SmallVec;
4
5use crate::traits::{IntoValue, SkipSerialization};
6
7/// Estimates the size in bytes this would be in JSON.
8pub fn estimate_size<T: IntoValue>(value: Option<&T>) -> usize {
9    let mut ser = SizeEstimatingSerializer::new();
10    if let Some(value) = value {
11        IntoValue::serialize_payload(value, &mut ser, SkipSerialization::default()).unwrap();
12    }
13    ser.size()
14}
15
16/// Estimates the size in bytes this would be in JSON, but does not recurse into objects or arrays.
17pub fn estimate_size_flat<T: IntoValue>(value: Option<&T>) -> usize {
18    let mut ser = SizeEstimatingSerializer::new();
19    ser.flat = true;
20    if let Some(value) = value {
21        IntoValue::serialize_payload(value, &mut ser, SkipSerialization::default()).unwrap();
22    }
23    ser.size()
24}
25
26/// Helper serializer that efficiently determines how much space something might take.
27///
28/// This counts in estimated bytes.
29#[derive(Default)]
30struct SizeEstimatingSerializer {
31    size: usize,
32    item_stack: SmallVec<[bool; 16]>,
33    flat: bool,
34}
35
36impl SizeEstimatingSerializer {
37    /// Creates a new serializer
38    pub fn new() -> Self {
39        Self::default()
40    }
41
42    /// Returns the calculated size
43    pub fn size(&self) -> usize {
44        self.size
45    }
46
47    #[inline]
48    fn count_size(&mut self, incr: usize) {
49        if self.flat && !self.item_stack.is_empty() {
50            return;
51        }
52
53        self.size += incr;
54    }
55
56    fn push(&mut self) {
57        self.item_stack.push(false);
58    }
59
60    fn pop(&mut self) {
61        self.item_stack.pop();
62    }
63
64    fn count_comma_sep(&mut self) {
65        if let Some(state) = self.item_stack.last_mut() {
66            if !*state {
67                *state = true;
68            } else {
69                self.count_size(1);
70            }
71        }
72    }
73}
74
75impl ser::Serializer for &mut SizeEstimatingSerializer {
76    type Ok = ();
77    type Error = Error;
78
79    type SerializeSeq = Self;
80    type SerializeTuple = Self;
81    type SerializeTupleStruct = Self;
82    type SerializeTupleVariant = Self;
83    type SerializeMap = Self;
84    type SerializeStruct = Self;
85    type SerializeStructVariant = Self;
86
87    #[inline]
88    fn serialize_bool(self, v: bool) -> Result<(), Error> {
89        self.count_size(if v { 4 } else { 5 });
90        Ok(())
91    }
92
93    #[inline(always)]
94    fn serialize_i8(self, v: i8) -> Result<(), Error> {
95        self.serialize_i64(i64::from(v))
96    }
97
98    #[inline(always)]
99    fn serialize_i16(self, v: i16) -> Result<(), Error> {
100        self.serialize_i64(i64::from(v))
101    }
102
103    #[inline(always)]
104    fn serialize_i32(self, v: i32) -> Result<(), Error> {
105        self.serialize_i64(i64::from(v))
106    }
107
108    #[inline]
109    fn serialize_i64(self, v: i64) -> Result<(), Error> {
110        self.size += &v.to_string().len();
111        Ok(())
112    }
113
114    #[inline(always)]
115    fn serialize_u8(self, v: u8) -> Result<(), Error> {
116        self.serialize_u64(u64::from(v))
117    }
118
119    #[inline(always)]
120    fn serialize_u16(self, v: u16) -> Result<(), Error> {
121        self.serialize_u64(u64::from(v))
122    }
123
124    #[inline(always)]
125    fn serialize_u32(self, v: u32) -> Result<(), Error> {
126        self.serialize_u64(u64::from(v))
127    }
128
129    #[inline]
130    fn serialize_u64(self, v: u64) -> Result<(), Error> {
131        self.count_size(v.to_string().len());
132        Ok(())
133    }
134
135    #[inline(always)]
136    fn serialize_f32(self, v: f32) -> Result<(), Error> {
137        self.serialize_f64(f64::from(v))
138    }
139
140    #[inline]
141    fn serialize_f64(self, v: f64) -> Result<(), Error> {
142        self.count_size(v.to_string().len());
143        Ok(())
144    }
145
146    #[inline(always)]
147    fn serialize_char(self, _v: char) -> Result<(), Error> {
148        self.count_size(1);
149        Ok(())
150    }
151
152    #[inline(always)]
153    fn serialize_str(self, v: &str) -> Result<(), Error> {
154        self.count_size(v.len() + 2);
155        Ok(())
156    }
157
158    fn serialize_bytes(self, v: &[u8]) -> Result<(), Error> {
159        use serde::ser::SerializeSeq;
160        let mut seq = self.serialize_seq(Some(v.len()))?;
161        for byte in v {
162            seq.serialize_element(byte)?;
163        }
164        seq.end()
165    }
166
167    #[inline(always)]
168    fn serialize_none(self) -> Result<(), Error> {
169        self.serialize_unit()
170    }
171
172    #[inline(always)]
173    fn serialize_some<T>(self, value: &T) -> Result<(), Error>
174    where
175        T: ?Sized + Serialize,
176    {
177        value.serialize(self)
178    }
179
180    #[inline(always)]
181    fn serialize_unit(self) -> Result<(), Error> {
182        self.count_size(4);
183        Ok(())
184    }
185
186    #[inline(always)]
187    fn serialize_unit_struct(self, _name: &'static str) -> Result<(), Error> {
188        self.serialize_unit()
189    }
190
191    #[inline(always)]
192    fn serialize_unit_variant(
193        self,
194        _name: &'static str,
195        _variant_index: u32,
196        variant: &'static str,
197    ) -> Result<(), Error> {
198        self.serialize_str(variant)
199    }
200
201    #[inline(always)]
202    fn serialize_newtype_struct<T>(self, _name: &'static str, value: &T) -> Result<(), Error>
203    where
204        T: ?Sized + Serialize,
205    {
206        value.serialize(self)
207    }
208
209    fn serialize_newtype_variant<T>(
210        self,
211        _name: &'static str,
212        _variant_index: u32,
213        variant: &'static str,
214        value: &T,
215    ) -> Result<(), Error>
216    where
217        T: ?Sized + Serialize,
218    {
219        // { x : y }
220        self.count_size(3);
221        variant.serialize(&mut *self)?;
222        value.serialize(&mut *self)?;
223        Ok(())
224    }
225
226    fn serialize_seq(self, _len: Option<usize>) -> Result<Self::SerializeSeq, Error> {
227        self.count_size(1);
228        self.push();
229        Ok(self)
230    }
231
232    #[inline(always)]
233    fn serialize_tuple(self, len: usize) -> Result<Self::SerializeTuple, Error> {
234        self.serialize_seq(Some(len))
235    }
236
237    #[inline(always)]
238    fn serialize_tuple_struct(
239        self,
240        _name: &'static str,
241        len: usize,
242    ) -> Result<Self::SerializeTupleStruct, Error> {
243        self.serialize_seq(Some(len))
244    }
245
246    fn serialize_tuple_variant(
247        self,
248        _name: &'static str,
249        _variant_index: u32,
250        variant: &'static str,
251        _len: usize,
252    ) -> Result<Self::SerializeTupleVariant, Error> {
253        // { x: [
254        self.count_size(3);
255        variant.serialize(&mut *self)?;
256        Ok(self)
257    }
258
259    fn serialize_map(self, _len: Option<usize>) -> Result<Self::SerializeMap, Error> {
260        // {
261        self.count_size(1);
262        self.push();
263        Ok(self)
264    }
265
266    #[inline(always)]
267    fn serialize_struct(
268        self,
269        _name: &'static str,
270        len: usize,
271    ) -> Result<Self::SerializeStruct, Error> {
272        self.serialize_map(Some(len))
273    }
274
275    fn serialize_struct_variant(
276        self,
277        _name: &'static str,
278        _variant_index: u32,
279        variant: &'static str,
280        _len: usize,
281    ) -> Result<Self::SerializeStructVariant, Error> {
282        // { x: {
283        self.count_size(3);
284        variant.serialize(&mut *self)?;
285        Ok(self)
286    }
287}
288
289impl ser::SerializeSeq for &mut SizeEstimatingSerializer {
290    type Ok = ();
291    type Error = Error;
292
293    fn serialize_element<T>(&mut self, value: &T) -> Result<(), Error>
294    where
295        T: ?Sized + Serialize,
296    {
297        self.count_comma_sep();
298        value.serialize(&mut **self)
299    }
300
301    fn end(self) -> Result<(), Error> {
302        self.pop();
303        self.count_size(1);
304        Ok(())
305    }
306}
307
308// Same thing but for tuples.
309impl ser::SerializeTuple for &mut SizeEstimatingSerializer {
310    type Ok = ();
311    type Error = Error;
312
313    fn serialize_element<T>(&mut self, value: &T) -> Result<(), Error>
314    where
315        T: ?Sized + Serialize,
316    {
317        self.count_comma_sep();
318        value.serialize(&mut **self)
319    }
320
321    #[inline(always)]
322    fn end(self) -> Result<(), Error> {
323        self.count_size(1);
324        Ok(())
325    }
326}
327
328// Same thing but for tuple structs.
329impl ser::SerializeTupleStruct for &mut SizeEstimatingSerializer {
330    type Ok = ();
331    type Error = Error;
332
333    fn serialize_field<T>(&mut self, value: &T) -> Result<(), Error>
334    where
335        T: ?Sized + Serialize,
336    {
337        self.count_comma_sep();
338        value.serialize(&mut **self)
339    }
340
341    #[inline(always)]
342    fn end(self) -> Result<(), Error> {
343        self.count_size(1);
344        Ok(())
345    }
346}
347
348impl ser::SerializeTupleVariant for &mut SizeEstimatingSerializer {
349    type Ok = ();
350    type Error = Error;
351
352    fn serialize_field<T>(&mut self, value: &T) -> Result<(), Error>
353    where
354        T: ?Sized + Serialize,
355    {
356        self.count_comma_sep();
357        value.serialize(&mut **self)
358    }
359
360    #[inline(always)]
361    fn end(self) -> Result<(), Error> {
362        self.count_size(2);
363        Ok(())
364    }
365}
366
367impl ser::SerializeMap for &mut SizeEstimatingSerializer {
368    type Ok = ();
369    type Error = Error;
370
371    fn serialize_key<T>(&mut self, key: &T) -> Result<(), Error>
372    where
373        T: ?Sized + Serialize,
374    {
375        self.count_comma_sep();
376        key.serialize(&mut **self)
377    }
378
379    fn serialize_value<T>(&mut self, value: &T) -> Result<(), Error>
380    where
381        T: ?Sized + Serialize,
382    {
383        self.count_size(1);
384        value.serialize(&mut **self)
385    }
386
387    #[inline(always)]
388    fn end(self) -> Result<(), Error> {
389        self.pop();
390        self.count_size(1);
391        Ok(())
392    }
393}
394
395impl ser::SerializeStruct for &mut SizeEstimatingSerializer {
396    type Ok = ();
397    type Error = Error;
398
399    fn serialize_field<T>(&mut self, key: &'static str, value: &T) -> Result<(), Error>
400    where
401        T: ?Sized + Serialize,
402    {
403        self.count_comma_sep();
404        self.count_size(2);
405        key.serialize(&mut **self)?;
406        value.serialize(&mut **self)
407    }
408
409    #[inline(always)]
410    fn end(self) -> Result<(), Error> {
411        self.count_size(1);
412        Ok(())
413    }
414}
415
416impl ser::SerializeStructVariant for &mut SizeEstimatingSerializer {
417    type Ok = ();
418    type Error = Error;
419
420    fn serialize_field<T>(&mut self, key: &'static str, value: &T) -> Result<(), Error>
421    where
422        T: ?Sized + Serialize,
423    {
424        self.count_comma_sep();
425        self.size += 2;
426        key.serialize(&mut **self)?;
427        value.serialize(&mut **self)
428    }
429
430    #[inline(always)]
431    fn end(self) -> Result<(), Error> {
432        self.count_size(2);
433        Ok(())
434    }
435}
436
437#[cfg(test)]
438mod tests {
439    use super::*;
440
441    use crate::annotated::Annotated;
442    use crate::value::{Object, Value};
443
444    #[test]
445    fn test_estimate_size() {
446        let json = r#"{"a":["Hello","World","aha","hmm",false,{"blub":42,"x":true},null]}"#;
447        let value = Annotated::<Object<Value>>::from_json(json).unwrap();
448        assert_eq!(estimate_size(value.value()), json.len());
449    }
450}