Merge pull request #1521 from fitzgen/anyref-heap-live-count

Utility for counting the number of live JsValues in the wasm-bindgen heap

Author: fitzgen

crates/cli-support/src/js/mod.rs

@@ -636,6 +636,37 @@ impl<'a> Context<'a> {
             ))
         })?;
 
+        self.bind("__wbindgen_anyref_heap_live_count", &|me| {
+            if me.config.anyref {
+                // Eventually we should add support to the anyref-xform to
+                // re-write calls to the imported
+                // `__wbindgen_anyref_heap_live_count` function into calls to
+                // the exported `__wbindgen_anyref_heap_live_count_impl`
+                // function, and to un-export that function.
+                //
+                // But for now, we just bounce wasm -> js -> wasm because it is
+                // easy.
+                Ok("function() {{ return wasm.__wbindgen_anyref_heap_live_count_impl(); }}".into())
+            } else {
+                me.expose_global_heap();
+                Ok(format!(
+                    "
+                    function() {{
+                        let free_count = 0;
+                        let next = heap_next;
+                        while (next < heap.length) {{
+                            free_count += 1;
+                            next = heap[next];
+                        }}
+                        return heap.length - free_count - {} - {};
+                    }}
+                    ",
+                    INITIAL_HEAP_OFFSET,
+                    INITIAL_HEAP_VALUES.len(),
+                ))
+            }
+        })?;
+
         self.bind("__wbindgen_debug_string", &|me| {
             me.expose_pass_string_to_wasm()?;
             me.expose_uint32_memory();

src/anyref.rs

@@ -1,8 +1,8 @@
-use std::slice;
-use std::vec::Vec;
-use std::ptr;
 use std::alloc::{self, Layout};
 use std::mem;
+use std::ptr;
+use std::slice;
+use std::vec::Vec;
 
 use crate::JsValue;
 
@@ -34,9 +34,7 @@ impl Slab {
         if ret == self.data.len() {
             if self.data.len() == self.data.capacity() {
                 let extra = 128;
-                let r = unsafe {
-                    __wbindgen_anyref_table_grow(extra)
-                };
+                let r = unsafe { __wbindgen_anyref_table_grow(extra) };
                 if r == -1 {
                     internal_error("table grow failure")
                 }
@@ -59,16 +57,8 @@ impl Slab {
                     if ptr.is_null() {
                         internal_error("allocation failure");
                     }
-                    ptr::copy_nonoverlapping(
-                        self.data.as_ptr(),
-                        ptr,
-                        self.data.len(),
-                    );
-                    let new_vec = Vec::from_raw_parts(
-                        ptr,
-                        self.data.len(),
-                        new_cap,
-                    );
+                    ptr::copy_nonoverlapping(self.data.as_ptr(), ptr, self.data.len());
+                    let new_vec = Vec::from_raw_parts(ptr, self.data.len(), new_cap);
                     let mut old = mem::replace(&mut self.data, new_vec);
                     old.set_len(0);
                 }
@@ -107,6 +97,20 @@ impl Slab {
             None => internal_error("slot out of bounds"),
         }
     }
+
+    fn live_count(&self) -> u32 {
+        let mut free_count = 0;
+        let mut next = self.head;
+        while next < self.data.len() {
+            debug_assert!((free_count as usize) < self.data.len());
+            free_count += 1;
+            match self.data.get(next) {
+                Some(n) => next = *n,
+                None => internal_error("slot out of bounds"),
+            };
+        }
+        self.data.len() as u32 - free_count - super::JSIDX_RESERVED
+    }
 }
 
 fn internal_error(msg: &str) -> ! {
@@ -135,19 +139,19 @@ fn internal_error(msg: &str) -> ! {
 // implementation that will be replaced once #55518 lands on stable.
 #[cfg(target_feature = "atomics")]
 mod tl {
-    use std::*; // hack to get `thread_local!` to work
     use super::Slab;
     use std::cell::Cell;
+    use std::*; // hack to get `thread_local!` to work
 
     thread_local!(pub static HEAP_SLAB: Cell<Slab> = Cell::new(Slab::new()));
 }
 
 #[cfg(not(target_feature = "atomics"))]
 mod tl {
+    use super::Slab;
     use std::alloc::{self, Layout};
     use std::cell::Cell;
     use std::ptr;
-    use super::Slab;
 
     pub struct HeapSlab;
     pub static HEAP_SLAB: HeapSlab = HeapSlab;
@@ -172,40 +176,57 @@ mod tl {
 }
 
 #[no_mangle]
-pub extern fn __wbindgen_anyref_table_alloc() -> usize {
-    tl::HEAP_SLAB.try_with(|slot| {
-        let mut slab = slot.replace(Slab::new());
-        let ret = slab.alloc();
-        slot.replace(slab);
-        ret
-    }).unwrap_or_else(|_| internal_error("tls access failure"))
+pub extern "C" fn __wbindgen_anyref_table_alloc() -> usize {
+    tl::HEAP_SLAB
+        .try_with(|slot| {
+            let mut slab = slot.replace(Slab::new());
+            let ret = slab.alloc();
+            slot.replace(slab);
+            ret
+        })
+        .unwrap_or_else(|_| internal_error("tls access failure"))
 }
 
 #[no_mangle]
-pub extern fn __wbindgen_anyref_table_dealloc(idx: usize) {
+pub extern "C" fn __wbindgen_anyref_table_dealloc(idx: usize) {
     if idx < super::JSIDX_RESERVED as usize {
-        return
+        return;
     }
     // clear this value from the table so while the table slot is un-allocated
     // we don't keep around a strong reference to a potentially large object
     unsafe {
         __wbindgen_anyref_table_set_null(idx);
     }
-    tl::HEAP_SLAB.try_with(|slot| {
-        let mut slab = slot.replace(Slab::new());
-        slab.dealloc(idx);
-        slot.replace(slab);
-    }).unwrap_or_else(|_| internal_error("tls access failure"))
+    tl::HEAP_SLAB
+        .try_with(|slot| {
+            let mut slab = slot.replace(Slab::new());
+            slab.dealloc(idx);
+            slot.replace(slab);
+        })
+        .unwrap_or_else(|_| internal_error("tls access failure"))
 }
 
 #[no_mangle]
-pub unsafe extern fn __wbindgen_drop_anyref_slice(ptr: *mut JsValue, len: usize) {
+pub unsafe extern "C" fn __wbindgen_drop_anyref_slice(ptr: *mut JsValue, len: usize) {
     for slot in slice::from_raw_parts_mut(ptr, len) {
         __wbindgen_anyref_table_dealloc(slot.idx as usize);
     }
 }
 
+// Implementation of `__wbindgen_anyref_heap_live_count` for when we are using
+// `anyref` instead of the JS `heap`.
+#[no_mangle]
+pub unsafe extern "C" fn __wbindgen_anyref_heap_live_count_impl() -> u32 {
+    tl::HEAP_SLAB
+        .try_with(|slot| {
+            let slab = slot.replace(Slab::new());
+            let count = slab.live_count();
+            slot.replace(slab);
+            count
+        })
+        .unwrap_or_else(|_| internal_error("tls access failure"))
+}
+
 // see comment in module above this in `link_mem_intrinsics`
 #[inline(never)]
-pub fn link_intrinsics() {
-}
+pub fn link_intrinsics() {}

src/lib.rs

@@ -489,19 +489,26 @@ externs! {
     extern "C" {
         fn __wbindgen_object_clone_ref(idx: u32) -> u32;
         fn __wbindgen_object_drop_ref(idx: u32) -> ();
+
         fn __wbindgen_string_new(ptr: *const u8, len: usize) -> u32;
         fn __wbindgen_number_new(f: f64) -> u32;
-        fn __wbindgen_number_get(idx: u32, invalid: *mut u8) -> f64;
+        fn __wbindgen_symbol_new(ptr: *const u8, len: usize) -> u32;
+
+        fn __wbindgen_anyref_heap_live_count() -> u32;
+
         fn __wbindgen_is_null(idx: u32) -> u32;
         fn __wbindgen_is_undefined(idx: u32) -> u32;
-        fn __wbindgen_boolean_get(idx: u32) -> u32;
-        fn __wbindgen_symbol_new(ptr: *const u8, len: usize) -> u32;
         fn __wbindgen_is_symbol(idx: u32) -> u32;
         fn __wbindgen_is_object(idx: u32) -> u32;
         fn __wbindgen_is_function(idx: u32) -> u32;
         fn __wbindgen_is_string(idx: u32) -> u32;
+
+        fn __wbindgen_number_get(idx: u32, invalid: *mut u8) -> f64;
+        fn __wbindgen_boolean_get(idx: u32) -> u32;
         fn __wbindgen_string_get(idx: u32, len: *mut usize) -> *mut u8;
+
         fn __wbindgen_debug_string(idx: u32, len: *mut usize) -> *mut u8;
+
         fn __wbindgen_throw(a: *const u8, b: usize) -> !;
         fn __wbindgen_rethrow(a: u32) -> !;
 
@@ -650,6 +657,56 @@ pub fn throw_val(s: JsValue) -> ! {
     }
 }
 
+/// Get the count of live `anyref`s / `JsValue`s in `wasm-bindgen`'s heap.
+///
+/// ## Usage
+///
+/// This is intended for debugging and writing tests.
+///
+/// To write a test that asserts against unnecessarily keeping `anref`s /
+/// `JsValue`s alive:
+///
+/// * get an initial live count,
+///
+/// * perform some series of operations or function calls that should clean up
+///   after themselves, and should not keep holding onto `anyref`s / `JsValue`s
+///   after completion,
+///
+/// * get the final live count,
+///
+/// * and assert that the initial and final counts are the same.
+///
+/// ## What is Counted
+///
+/// Note that this only counts the *owned* `anyref`s / `JsValue`s that end up in
+/// `wasm-bindgen`'s heap. It does not count borrowed `anyref`s / `JsValue`s
+/// that are on its stack.
+///
+/// For example, these `JsValue`s are accounted for:
+///
+/// ```ignore
+/// #[wasm_bindgen]
+/// pub fn my_function(this_is_counted: JsValue) {
+///     let also_counted = JsValue::from_str("hi");
+///     assert!(wasm_bindgen::anyref_heap_live_count() >= 2);
+/// }
+/// ```
+///
+/// While this borrowed `JsValue` ends up on the stack, not the heap, and
+/// therefore is not accounted for:
+///
+/// ```ignore
+/// #[wasm_bindgen]
+/// pub fn my_other_function(this_is_not_counted: &JsValue) {
+///     // ...
+/// }
+/// ```
+pub fn anyref_heap_live_count() -> u32 {
+    unsafe {
+        __wbindgen_anyref_heap_live_count()
+    }
+}
+
 /// An extension trait for `Option<T>` and `Result<T, E>` for unwraping the `T`
 /// value, or throwing a JS error if it is not available.
 ///

tests/headless/anyref_heap_live_count.rs

@@ -0,0 +1,20 @@
+use wasm_bindgen::prelude::*;
+use wasm_bindgen_test::*;
+
+// This test is in the headless suite so that we can test the `anyref` table
+// implementation of `anyref_heap_live_count` (as opposed to the JS `heap`
+// implementation) in Firefox.
+#[wasm_bindgen_test]
+fn test_anyref_heap_live_count() {
+    let initial = wasm_bindgen::anyref_heap_live_count();
+
+    let after_alloc = {
+        let _vals: Vec<_> = (0..10).map(JsValue::from).collect();
+        wasm_bindgen::anyref_heap_live_count()
+    };
+
+    let after_dealloc = wasm_bindgen::anyref_heap_live_count();
+
+    assert_eq!(initial, after_dealloc);
+    assert_eq!(initial + 10, after_alloc);
+}

tests/headless/main.rs

@@ -49,3 +49,4 @@ pub fn import_export_same_name() {
 
 pub mod snippets;
 pub mod modules;
+pub mod anyref_heap_live_count;