Simplify allocation-recording in test suite using new ts_set_allocator API

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 tree_sitter::{allocations, LogType, Node, Parser, Tree};
+use std::fs;
+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::fixtures::{get_language, get_test_grammar, get_test_language};
+use super::helpers::{
+    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 tree_sitter::{allocations, Parser};
+use super::helpers::{allocations, fixtures::get_language};
+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,
-    QueryErrorKind, QueryMatch, QueryPredicate, QueryPredicateArg, QueryProperty,
+    Language, Node, Parser, Point, Query, QueryCapture, QueryCursor, QueryError, QueryErrorKind,
+    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 std::ffi::CStr;
-use std::ffi::CString;
-use std::{fs, ptr, slice, str};
-use tree_sitter::{allocations, Point};
-use tree_sitter_tags::c_lib as c;
-use tree_sitter_tags::{Error, TagsConfiguration, TagsContext};
+use super::helpers::{
+    allocations,
+    fixtures::{get_language, get_language_queries_path},
+};
+use std::{
+    ffi::{CStr, CString},
+    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#"
 (