diff --git a/malloc_size_of/src/impls.rs b/malloc_size_of/src/impls.rs index 2898455..691b993 100644 --- a/malloc_size_of/src/impls.rs +++ b/malloc_size_of/src/impls.rs @@ -23,7 +23,7 @@ use core::sync::atomic::{AtomicU16, AtomicU32, AtomicU64, AtomicU8, AtomicUsize} use alloc::borrow::{Cow, ToOwned}; use alloc::boxed::Box; -use alloc::collections::{BTreeMap, VecDeque}; +use alloc::collections::{BTreeMap, BTreeSet, VecDeque}; use alloc::string::String; use alloc::vec::Vec; @@ -318,6 +318,34 @@ where } } +impl MallocShallowSizeOf for BTreeSet +where + T: Ord, +{ + fn shallow_size_of(&self, ops: &mut MallocSizeOfOps) -> usize { + if ops.has_malloc_enclosing_size_of() { + self.iter() + .next() + .map_or(0, |v| unsafe { ops.malloc_enclosing_size_of(v) }) + } else { + self.len() * (size_of::() + size_of::()) + } + } +} + +impl MallocSizeOf for BTreeSet +where + T: Ord + MallocSizeOf, +{ + fn size_of(&self, ops: &mut MallocSizeOfOps) -> usize { + let mut n = self.shallow_size_of(ops); + for v in self.iter() { + n += v.size_of(ops); + } + n + } +} + #[cfg(feature = "std")] impl MallocShallowSizeOf for HashSet where