Simplify allocation-recording in test suite using new ts_set_allocator API

cli/src/tests/helpers/allocations.rs

@@ -0,0 +1,119 @@
+use std::{
+    collections::HashMap,
+    os::raw::c_void,
+    sync::{
+        atomic::{AtomicBool, AtomicU64, Ordering::SeqCst},
+        Mutex,
+    },
+};
+
+#[ctor::ctor]
+unsafe fn initialize_allocation_recording() {
+    tree_sitter::set_allocator(
+        Some(ts_record_malloc),
+        Some(ts_record_calloc),
+        Some(ts_record_realloc),
+        Some(ts_record_free),
+    );
+}
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+struct Allocation(*const c_void);
+unsafe impl Send for Allocation {}
+unsafe impl Sync for Allocation {}
+
+#[derive(Default)]
+struct AllocationRecorder {
+    enabled: AtomicBool,
+    allocation_count: AtomicU64,
+    outstanding_allocations: Mutex<HashMap<Allocation, u64>>,
+}
+
+thread_local! {
+    static RECORDER: AllocationRecorder = Default::default();
+}
+
+extern "C" {
+    fn malloc(size: usize) -> *mut c_void;
+    fn calloc(count: usize, size: usize) -> *mut c_void;
+    fn realloc(ptr: *mut c_void, size: usize) -> *mut c_void;
+    fn free(ptr: *mut c_void);
+}
+
+pub fn record<T>(f: impl FnOnce() -> T) -> T {
+    RECORDER.with(|recorder| {
+        recorder.enabled.store(true, SeqCst);
+        recorder.allocation_count.store(0, SeqCst);
+        recorder.outstanding_allocations.lock().unwrap().clear();
+    });
+
+    let value = f();
+
+    let outstanding_allocation_indices = RECORDER.with(|recorder| {
+        recorder.enabled.store(false, SeqCst);
+        recorder.allocation_count.store(0, SeqCst);
+        recorder
+            .outstanding_allocations
+            .lock()
+            .unwrap()
+            .drain()
+            .map(|e| e.1)
+            .collect::<Vec<_>>()
+    });
+    if !outstanding_allocation_indices.is_empty() {
+        panic!(
+            "Leaked allocation indices: {:?}",
+            outstanding_allocation_indices
+        );
+    }
+    value
+}
+
+fn record_alloc(ptr: *mut c_void) {
+    RECORDER.with(|recorder| {
+        if recorder.enabled.load(SeqCst) {
+            let count = recorder.allocation_count.fetch_add(1, SeqCst);
+            recorder
+                .outstanding_allocations
+                .lock()
+                .unwrap()
+                .insert(Allocation(ptr), count);
+        }
+    });
+}
+
+fn record_dealloc(ptr: *mut c_void) {
+    RECORDER.with(|recorder| {
+        if recorder.enabled.load(SeqCst) {
+            recorder
+                .outstanding_allocations
+                .lock()
+                .unwrap()
+                .remove(&Allocation(ptr));
+        }
+    });
+}
+
+unsafe extern "C" fn ts_record_malloc(size: usize) -> *mut c_void {
+    let result = malloc(size);
+    record_alloc(result);
+    result
+}
+
+unsafe extern "C" fn ts_record_calloc(count: usize, size: usize) -> *mut c_void {
+    let result = calloc(count, size);
+    record_alloc(result);
+    result
+}
+
+unsafe extern "C" fn ts_record_realloc(ptr: *mut c_void, size: usize) -> *mut c_void {
+    record_dealloc(ptr);
+    let result = realloc(ptr, size);
+    record_alloc(result);
+    result
+}
+
+unsafe extern "C" fn ts_record_free(ptr: *mut c_void) {
+    record_dealloc(ptr);
+    free(ptr);
+}

cli/src/tests/helpers/mod.rs

@@ -1,3 +1,4 @@
+pub(super) mod allocations;
 pub(super) mod edits;
 pub(super) mod fixtures;
 pub(super) mod query_helpers;