Simplify allocation-recording in test suite using new ts_set_allocator API

Cargo.lock

@@ -591,12 +591,6 @@ version = "2.5.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "75ce4f9dc4a41b4c3476cc925f1efb11b66df373a8fde5d4b8915fa91b5d995e"
 
-[[package]]
-name = "spin"
-version = "0.7.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "13287b4da9d1207a4f4929ac390916d64eacfe236a487e9a9f5b3be392be5162"
-
 [[package]]
 name = "strsim"
 version = "0.8.0"
@@ -720,7 +714,6 @@ dependencies = [
  "cc",
  "lazy_static",
  "regex",
- "spin",
 ]
 
 [[package]]
@@ -731,6 +724,7 @@ dependencies = [
  "anyhow",
  "atty",
  "clap",
+ "ctor",
  "difference",
  "dirs",
  "glob",

cli/Cargo.toml

@@ -43,11 +43,6 @@ which = "4.1.0"
 version = "0.20"
 path = "../lib"
 
-[dev-dependencies.tree-sitter]
-version = "0.20"
-path = "../lib"
-features = ["allocation-tracking"]
-
 [dependencies.tree-sitter-config]
 version = "0.19.0"
 path = "config"
@@ -76,6 +71,7 @@ features = ["std"]
 rand = "0.8"
 tempfile = "3"
 pretty_assertions = "0.7.2"
+ctor = "0.1"
 
 [build-dependencies]
 toml = "0.5"

cli/src/tests/corpus_test.rs

@@ -1,4 +1,5 @@
 use super::helpers::{
+    allocations,
     edits::{get_random_edit, invert_edit},
     fixtures::{fixtures_dir, get_language, get_test_language},
     random::Rand,
@@ -11,8 +12,8 @@ use crate::{
     test::{parse_tests, print_diff, print_diff_key, strip_sexp_fields, TestEntry},
     util,
 };
-use std::{fs, usize};
+use std::fs;
-use tree_sitter::{allocations, LogType, Node, Parser, Tree};
+use tree_sitter::{LogType, Node, Parser, Tree};
 
 const EDIT_COUNT: usize = 3;
 const TRIAL_COUNT: usize = 10;

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;

cli/src/tests/parser_test.rs

@@ -1,10 +1,13 @@
-use super::helpers::edits::ReadRecorder;
+use super::helpers::{
-use super::helpers::fixtures::{get_language, get_test_grammar, get_test_language};
+    allocations,
+    edits::ReadRecorder,
+    fixtures::{get_language, get_test_grammar, get_test_language},
+};
 use crate::generate::generate_parser_for_grammar;
 use crate::parse::{perform_edit, Edit};
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::{thread, time};
-use tree_sitter::{allocations, IncludedRangesError, InputEdit, LogType, Parser, Point, Range};
+use tree_sitter::{IncludedRangesError, InputEdit, LogType, Parser, Point, Range};
 
 #[test]
 fn test_parsing_simple_string() {

cli/src/tests/pathological_test.rs

@@ -1,5 +1,5 @@
-use super::helpers::fixtures::get_language;
+use super::helpers::{allocations, fixtures::get_language};
-use tree_sitter::{allocations, Parser};
+use tree_sitter::Parser;
 
 #[test]
 fn test_pathological_example_1() {

cli/src/tests/query_test.rs

@@ -1,4 +1,5 @@
 use super::helpers::{
+    allocations,
     fixtures::get_language,
     query_helpers::{Match, Pattern},
 };
@@ -6,8 +7,8 @@ use lazy_static::lazy_static;
 use rand::{prelude::StdRng, SeedableRng};
 use std::{env, fmt::Write};
 use tree_sitter::{
-    allocations, Language, Node, Parser, Point, Query, QueryCapture, QueryCursor, QueryError,
+    Language, Node, Parser, Point, Query, QueryCapture, QueryCursor, QueryError, QueryErrorKind,
-    QueryErrorKind, QueryMatch, QueryPredicate, QueryPredicateArg, QueryProperty,
+    QueryMatch, QueryPredicate, QueryPredicateArg, QueryProperty,
 };
 
 lazy_static! {

cli/src/tests/tags_test.rs

@@ -1,10 +1,13 @@
-use super::helpers::fixtures::{get_language, get_language_queries_path};
+use super::helpers::{
-use std::ffi::CStr;
+    allocations,
-use std::ffi::CString;
+    fixtures::{get_language, get_language_queries_path},
-use std::{fs, ptr, slice, str};
+};
-use tree_sitter::{allocations, Point};
+use std::{
-use tree_sitter_tags::c_lib as c;
+    ffi::{CStr, CString},
-use tree_sitter_tags::{Error, TagsConfiguration, TagsContext};
+    fs, ptr, slice, str,
+};
+use tree_sitter::Point;
+use tree_sitter_tags::{c_lib as c, Error, TagsConfiguration, TagsContext};
 
 const PYTHON_TAG_QUERY: &'static str = r#"
 (

lib/Cargo.toml

@@ -22,17 +22,11 @@ include = [
 ]
 
 [dependencies]
-lazy_static = { version="1.2.0", optional=true }
+lazy_static = { version = "1.2.0", optional = true }
 regex = "1"
-spin = { version="0.7", optional=true }
 
 [build-dependencies]
 cc = "^1.0.58"
 
 [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 = ["lazy_static", "spin"]

lib/binding_rust/allocations.rs

@@ -1,120 +0,0 @@
-use spin::Mutex;
-use std::{
-    collections::HashMap,
-    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>,
-}
-
-thread_local! {
-    static RECORDER: Mutex<AllocationRecorder> = Default::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 record<T>(f: impl FnOnce() -> T) -> T {
-    RECORDER.with(|recorder| {
-        let mut recorder = recorder.lock();
-        recorder.enabled = true;
-        recorder.allocation_count = 0;
-        recorder.outstanding_allocations.clear();
-    });
-
-    let value = f();
-
-    let outstanding_allocation_indices = RECORDER.with(|recorder| {
-        let mut recorder = recorder.lock();
-        recorder.enabled = false;
-        recorder.allocation_count = 0;
-        recorder
-            .outstanding_allocations
-            .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| {
-        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) {
-    RECORDER.with(|recorder| {
-        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 {
-    RECORDER.with(|recorder| {
-        let mut recorder = recorder.lock();
-        let was_enabled = recorder.enabled;
-        recorder.enabled = enabled;
-        was_enabled
-    })
-}

lib/binding_rust/bindings.rs

@@ -1,4 +1,4 @@
-/* automatically generated by rust-bindgen 0.59.1 */
+/* automatically generated by rust-bindgen 0.59.2 */
 
 pub type TSSymbol = u16;
 pub type TSFieldId = u16;
@@ -29,11 +29,11 @@ pub struct TSQueryCursor {
 }
 pub const TSInputEncoding_TSInputEncodingUTF8: TSInputEncoding = 0;
 pub const TSInputEncoding_TSInputEncodingUTF16: TSInputEncoding = 1;
-pub type TSInputEncoding = u32;
+pub type TSInputEncoding = ::std::os::raw::c_uint;
 pub const TSSymbolType_TSSymbolTypeRegular: TSSymbolType = 0;
 pub const TSSymbolType_TSSymbolTypeAnonymous: TSSymbolType = 1;
 pub const TSSymbolType_TSSymbolTypeAuxiliary: TSSymbolType = 2;
-pub type TSSymbolType = u32;
+pub type TSSymbolType = ::std::os::raw::c_uint;
 #[repr(C)]
 #[derive(Debug, Copy, Clone)]
 pub struct TSPoint {
@@ -64,7 +64,7 @@ pub struct TSInput {
 }
 pub const TSLogType_TSLogTypeParse: TSLogType = 0;
 pub const TSLogType_TSLogTypeLex: TSLogType = 1;
-pub type TSLogType = u32;
+pub type TSLogType = ::std::os::raw::c_uint;
 #[repr(C)]
 #[derive(Debug, Copy, Clone)]
 pub struct TSLogger {
@@ -118,7 +118,7 @@ pub struct TSQueryMatch {
 pub const TSQueryPredicateStepType_TSQueryPredicateStepTypeDone: TSQueryPredicateStepType = 0;
 pub const TSQueryPredicateStepType_TSQueryPredicateStepTypeCapture: TSQueryPredicateStepType = 1;
 pub const TSQueryPredicateStepType_TSQueryPredicateStepTypeString: TSQueryPredicateStepType = 2;
-pub type TSQueryPredicateStepType = u32;
+pub type TSQueryPredicateStepType = ::std::os::raw::c_uint;
 #[repr(C)]
 #[derive(Debug, Copy, Clone)]
 pub struct TSQueryPredicateStep {
@@ -132,7 +132,7 @@ pub const TSQueryError_TSQueryErrorField: TSQueryError = 3;
 pub const TSQueryError_TSQueryErrorCapture: TSQueryError = 4;
 pub const TSQueryError_TSQueryErrorStructure: TSQueryError = 5;
 pub const TSQueryError_TSQueryErrorLanguage: TSQueryError = 6;
-pub type TSQueryError = u32;
+pub type TSQueryError = ::std::os::raw::c_uint;
 extern "C" {
     #[doc = " Create a new parser."]
     pub fn ts_parser_new() -> *mut TSParser;
@@ -823,6 +823,37 @@ extern "C" {
     #[doc = " See also `ts_parser_set_language`."]
     pub fn ts_language_version(arg1: *const TSLanguage) -> u32;
 }
+extern "C" {
+    #[doc = " Set the allocation functions used by the library."]
+    #[doc = ""]
+    #[doc = " By default, Tree-sitter uses the standard libc allocation functions,"]
+    #[doc = " but aborts the process when an allocation fails. This function lets"]
+    #[doc = " you supply alternative allocation functions at runtime."]
+    #[doc = ""]
+    #[doc = " If you pass `NULL` for any parameter, Tree-sitter will switch back to"]
+    #[doc = " its default implementation of that function."]
+    #[doc = ""]
+    #[doc = " If you call this function after the library has already been used, then"]
+    #[doc = " you must ensure that either:"]
+    #[doc = "  1. All the existing objects have been freed."]
+    #[doc = "  2. The new allocator shares its state with the old one, so it is capable"]
+    #[doc = "     of freeing memory that was allocated by the old allocator."]
+    pub fn ts_set_allocator(
+        new_malloc: ::std::option::Option<
+            unsafe extern "C" fn(arg1: usize) -> *mut ::std::os::raw::c_void,
+        >,
+        new_calloc: ::std::option::Option<
+            unsafe extern "C" fn(arg1: usize, arg2: usize) -> *mut ::std::os::raw::c_void,
+        >,
+        new_realloc: ::std::option::Option<
+            unsafe extern "C" fn(
+                arg1: *mut ::std::os::raw::c_void,
+                arg2: usize,
+            ) -> *mut ::std::os::raw::c_void,
+        >,
+        new_free: ::std::option::Option<unsafe extern "C" fn(arg1: *mut ::std::os::raw::c_void)>,
+    );
+}
 
 pub const TREE_SITTER_LANGUAGE_VERSION: usize = 13;
 pub const TREE_SITTER_MIN_COMPATIBLE_LANGUAGE_VERSION: usize = 13;

lib/binding_rust/build.rs

@@ -1,5 +1,3 @@
-extern crate cc;
-
 use std::path::{Path, PathBuf};
 use std::{env, fs};
 
@@ -19,13 +17,6 @@ fn main() {
         }
     }
 
-    let mut config = cc::Build::new();
-
-    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");
     for entry in fs::read_dir(&src_path).unwrap() {
         let entry = entry.unwrap();
@@ -33,7 +24,7 @@ fn main() {
         println!("cargo:rerun-if-changed={}", path.to_str().unwrap());
     }
 
-    config
+    cc::Build::new()
         .flag_if_supported("-std=c99")
         .flag_if_supported("-Wno-unused-parameter")
         .include(src_path)

lib/binding_rust/lib.rs

@@ -1,9 +1,6 @@
 mod ffi;
 mod util;
 
-#[cfg(feature = "allocation-tracking")]
-pub mod allocations;
-
 #[cfg(unix)]
 use std::os::unix::io::AsRawFd;
 
@@ -1040,7 +1037,7 @@ impl<'tree> Node<'tree> {
             .to_str()
             .unwrap()
             .to_string();
-        unsafe { util::free_ptr(c_string as *mut c_void) };
+        unsafe { (FREE_FN)(c_string as *mut c_void) };
         result
     }
 
@@ -2164,6 +2161,22 @@ impl fmt::Display for QueryError {
     }
 }
 
+extern "C" {
+    fn free(ptr: *mut c_void);
+}
+
+static mut FREE_FN: unsafe extern "C" fn(ptr: *mut c_void) = free;
+
+pub unsafe fn set_allocator(
+    new_malloc: Option<unsafe extern "C" fn(usize) -> *mut c_void>,
+    new_calloc: Option<unsafe extern "C" fn(usize, usize) -> *mut c_void>,
+    new_realloc: Option<unsafe extern "C" fn(*mut c_void, usize) -> *mut c_void>,
+    new_free: Option<unsafe extern "C" fn(*mut c_void)>,
+) {
+    FREE_FN = new_free.unwrap_or(free);
+    ffi::ts_set_allocator(new_malloc, new_calloc, new_realloc, new_free);
+}
+
 impl error::Error for IncludedRangesError {}
 impl error::Error for LanguageError {}
 impl error::Error for QueryError {}

lib/binding_rust/util.rs

@@ -1,19 +1,6 @@
+use super::FREE_FN;
 use std::os::raw::c_void;
 
-#[cfg(not(feature = "allocation-tracking"))]
-extern "C" {
-    /// Normally, use `free(1)` to free memory allocated from C.
-    #[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")]
-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,
@@ -50,8 +37,6 @@ impl<T: Copy> ExactSizeIterator for CBufferIter<T> {}
 
 impl<T> Drop for CBufferIter<T> {
     fn drop(&mut self) {
-        unsafe {
+        unsafe { (FREE_FN)(self.ptr as *mut c_void) };
-            free_ptr(self.ptr as *mut c_void);
-        }
     }
 }

lib/include/tree_sitter/api.h

@@ -138,22 +138,6 @@ typedef enum {
 /* Section - Parser */
 /********************/
 
-/**
- * Switch to a new allocator.
- *
- * This function can be invoked more than once. However, the application needs
- * to pay attention to the memory allocated by the old allocator but might be
- * freed by the new one.
- *
- * Specifically, the application either,
- *  1. ensures all parsers and trees are freed before calling it;
- *  2. provides an allocator that shares its state with the old allocator.
- */
-void ts_set_allocator(void *(*new_malloc)(size_t),
-		      void *(*new_calloc)(size_t, size_t),
-		      void *(*new_realloc)(void *, size_t),
-		      void (*new_free)(void *));
-
 /**
  * Create a new parser.
  */
@@ -912,6 +896,33 @@ TSSymbolType ts_language_symbol_type(const TSLanguage *, TSSymbol);
  */
 uint32_t ts_language_version(const TSLanguage *);
 
+/**********************************/
+/* Section - Global Configuration */
+/**********************************/
+
+/**
+ * Set the allocation functions used by the library.
+ *
+ * By default, Tree-sitter uses the standard libc allocation functions,
+ * but aborts the process when an allocation fails. This function lets
+ * you supply alternative allocation functions at runtime.
+ * 
+ * If you pass `NULL` for any parameter, Tree-sitter will switch back to
+ * its default implementation of that function.
+ * 
+ * If you call this function after the library has already been used, then
+ * you must ensure that either:
+ *  1. All the existing objects have been freed.
+ *  2. The new allocator shares its state with the old one, so it is capable
+ *     of freeing memory that was allocated by the old allocator.
+ */
+void ts_set_allocator(
+  void *(*new_malloc)(size_t),
+	void *(*new_calloc)(size_t, size_t),
+	void *(*new_realloc)(void *, size_t),
+	void (*new_free)(void *)
+);
+
 #ifdef __cplusplus
 }
 #endif

lib/src/alloc.c

@@ -1,29 +1,7 @@
 #include "alloc.h"
-
-#if defined(TREE_SITTER_ALLOCATION_TRACKING)
-
-void *ts_record_malloc(size_t);
-void *ts_record_calloc(size_t, size_t);
-void *ts_record_realloc(void *, size_t);
-void ts_record_free(void *);
-bool ts_toggle_allocation_recording(bool);
-
-void *(*ts_current_malloc)(size_t) = ts_record_malloc;
-void *(*ts_current_calloc)(size_t, size_t) = ts_record_calloc;
-void *(*ts_current_realloc)(void *, size_t) = ts_record_realloc;
-void (*ts_current_free)(void *) = ts_record_free;
-
-#else
-
 #include <stdlib.h>
 
-bool ts_toggle_allocation_recording(bool value) {
+static void *ts_malloc_default(size_t size) {
-  (void)value;
-  return false;
-}
-
-
-static inline void *ts_malloc_default(size_t size) {
   void *result = malloc(size);
   if (size > 0 && !result) {
     fprintf(stderr, "tree-sitter failed to allocate %zu bytes", size);
@@ -32,7 +10,7 @@ static inline void *ts_malloc_default(size_t size) {
   return result;
 }
 
-static inline void *ts_calloc_default(size_t count, size_t size) {
+static void *ts_calloc_default(size_t count, size_t size) {
   void *result = calloc(count, size);
   if (count > 0 && !result) {
     fprintf(stderr, "tree-sitter failed to allocate %zu bytes", count * size);
@@ -41,7 +19,7 @@ static inline void *ts_calloc_default(size_t count, size_t size) {
   return result;
 }
 
-static inline void *ts_realloc_default(void *buffer, size_t size) {
+static void *ts_realloc_default(void *buffer, size_t size) {
   void *result = realloc(buffer, size);
   if (size > 0 && !result) {
     fprintf(stderr, "tree-sitter failed to reallocate %zu bytes", size);
@@ -50,27 +28,21 @@ static inline void *ts_realloc_default(void *buffer, size_t size) {
   return result;
 }
 
-static inline void ts_free_default(void *buffer) {
-  free(buffer);
-}
-
-
 // Allow clients to override allocation functions dynamically
 void *(*ts_current_malloc)(size_t) = ts_malloc_default;
 void *(*ts_current_calloc)(size_t, size_t) = ts_calloc_default;
 void *(*ts_current_realloc)(void *, size_t) = ts_realloc_default;
-void (*ts_current_free)(void *) = ts_free_default;
+void (*ts_current_free)(void *) = free;
 
-#endif // defined(TREE_SITTER_ALLOCATION_TRACKING)
+void ts_set_allocator(
-
+  void *(*new_malloc)(size_t),
-void ts_set_allocator(void *(*new_malloc)(size_t),
+  void *(*new_calloc)(size_t, size_t),
-                      void *(*new_calloc)(size_t, size_t),
+  void *(*new_realloc)(void *, size_t),
-                      void *(*new_realloc)(void *, size_t),
+  void (*new_free)(void *)
-                      void (*new_free)(void *))
+) {
-{
+  ts_current_malloc = new_malloc ? new_malloc : ts_malloc_default;
-  ts_current_malloc = new_malloc;
+  ts_current_calloc = new_calloc ? new_calloc : ts_calloc_default;
-  ts_current_calloc = new_calloc;
+  ts_current_realloc = new_realloc ? new_realloc : ts_realloc_default;
-  ts_current_realloc = new_realloc;
+  ts_current_free = new_free ? new_free : free;
-  ts_current_free = new_free;
 }
 

lib/src/alloc.h

@@ -30,8 +30,6 @@ extern void (*ts_current_free)(void *);
 #define ts_free    ts_current_free
 #endif
 
-bool ts_toggle_allocation_recording(bool);
-
 #ifdef __cplusplus
 }
 #endif

lib/src/stack.c

@@ -722,7 +722,6 @@ void ts_stack_clear(Stack *self) {
 
 bool ts_stack_print_dot_graph(Stack *self, const TSLanguage *language, FILE *f) {
   array_reserve(&self->iterators, 32);
-  bool was_recording_allocations = ts_toggle_allocation_recording(false);
   if (!f) f = stderr;
 
   fprintf(f, "digraph stack {\n");
@@ -851,7 +850,6 @@ bool ts_stack_print_dot_graph(Stack *self, const TSLanguage *language, FILE *f)
   fprintf(f, "}\n");
 
   array_delete(&visited_nodes);
-  ts_toggle_allocation_recording(was_recording_allocations);
   return true;
 }