Implement get-size2 for LhsValue

Author: utkarshgupta137

Cargo.lock

@@ -24,6 +24,7 @@ criterion = "0.8.2"
 dyn-clone = "1.0.20"
 erased-serde = "0.4.9"
 fnv = "1.0.7"
+get-size2 = "0.10.3"
 getrandom = "0.3.4"
 indoc = "2.0.7"
 js-sys = "0.3.85"

Cargo.toml

@@ -26,6 +26,7 @@ cidr.workspace = true
 dyn-clone.workspace = true
 erased-serde.workspace = true
 fnv.workspace = true
+get-size2 = { workspace = true, optional = true }
 memchr.workspace = true
 rand.workspace = true
 regex-automata = { workspace = true, optional = true }
@@ -46,5 +47,6 @@ serde_json.workspace = true
 getrandom.workspace = true
 
 [features]
-default = ["regex"]
+default = ["get-size2", "regex"]
+get-size2 = ["dep:get-size2"]
 regex = ["dep:regex-automata"]

engine/Cargo.toml

@@ -4,6 +4,8 @@ use crate::lhs_types::AsRefIterator;
 use crate::types::{
     CompoundType, GetType, IntoValue, LhsValue, LhsValueSeed, Type, TypeMismatchError,
 };
+#[cfg(feature = "get-size2")]
+use get_size2::{GetSize, GetSizeTracker};
 use serde::de::{self, DeserializeSeed, Deserializer, SeqAccess, Visitor};
 use serde::ser::SerializeSeq;
 use serde::{Serialize, Serializer};
@@ -20,6 +22,16 @@ pub(crate) enum InnerArray<'a> {
     Borrowed(&'a [LhsValue<'a>]),
 }
 
+#[cfg(feature = "get-size2")]
+impl GetSize for InnerArray<'_> {
+    fn get_heap_size_with_tracker<T: GetSizeTracker>(&self, tracker: T) -> (usize, T) {
+        match self {
+            Self::Owned(values) => values.get_heap_size_with_tracker(tracker),
+            Self::Borrowed(_) => (0, tracker),
+        }
+    }
+}
+
 impl<'a> InnerArray<'a> {
     #[inline]
     const fn new() -> Self {
@@ -81,6 +93,13 @@ pub struct Array<'a> {
     pub(crate) data: InnerArray<'a>,
 }
 
+#[cfg(feature = "get-size2")]
+impl GetSize for Array<'_> {
+    fn get_heap_size_with_tracker<T: GetSizeTracker>(&self, tracker: T) -> (usize, T) {
+        self.data.get_heap_size_with_tracker(tracker)
+    }
+}
+
 impl<'a> Array<'a> {
     /// Creates a new array
     pub fn new(val_type: impl Into<CompoundType>) -> Self {

engine/src/lhs_types/array.rs

@@ -1,3 +1,5 @@
+#[cfg(feature = "get-size2")]
+use get_size2::{GetSize, GetSizeTracker};
 use serde::de::Visitor;
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
 use std::borrow::{Borrow, Cow};
@@ -13,6 +15,16 @@ pub enum Bytes<'a> {
     Owned(Box<[u8]>),
 }
 
+#[cfg(feature = "get-size2")]
+impl GetSize for Bytes<'_> {
+    fn get_heap_size_with_tracker<T: GetSizeTracker>(&self, tracker: T) -> (usize, T) {
+        match self {
+            Self::Owned(bytes) => bytes.get_heap_size_with_tracker(tracker),
+            Self::Borrowed(_) => (0, tracker),
+        }
+    }
+}
+
 impl<'a> Bytes<'a> {
     /// Clones self into a fully owned byte string.
     #[inline]

engine/src/lhs_types/bytes.rs

@@ -3,6 +3,8 @@ use super::array::InnerArray;
 use crate::TypeMismatchError;
 use crate::lhs_types::{AsRefIterator, Bytes};
 use crate::types::{CompoundType, GetType, IntoValue, LhsValue, LhsValueSeed, Type};
+#[cfg(feature = "get-size2")]
+use get_size2::{GetSize, GetSizeTracker};
 use serde::de::{self, DeserializeSeed, Deserializer, MapAccess, SeqAccess, Visitor};
 use serde::ser::{SerializeMap, SerializeSeq};
 use serde::{Serialize, Serializer};
@@ -17,6 +19,16 @@ pub(crate) enum InnerMap<'a> {
     Borrowed(&'a BTreeMap<Box<[u8]>, LhsValue<'a>>),
 }
 
+#[cfg(feature = "get-size2")]
+impl GetSize for InnerMap<'_> {
+    fn get_heap_size_with_tracker<T: GetSizeTracker>(&self, tracker: T) -> (usize, T) {
+        match self {
+            Self::Owned(values) => values.get_heap_size_with_tracker(tracker),
+            Self::Borrowed(_) => (0, tracker),
+        }
+    }
+}
+
 impl<'a> InnerMap<'a> {
     #[inline]
     const fn new() -> Self {
@@ -79,6 +91,13 @@ pub struct Map<'a> {
     pub(crate) data: InnerMap<'a>,
 }
 
+#[cfg(feature = "get-size2")]
+impl GetSize for Map<'_> {
+    fn get_heap_size_with_tracker<T: GetSizeTracker>(&self, tracker: T) -> (usize, T) {
+        self.data.get_heap_size_with_tracker(tracker)
+    }
+}
+
 impl<'a> Map<'a> {
     /// Creates a new map
     pub fn new(val_type: impl Into<CompoundType>) -> Self {

engine/src/lhs_types/map.rs

@@ -5,6 +5,8 @@ use crate::rhs_types::{
 };
 use crate::scheme::{FieldIndex, IndexAccessError};
 use crate::strict_partial_ord::StrictPartialOrd;
+#[cfg(feature = "get-size2")]
+use get_size2::{GetSize, GetSizeTracker};
 use serde::de::{DeserializeSeed, Deserializer};
 use serde::{Deserialize, Serialize, Serializer};
 use std::cmp::Ordering;
@@ -1062,6 +1064,18 @@ declare_types!(
     Map[CompoundType](#[serde(skip_deserializing)] Map<'a> | UninhabitedMap | UninhabitedMap),
 );
 
+#[cfg(feature = "get-size2")]
+impl GetSize for LhsValue<'_> {
+    fn get_heap_size_with_tracker<T: GetSizeTracker>(&self, tracker: T) -> (usize, T) {
+        match self {
+            Self::Bool(_) | Self::Int(_) | Self::Ip(_) => (0, tracker),
+            Self::Bytes(value) => value.get_heap_size_with_tracker(tracker),
+            Self::Array(value) => value.get_heap_size_with_tracker(tracker),
+            Self::Map(value) => value.get_heap_size_with_tracker(tracker),
+        }
+    }
+}
+
 #[test]
 fn test_lhs_value_deserialize() {
     use std::str::FromStr;
@@ -1176,3 +1190,94 @@ fn test_type_deserialize() {
 fn test_size_of_lhs_value() {
     assert_eq!(std::mem::size_of::<LhsValue<'_>>(), 48);
 }
+
+#[cfg(feature = "get-size2")]
+#[test]
+fn test_lhs_value_get_size() {
+    let stack_size = std::mem::size_of::<LhsValue<'static>>();
+
+    assert_eq!(LhsValue::Bool(false).get_size(), stack_size);
+    assert_eq!(LhsValue::Int(42).get_size(), stack_size);
+    assert_eq!(
+        LhsValue::Ip(IpAddr::V4(Ipv4Addr::LOCALHOST)).get_size(),
+        stack_size
+    );
+
+    let borrowed_bytes = LhsValue::Bytes(Bytes::Borrowed(b"borrowed"));
+    assert_eq!(borrowed_bytes.get_size(), stack_size);
+
+    let owned_bytes = LhsValue::Bytes(Bytes::Owned(Box::from(&b"owned"[..])));
+    assert_eq!(owned_bytes.get_size(), stack_size + b"owned".len());
+}
+
+#[cfg(feature = "get-size2")]
+#[test]
+fn test_array_lhs_value_get_size() {
+    let stack_size = std::mem::size_of::<LhsValue<'static>>();
+
+    let mut borrowed_values = Vec::with_capacity(4);
+    borrowed_values.push(LhsValue::Bytes(Bytes::Borrowed(b"borrowed")));
+    let borrowed_values_capacity = borrowed_values.capacity();
+    let owned_array_of_borrowed = Array::try_from_vec(Type::Bytes, borrowed_values).unwrap();
+    assert_eq!(
+        LhsValue::Array(owned_array_of_borrowed.as_ref()).get_size(),
+        stack_size
+    );
+    assert_eq!(
+        LhsValue::Array(owned_array_of_borrowed).get_size(),
+        stack_size + borrowed_values_capacity * stack_size
+    );
+
+    let mut owned_values = Vec::with_capacity(4);
+    owned_values.push(LhsValue::Bytes(Bytes::Owned(Box::from(&b"owned"[..]))));
+    let owned_values_capacity = owned_values.capacity();
+    let owned_array_of_owned = Array::try_from_vec(Type::Bytes, owned_values).unwrap();
+    assert_eq!(
+        LhsValue::Array(owned_array_of_owned.as_ref()).get_size(),
+        stack_size
+    );
+    assert_eq!(
+        LhsValue::Array(owned_array_of_owned).get_size(),
+        stack_size + owned_values_capacity * stack_size + b"owned".len()
+    );
+}
+
+#[cfg(feature = "get-size2")]
+#[test]
+fn test_map_lhs_value_get_size() {
+    let stack_size = std::mem::size_of::<LhsValue<'static>>();
+
+    let owned_map_of_borrowed = Map::try_from_iter(
+        Type::Bytes,
+        [Ok::<_, TypeMismatchError>((
+            Box::from(b"key".as_slice()),
+            Bytes::Borrowed(b"borrowed"),
+        ))],
+    )
+    .unwrap();
+    assert_eq!(
+        LhsValue::Map(owned_map_of_borrowed.as_ref()).get_size(),
+        stack_size
+    );
+    assert_eq!(
+        LhsValue::Map(owned_map_of_borrowed).get_size(),
+        stack_size + std::mem::size_of::<Box<[u8]>>() + b"key".len() + stack_size
+    );
+
+    let owned_map_of_owned = Map::try_from_iter(
+        Type::Bytes,
+        [Ok::<_, TypeMismatchError>((
+            Box::from(b"key".as_slice()),
+            Bytes::Owned(Box::from(&b"owned"[..])),
+        ))],
+    )
+    .unwrap();
+    assert_eq!(
+        LhsValue::Map(owned_map_of_owned.as_ref()).get_size(),
+        stack_size
+    );
+    assert_eq!(
+        LhsValue::Map(owned_map_of_owned).get_size(),
+        stack_size + std::mem::size_of::<Box<[u8]>>() + b"key".len() + stack_size + b"owned".len()
+    );
+}