diff --git a/src/lib.rs b/src/lib.rs index 43551a9..59a8f08 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,6 +4,8 @@ //! Data structure measurement. +#![cfg_attr(feature = "unstable", feature(rc_counts, arc_counts))] + #[cfg(target_os = "windows")] extern crate kernel32; @@ -16,7 +18,7 @@ use std::collections::{BTreeMap, HashMap, LinkedList, VecDeque}; use std::hash::BuildHasher; use std::hash::Hash; use std::marker::PhantomData; -use std::mem::size_of; +use std::mem::{size_of, size_of_val}; use std::net::{Ipv4Addr, Ipv6Addr}; use std::os::raw::c_void; use std::sync::Arc; @@ -45,7 +47,7 @@ unsafe fn heap_size_of_impl(ptr: *const c_void) -> usize { // this function doesn't modify the contents of the block that `ptr` points to, so we use // `*const c_void` here. extern "C" { - #[cfg_attr(any(prefixed_jemalloc, target_os = "macos", target_os = "android"), link_name = "je_malloc_usable_size")] + #[cfg_attr(any(prefixed_jemalloc, target_os = "macos", target_os = "android"), link_name = "je_malloc_usable_size")] fn malloc_usable_size(ptr: *const c_void) -> usize; } malloc_usable_size(ptr) @@ -150,12 +152,6 @@ impl HeapSizeOf for (T, U) { } } -impl HeapSizeOf for Arc { - fn heap_size_of_children(&self) -> usize { - (**self).heap_size_of_children() - } -} - impl HeapSizeOf for RefCell { fn heap_size_of_children(&self) -> usize { self.borrow().heap_size_of_children() @@ -185,16 +181,44 @@ impl HeapSizeOf for VecDeque { } } -impl HeapSizeOf for Vec> { +impl HeapSizeOf for Rc { + #[cfg(not(feature = "unstable"))] fn heap_size_of_children(&self) -> usize { - // The fate of measuring Rc is still undecided, but we still want to measure - // the space used for storing them. - unsafe { - heap_size_of(self.as_ptr() as *const c_void) + 0 + } + + #[cfg(feature = "unstable")] + fn heap_size_of_children(&self) -> usize { + if Rc::strong_count(self) == 1 { + heap_size_of_rcbox::(self) + } else { + 0 } } } +impl HeapSizeOf for Arc { + #[cfg(not(feature = "unstable"))] + fn heap_size_of_children(&self) -> usize { + 0 + } + + #[cfg(feature = "unstable")] + fn heap_size_of_children(&self) -> usize { + if Arc::strong_count(self) == 1 { + heap_size_of_rcbox::(self) + } else { + 0 + } + } +} + +/// Arc and Rc heap-allocate `(usize, usize, T)` with strong and weak refcounts. +#[cfg(feature = "unstable")] +fn heap_size_of_rcbox(x: &T) -> usize { + size_of::() * 2 + size_of_val(x) + x.heap_size_of_children() +} + #[cfg(feature = "unstable")] impl HeapSizeOf for HashMap where K: Eq + Hash, S: BuildHasher { diff --git a/tests/tests.rs b/tests/tests.rs index 67ca98d..60b603f 100644 --- a/tests/tests.rs +++ b/tests/tests.rs @@ -7,6 +7,7 @@ extern crate heapsize; use heapsize::{HeapSizeOf, heap_size_of}; +use std::mem::size_of; use std::os::raw::c_void; pub const EMPTY: *mut () = 0x1 as *mut (); @@ -141,13 +142,14 @@ fn test_heap_size() { let x = Some(Box::new(0i64)); assert_eq!(x.heap_size_of_children(), 8); - // Not on the heap. let x = ::std::sync::Arc::new(0i64); - assert_eq!(x.heap_size_of_children(), 0); + assert_eq!(x.heap_size_of_children(), + if cfg!(feature = "unstable") { size_of::() * 2 + 8 } else { 0 }); - // The `Arc` is not on the heap, but the Box is. + // `ArcInner>` has 2 refcounts + a pointer on the heap. let x = ::std::sync::Arc::new(Box::new(0i64)); - assert_eq!(x.heap_size_of_children(), 8); + assert_eq!(x.heap_size_of_children(), + if cfg!(feature = "unstable") { size_of::() * 3 + 8 } else { 0 }); // Zero elements, no heap storage. let x: Vec = vec![];