Move allocation tracking into lib crate

We have several test cases defined in the `cli` crate that depend on the
`lib` crate's `allocation-tracking` feature.  The implementation of the
actual allocation tracker used to live in the `cli` crate, close to the
test cases that use it.  The `allocation-tracking` feature in the `lib`
crate was just used to tell the tree-sitter implementation to expect
that the allocation tracker exists, and to use it.

That pattern meant that we had a circular dependency: `cli` depends on
`lib`, but `lib` required some code that was implemented in `cli`.
That, in turn, caused linker errors — but only when compiling in certain
configurations! [1]

This patch moves all of the allocation tracking implementation into the
`lib` crate, gated on the existing `allocation-tracking` feature, which
fixes the circular dependency.

Note that this patch does **not** fix the fact that feature unification
causes the `lib` crate to be built with the `allocation-tracking`
feature enabled, even though it's not a default.  Fixing that depends on
the forthcoming version 2 feature resolver [2], or using the `dev_dep`
workaround [3] in the meantime.

[1] https://github.com/tree-sitter/tree-sitter/issues/919
[2] https://doc.rust-lang.org/nightly/cargo/reference/features.html#feature-resolver-version-2
[3] https://github.com/tree-sitter/tree-sitter/issues/919#issuecomment-777107086

lib/binding_rust/allocations.rs

@@ -0,0 +1,118 @@
+use lazy_static::lazy_static;
+use spin::Mutex;
+use std::collections::HashMap;
+use std::env;
+use std::os::raw::{c_ulong, c_void};
+
+#[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: bool,
+    allocation_count: u64,
+    outstanding_allocations: HashMap<Allocation, u64>,
+}
+
+lazy_static! {
+    static ref RECORDER: Mutex<AllocationRecorder> = Mutex::new(AllocationRecorder::default());
+}
+
+extern "C" {
+    fn malloc(size: c_ulong) -> *mut c_void;
+    fn calloc(count: c_ulong, size: c_ulong) -> *mut c_void;
+    fn realloc(ptr: *mut c_void, size: c_ulong) -> *mut c_void;
+    fn free(ptr: *mut c_void);
+}
+
+pub fn start_recording() {
+    let mut recorder = RECORDER.lock();
+    recorder.allocation_count = 0;
+    recorder.outstanding_allocations.clear();
+
+    if env::var("RUST_TEST_THREADS").map_or(false, |s| s == "1") {
+        recorder.enabled = true;
+    } else {
+        panic!("This test must be run with RUST_TEST_THREADS=1. Use script/test.");
+    }
+}
+
+pub fn stop_recording() {
+    let mut recorder = RECORDER.lock();
+    recorder.enabled = false;
+
+    if !recorder.outstanding_allocations.is_empty() {
+        let mut allocation_indices = recorder
+            .outstanding_allocations
+            .iter()
+            .map(|e| e.1)
+            .collect::<Vec<_>>();
+        allocation_indices.sort_unstable();
+        panic!("Leaked allocation indices: {:?}", allocation_indices);
+    }
+}
+
+pub fn record(f: impl FnOnce()) {
+    start_recording();
+    f();
+    stop_recording();
+}
+
+fn record_alloc(ptr: *mut c_void) {
+    let mut recorder = RECORDER.lock();
+    if recorder.enabled {
+        let count = recorder.allocation_count;
+        recorder.allocation_count += 1;
+        recorder
+            .outstanding_allocations
+            .insert(Allocation(ptr), count);
+    }
+}
+
+fn record_dealloc(ptr: *mut c_void) {
+    let mut recorder = RECORDER.lock();
+    if recorder.enabled {
+        recorder.outstanding_allocations.remove(&Allocation(ptr));
+    }
+}
+
+#[no_mangle]
+pub extern "C" fn ts_record_malloc(size: c_ulong) -> *const c_void {
+    let result = unsafe { malloc(size) };
+    record_alloc(result);
+    result
+}
+
+#[no_mangle]
+pub extern "C" fn ts_record_calloc(count: c_ulong, size: c_ulong) -> *const c_void {
+    let result = unsafe { calloc(count, size) };
+    record_alloc(result);
+    result
+}
+
+#[no_mangle]
+pub extern "C" fn ts_record_realloc(ptr: *mut c_void, size: c_ulong) -> *const c_void {
+    record_dealloc(ptr);
+    let result = unsafe { realloc(ptr, size) };
+    record_alloc(result);
+    result
+}
+
+// This needs to be unsafe because it's reexported as crate::util::free_ptr, which is mapped to
+// libc's `free` function when the allocation-tracking feature is disabled.  Since `free` is
+// unsafe, this function needs to be too.
+#[no_mangle]
+pub unsafe extern "C" fn ts_record_free(ptr: *mut c_void) {
+    record_dealloc(ptr);
+    free(ptr);
+}
+
+#[no_mangle]
+pub extern "C" fn ts_toggle_allocation_recording(enabled: bool) -> bool {
+    let mut recorder = RECORDER.lock();
+    let was_enabled = recorder.enabled;
+    recorder.enabled = enabled;
+    was_enabled
+}

lib/binding_rust/lib.rs

@@ -1,6 +1,9 @@
 mod ffi;
 mod util;
 
+#[cfg(feature = "allocation-tracking")]
+pub mod allocations;
+
 #[cfg(unix)]
 use std::os::unix::io::AsRawFd;
 

lib/binding_rust/util.rs

@@ -1,21 +1,19 @@
 use std::os::raw::c_void;
 
+#[cfg(not(feature = "allocation-tracking"))]
 extern "C" {
     /// Normally, use `free(1)` to free memory allocated from C.
-    #[cfg(not(feature = "allocation-tracking"))]
     #[link_name = "free"]
     pub fn free_ptr(ptr: *mut c_void);
-
-    /// When the `allocation-tracking` feature is enabled, the C library is compiled with
-    /// the `TREE_SITTER_TEST` macro, so all calls to `malloc`, `free`, etc are linked
-    /// against wrapper functions called `ts_record_malloc`, `ts_record_free`, etc.
-    /// When freeing buffers allocated from C, use the wrapper `free` function.
-    #[cfg(feature = "allocation-tracking")]
-    #[link_name = "ts_record_free"]
-    pub fn free_ptr(ptr: *mut c_void);
-
 }
 
+/// When the `allocation-tracking` feature is enabled, the C library is compiled with
+/// the `TREE_SITTER_TEST` macro, so all calls to `malloc`, `free`, etc are linked
+/// against wrapper functions called `ts_record_malloc`, `ts_record_free`, etc.
+/// When freeing buffers allocated from C, use the wrapper `free` function.
+#[cfg(feature = "allocation-tracking")]
+pub use crate::allocations::ts_record_free as free_ptr;
+
 /// A raw pointer and a length, exposed as an iterator.
 pub struct CBufferIter<T> {
     ptr: *mut T,