vector_common/
byte_size_of.rs

1use std::{
2    collections::{BTreeMap, BTreeSet},
3    mem,
4};
5
6use bytes::{Bytes, BytesMut};
7use chrono::{DateTime, Utc};
8use serde_json::{value::RawValue, Value as JsonValue};
9use smallvec::SmallVec;
10use vrl::value::{KeyString, Value};
11
12pub trait ByteSizeOf {
13    /// Returns the in-memory size of this type
14    ///
15    /// This function returns the total number of bytes that
16    /// [`std::mem::size_of`] does in addition to any interior allocated
17    /// bytes. Its default implementation is `std::mem::size_of` +
18    /// `ByteSizeOf::allocated_bytes`.
19    fn size_of(&self) -> usize {
20        mem::size_of_val(self) + self.allocated_bytes()
21    }
22
23    /// Returns the allocated bytes of this type
24    ///
25    /// This function returns the total number of bytes that have been allocated
26    /// interior to this type instance. It does not include any bytes that are
27    /// captured by [`std::mem::size_of`] except for any bytes that are interior
28    /// to this type. For instance, `BTreeMap<String, Vec<u8>>` would count all
29    /// bytes for `String` and `Vec<u8>` instances but not the exterior bytes
30    /// for `BTreeMap`.
31    fn allocated_bytes(&self) -> usize;
32}
33
34impl<T> ByteSizeOf for &T
35where
36    T: ByteSizeOf,
37{
38    fn allocated_bytes(&self) -> usize {
39        (*self).size_of()
40    }
41}
42
43impl ByteSizeOf for Bytes {
44    fn allocated_bytes(&self) -> usize {
45        self.len()
46    }
47}
48
49impl ByteSizeOf for BytesMut {
50    fn allocated_bytes(&self) -> usize {
51        self.len()
52    }
53}
54
55impl ByteSizeOf for String {
56    fn allocated_bytes(&self) -> usize {
57        self.len()
58    }
59}
60
61impl ByteSizeOf for KeyString {
62    fn allocated_bytes(&self) -> usize {
63        self.len()
64    }
65}
66
67impl ByteSizeOf for &str {
68    fn allocated_bytes(&self) -> usize {
69        0
70    }
71}
72
73impl<K, V> ByteSizeOf for BTreeMap<K, V>
74where
75    K: ByteSizeOf,
76    V: ByteSizeOf,
77{
78    fn allocated_bytes(&self) -> usize {
79        self.iter()
80            .fold(0, |acc, (k, v)| acc + k.size_of() + v.size_of())
81    }
82}
83
84impl<T> ByteSizeOf for BTreeSet<T>
85where
86    T: ByteSizeOf,
87{
88    fn allocated_bytes(&self) -> usize {
89        self.iter().map(ByteSizeOf::size_of).sum()
90    }
91}
92
93impl<T> ByteSizeOf for Vec<T>
94where
95    T: ByteSizeOf,
96{
97    fn allocated_bytes(&self) -> usize {
98        self.iter().map(ByteSizeOf::size_of).sum()
99    }
100}
101
102impl<A: smallvec::Array> ByteSizeOf for SmallVec<A>
103where
104    A::Item: ByteSizeOf,
105{
106    fn allocated_bytes(&self) -> usize {
107        self.iter().map(ByteSizeOf::size_of).sum()
108    }
109}
110
111impl<T> ByteSizeOf for &[T]
112where
113    T: ByteSizeOf,
114{
115    fn allocated_bytes(&self) -> usize {
116        self.iter().map(ByteSizeOf::size_of).sum()
117    }
118}
119
120impl<T, const N: usize> ByteSizeOf for [T; N]
121where
122    T: ByteSizeOf,
123{
124    fn size_of(&self) -> usize {
125        self.allocated_bytes()
126    }
127
128    fn allocated_bytes(&self) -> usize {
129        self.iter().map(ByteSizeOf::size_of).sum()
130    }
131}
132
133impl<T> ByteSizeOf for Option<T>
134where
135    T: ByteSizeOf,
136{
137    fn allocated_bytes(&self) -> usize {
138        self.as_ref().map_or(0, ByteSizeOf::allocated_bytes)
139    }
140}
141
142macro_rules! num {
143    ($t:ty) => {
144        impl ByteSizeOf for $t {
145            fn allocated_bytes(&self) -> usize {
146                0
147            }
148        }
149    };
150}
151
152num!(u8);
153num!(u16);
154num!(u32);
155num!(u64);
156num!(u128);
157num!(i8);
158num!(i16);
159num!(i32);
160num!(i64);
161num!(i128);
162num!(f32);
163num!(f64);
164
165impl ByteSizeOf for Box<RawValue> {
166    fn allocated_bytes(&self) -> usize {
167        self.get().len()
168    }
169}
170
171impl ByteSizeOf for JsonValue {
172    fn allocated_bytes(&self) -> usize {
173        match self {
174            JsonValue::Null | JsonValue::Bool(_) | JsonValue::Number(_) => 0,
175            JsonValue::String(s) => s.len(),
176            JsonValue::Array(a) => a.size_of(),
177            JsonValue::Object(o) => o.iter().map(|(k, v)| k.size_of() + v.size_of()).sum(),
178        }
179    }
180}
181
182impl ByteSizeOf for Value {
183    fn allocated_bytes(&self) -> usize {
184        match self {
185            Value::Bytes(bytes) => bytes.len(),
186            Value::Object(map) => map.size_of(),
187            Value::Array(arr) => arr.size_of(),
188            _ => 0,
189        }
190    }
191}
192
193impl ByteSizeOf for DateTime<Utc> {
194    fn allocated_bytes(&self) -> usize {
195        0
196    }
197}
198
199impl<K, V> ByteSizeOf for indexmap::IndexMap<K, V>
200where
201    K: ByteSizeOf,
202    V: ByteSizeOf,
203{
204    fn allocated_bytes(&self) -> usize {
205        self.iter()
206            .fold(0, |acc, (k, v)| acc + k.size_of() + v.size_of())
207    }
208}
209
210impl<T> ByteSizeOf for indexmap::IndexSet<T>
211where
212    T: ByteSizeOf,
213{
214    fn allocated_bytes(&self) -> usize {
215        self.iter().map(ByteSizeOf::size_of).sum()
216    }
217}