Simplify setup for enabling/disabling allocation recording in the C lib

cli/Cargo.toml

@@ -41,6 +41,11 @@ webbrowser = "0.5.1"
 version = ">= 0.17.0"
 path = "../lib"
 
+[dev-dependencies.tree-sitter]
+version = ">= 0.17.0"
+path = "../lib"
+features = ["allocation-tracking"]
+
 [dependencies.tree-sitter-highlight]
 version = ">= 0.3.0"
 path = "../highlight"

cli/src/allocations_stubs.rs

@@ -0,0 +1,40 @@
+// In all dev builds, the tree-sitter library is built with the `allocation-tracking`
+// feature enabled. This causes the library to link against a set of externally
+// defined C functions like `ts_record_malloc` and `ts_record_free`. In tests, these
+// are defined to actually keep track of outstanding allocations. But when not running
+// tests, the symbols still need to be defined. This file provides pass-through
+// implementations of all of these functions.
+
+use std::os::raw::c_void;
+
+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);
+}
+
+#[no_mangle]
+unsafe extern "C" fn ts_record_malloc(size: usize) -> *const c_void {
+    malloc(size)
+}
+
+#[no_mangle]
+unsafe extern "C" fn ts_record_calloc(count: usize, size: usize) -> *const c_void {
+    calloc(count, size)
+}
+
+#[no_mangle]
+unsafe extern "C" fn ts_record_realloc(ptr: *mut c_void, size: usize) -> *const c_void {
+    realloc(ptr, size)
+}
+
+#[no_mangle]
+unsafe extern "C" fn ts_record_free(ptr: *mut c_void) {
+    free(ptr)
+}
+
+#[no_mangle]
+extern "C" fn ts_toggle_allocation_recording(_: bool) -> bool {
+    false
+}

cli/src/lib.rs

@@ -16,3 +16,6 @@ pub mod web_ui;
 
 #[cfg(test)]
 mod tests;
+
+#[cfg(not(test))]
+mod allocations_stubs;

cli/src/tests/helpers/allocations.rs

@@ -4,6 +4,7 @@
 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)]
@@ -31,9 +32,14 @@ extern "C" {
 
 pub fn start_recording() {
     let mut recorder = RECORDER.lock();
-    recorder.enabled = true;
     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() {

lib/Cargo.toml

@@ -28,3 +28,8 @@ cc = "1.0"
 
 [lib]
 path = "binding_rust/lib.rs"
+
+# This feature is only useful for testing the Tree-sitter library itself.
+# It is exposed because all of Tree-sitter's tests live in the Tree-sitter CLI crate.
+[features]
+allocation-tracking = []

lib/binding_rust/build.rs

@@ -21,9 +21,9 @@ fn main() {
 
     let mut config = cc::Build::new();
 
-    println!("cargo:rerun-if-env-changed=PROFILE");
-    if env::var("PROFILE").map_or(false, |s| s == "debug") {
-        config.define("TREE_SITTER_TEST", "");
+    println!("cargo:rerun-if-env-changed=CARGO_FEATURE_ALLOCATION_TRACKING");
+    if env::var("CARGO_FEATURE_ALLOCATION_TRACKING").is_ok() {
+        config.define("TREE_SITTER_ALLOCATION_TRACKING", "");
     }
 
     let src_path = Path::new("src");

lib/binding_rust/util.rs

@@ -1,71 +1,22 @@
 use std::os::raw::c_void;
 
 extern "C" {
-    /// In *Release* builds, the C library links directly against `malloc` and `free`.
-    ///
-    /// When freeing memory that was allocated by C code, use `free` directly.
-    #[cfg(not(debug_assertions))]
+    /// 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);
 
-    /// In *Test* builds, 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. These symbols are defined
-    /// in the `tree_sitter_cli::tests::helpers::allocations` module.
-    ///
-    /// When freeing memory that was allocated by C code, use the `free` function
-    /// from that module.
-    #[cfg(debug_assertions)]
+    /// 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);
 
-    /// In *Debug* builds, the C library is compiled the same as in test builds: using
-    /// the wrapper functions. This prevents the C library from having to be recompiled
-    /// constantly when switching between running tests and compiling with RLS.
-    ///
-    /// But we don't want to actually record allocations when running the library in
-    /// debug mode, so we define symbols like `ts_record_malloc` to just delegate to
-    /// the normal `malloc` functions.
-    #[cfg(all(debug_assertions, not(test)))]
-    fn malloc(size: usize) -> *mut c_void;
-    #[cfg(all(debug_assertions, not(test)))]
-    fn calloc(count: usize, size: usize) -> *mut c_void;
-    #[cfg(all(debug_assertions, not(test)))]
-    fn realloc(ptr: *mut c_void, size: usize) -> *mut c_void;
-    #[cfg(all(debug_assertions, not(test)))]
-    fn free(ptr: *mut c_void);
-}
-
-#[cfg(all(debug_assertions, not(test)))]
-#[no_mangle]
-unsafe extern "C" fn ts_record_malloc(size: usize) -> *const c_void {
-    malloc(size)
-}
-
-#[cfg(all(debug_assertions, not(test)))]
-#[no_mangle]
-unsafe extern "C" fn ts_record_calloc(count: usize, size: usize) -> *const c_void {
-    calloc(count, size)
-}
-
-#[cfg(all(debug_assertions, not(test)))]
-#[no_mangle]
-unsafe extern "C" fn ts_record_realloc(ptr: *mut c_void, size: usize) -> *const c_void {
-    realloc(ptr, size)
-}
-
-#[cfg(all(debug_assertions, not(test)))]
-#[no_mangle]
-unsafe extern "C" fn ts_record_free(ptr: *mut c_void) {
-    free(ptr)
-}
-
-#[cfg(all(debug_assertions, not(test)))]
-#[no_mangle]
-extern "C" fn ts_toggle_allocation_recording(_: bool) -> bool {
-    false
 }
 
+/// A raw pointer and a length, exposed as an iterator.
 pub struct CBufferIter<T> {
     ptr: *mut T,
     count: usize,

lib/src/alloc.h

@@ -9,7 +9,7 @@ extern "C" {
 #include <stdbool.h>
 #include <stdio.h>
 
-#if defined(TREE_SITTER_TEST)
+#if defined(TREE_SITTER_ALLOCATION_TRACKING)
 
 void *ts_record_malloc(size_t);
 void *ts_record_calloc(size_t, size_t);

script/test

@@ -31,7 +31,6 @@ OPTIONS
 EOF
 }
 
-export TREE_SITTER_TEST=1
 export RUST_TEST_THREADS=1
 export RUST_BACKTRACE=full
 

script/test.cmd

@@ -1,7 +1,6 @@
 @echo off
 
 setlocal
-set TREE_SITTER_TEST=1
 set RUST_TEST_THREADS=1
 set RUST_BACKTRACE=full
 cargo test -p tree-sitter-cli "%~1" -- --nocapture