Add the anyref_heap_live_count function

This is useful for debugging and writing tests that assert various operations do
not leak `JsValue`s.

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

@@ -107,6 +107,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) -> ! {
@@ -205,6 +219,20 @@ pub unsafe extern fn __wbindgen_drop_anyref_slice(ptr: *mut JsValue, len: 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() {

src/lib.rs

@@ -494,6 +494,8 @@ externs! {
         fn __wbindgen_number_new(f: f64) -> u32;
         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_is_symbol(idx: u32) -> u32;
@@ -655,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;