Properly support Windows

We use HeapSize with GetProcessHeap() and flags 0, like Rust itself.

Author: nox

Cargo.toml

@@ -6,5 +6,8 @@ description = "Infrastructure for measuring the total runtime size of an object
 license = "MPL-2.0"
 repository = "https://github.com/servo/heapsize"
 
+[dependencies]
+kernel32-sys = "0.2.1"
+
 [features]
 unstable = []

src/lib.rs

@@ -4,6 +4,11 @@
 
 //! Data structure measurement.
 
+#[cfg(target_os = "windows")]
+extern crate kernel32;
+
+#[cfg(target_os = "windows")]
+use kernel32::{GetProcessHeap, HeapSize, HeapValidate};
 use std::borrow::Cow;
 use std::cell::{Cell, RefCell};
 use std::collections::{BTreeMap, HashMap, LinkedList, VecDeque};
@@ -18,38 +23,42 @@ use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, AtomicIsize, AtomicUsize};
 use std::rc::Rc;
 
-#[cfg(not(target_os = "windows"))]
-extern {
-    // Get the size of a heap block.
-    //
-    // Ideally Rust would expose a function like this in std::rt::heap, which would avoid the
-    // jemalloc dependence.
-    //
-    // The C prototype is `je_malloc_usable_size(JEMALLOC_USABLE_SIZE_CONST void *ptr)`. On some
-    // platforms `JEMALLOC_USABLE_SIZE_CONST` is `const` and on some it is empty. But in practice
-    // this function doesn't modify the contents of the block that `ptr` points to, so we use
-    // `*const c_void` here.
-    fn je_malloc_usable_size(ptr: *const c_void) -> usize;
-}
-
-/// A wrapper for je_malloc_usable_size that handles `EMPTY` and returns `usize`.
+/// Get the size of a heap block.
+///
+/// Ideally Rust would expose a function like this in std::rt::heap.
 ///
 /// `unsafe` because the caller must ensure that the pointer is from jemalloc.
 /// FIXME: This probably interacts badly with custom allocators:
 /// https://doc.rust-lang.org/book/custom-allocators.html
-#[cfg(not(target_os = "windows"))]
 pub unsafe fn heap_size_of(ptr: *const c_void) -> usize {
     if ptr == 0x01 as *const c_void {
         0
     } else {
-        je_malloc_usable_size(ptr)
+        heap_size_of_impl(ptr)
     }
 }
 
-/// FIXME: Need to implement heap size support on Windows.
+#[cfg(not(target_os = "windows"))]
+unsafe fn heap_size_of_impl(ptr: *const c_void) -> usize {
+    // The C prototype is `je_malloc_usable_size(JEMALLOC_USABLE_SIZE_CONST void *ptr)`. On some
+    // platforms `JEMALLOC_USABLE_SIZE_CONST` is `const` and on some it is empty. But in practice
+    // this function doesn't modify the contents of the block that `ptr` points to, so we use
+    // `*const c_void` here.
+    extern "C" {
+        fn je_malloc_usable_size(ptr: *const c_void) -> usize;
+    }
+    je_malloc_usable_size(ptr)
+}
+
 #[cfg(target_os = "windows")]
-pub unsafe fn heap_size_of(ptr: *const c_void) -> usize {
-    0
+pub unsafe fn heap_size_of_impl(mut ptr: *const c_void) -> usize {
+    let heap = GetProcessHeap();
+
+    if HeapValidate(heap, 0, ptr) == 0 {
+        ptr = *(ptr as *const *const c_void).offset(-1);
+    }
+
+    HeapSize(heap, 0, ptr) as usize
 }
 
 // The simplest trait for measuring the size of heap data structures. More complex traits that

tests/tests.rs

@@ -2,7 +2,7 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this
  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
 
-#![cfg_attr(feature= "unstable", feature(alloc, heap_api))]
+#![cfg_attr(feature= "unstable", feature(alloc, heap_api, repr_simd))]
 
 extern crate heapsize;
 
@@ -18,11 +18,15 @@ mod unstable {
     use heapsize::heap_size_of;
     use std::os::raw::c_void;
 
+    #[repr(C, simd)]
+    struct OverAligned(u64, u64, u64, u64);
+
     #[test]
     fn check_empty() {
         assert_eq!(::EMPTY, alloc::heap::EMPTY);
     }
 
+    #[cfg(not(target_os = "windows"))]
     #[test]
     fn test_alloc() {
         unsafe {
@@ -40,8 +44,53 @@ mod unstable {
             let x = alloc::heap::allocate(1024 * 1024, 0);
             assert_eq!(heap_size_of(x as *const c_void), 1024 * 1024);
             alloc::heap::deallocate(x, 1024 * 1024, 0);
+
+            // An overaligned 1MiB request is allocated exactly.
+            let x = alloc::heap::allocate(1024 * 1024, 32);
+            assert_eq!(heap_size_of(x as *const c_void), 1024 * 1024);
+            alloc::heap::deallocate(x, 1024 * 1024, 32);
+        }
+    }
+
+    #[cfg(target_os = "windows")]
+    #[test]
+    fn test_alloc() {
+        unsafe {
+            // A 64 byte request is allocated exactly.
+            let x = alloc::heap::allocate(64, 0);
+            assert_eq!(heap_size_of(x as *const c_void), 64);
+            alloc::heap::deallocate(x, 64, 0);
+
+            // A 255 byte request is allocated exactly.
+            let x = alloc::heap::allocate(255, 0);
+            assert_eq!(heap_size_of(x as *const c_void), 255);
+            alloc::heap::deallocate(x, 255, 0);
+
+            // A 1MiB request is allocated exactly.
+            let x = alloc::heap::allocate(1024 * 1024, 0);
+            assert_eq!(heap_size_of(x as *const c_void), 1024 * 1024);
+            alloc::heap::deallocate(x, 1024 * 1024, 0);
+
+            // An overaligned 1MiB request is over-allocated.
+            let x = alloc::heap::allocate(1024 * 1024, 32);
+            assert_eq!(heap_size_of(x as *const c_void), 1024 * 1024 + 32);
+            alloc::heap::deallocate(x, 1024 * 1024, 32);
         }
     }
+
+    #[cfg(not(target_os = "windows"))]
+    #[test]
+    fn test_simd() {
+        let x = Box::new(OverAligned(0, 0, 0, 0));
+        assert_eq!(unsafe { heap_size_of(&*x as *const _ as *const c_void) }, 32);
+    }
+
+    #[cfg(target_os = "windows")]
+    #[test]
+    fn test_simd() {
+        let x = Box::new(OverAligned(0, 0, 0, 0));
+        assert_eq!(unsafe { heap_size_of(&*x as *const _ as *const c_void) }, 32 + 32);
+    }
 }
 
 #[test]