Assert no memory leaks by stubbing malloc/free in the test suite

cli/src/tests/allocations.rs

@@ -0,0 +1,104 @@
+#![cfg(test)]
+#![allow(dead_code)]
+
+use spin::Mutex;
+use std::collections::HashMap;
+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.enabled = true;
+    recorder.allocation_count = 0;
+    recorder.outstanding_allocations.clear();
+}
+
+pub fn stop_recording() {
+    let mut recorder = RECORDER.lock();
+    recorder.enabled = false;
+
+    if !recorder.outstanding_allocations.is_empty() {
+        panic!(
+            "Leaked allocation indices: {:?}",
+            recorder
+                .outstanding_allocations
+                .iter()
+                .map(|e| e.1)
+                .collect::<Vec<_>>()
+        );
+    }
+}
+
+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]
+extern "C" fn ts_record_malloc(size: c_ulong) -> *const c_void {
+    let result = unsafe { malloc(size) };
+    record_alloc(result);
+    result
+}
+
+#[no_mangle]
+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]
+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
+}
+
+#[no_mangle]
+extern "C" fn ts_record_free(ptr: *mut c_void) {
+    record_dealloc(ptr);
+    unsafe { free(ptr) };
+}
+
+#[no_mangle]
+extern "C" fn ts_record_allocations_toggle() {
+    let mut recorder = RECORDER.lock();
+    recorder.enabled = !recorder.enabled;
+}

cli/src/tests/corpuses.rs

@@ -1,9 +1,10 @@
+use super::allocations;
 use super::fixtures::{fixtures_dir, get_language, get_test_language};
 use crate::generate;
 use crate::test::{parse_tests, print_diff, print_diff_key, TestEntry};
 use crate::util;
 use std::fs;
-use tree_sitter::{LogType, Parser};
+use tree_sitter::{Language, LogType, Parser};
 
 const LANGUAGES: &'static [&'static str] = &[
     "bash",
@@ -27,8 +28,6 @@ lazy_static! {
 
 #[test]
 fn test_real_language_corpus_files() {
-    let mut log_session = None;
-    let mut parser = get_parser(&mut log_session, "log1.html");
     let grammars_dir = fixtures_dir().join("grammars");
 
     let mut did_fail = false;
@@ -44,8 +43,7 @@ fn test_real_language_corpus_files() {
         let language = get_language(language_name);
         let corpus_dir = grammars_dir.join(language_name).join("corpus");
         let test = parse_tests(&corpus_dir).unwrap();
-        parser.set_language(language).unwrap();
-        did_fail |= run_mutation_tests(&mut parser, test);
+        did_fail |= run_mutation_tests(language, test);
     }
 
     if did_fail {
@@ -55,8 +53,6 @@ fn test_real_language_corpus_files() {
 
 #[test]
 fn test_error_corpus_files() {
-    let mut log_session = None;
-    let mut parser = get_parser(&mut log_session, "log2.html");
     let corpus_dir = fixtures_dir().join("error_corpus");
 
     let mut did_fail = false;
@@ -74,8 +70,7 @@ fn test_error_corpus_files() {
 
         let test = parse_tests(&entry.path()).unwrap();
         let language = get_language(&language_name);
-        parser.set_language(language).unwrap();
-        did_fail |= run_mutation_tests(&mut parser, test);
+        did_fail |= run_mutation_tests(language, test);
     }
 
     if did_fail {
@@ -85,8 +80,6 @@ fn test_error_corpus_files() {
 
 #[test]
 fn test_feature_corpus_files() {
-    let mut log_session = None;
-    let mut parser = get_parser(&mut log_session, "log3.html");
     let test_grammars_dir = fixtures_dir().join("test_grammars");
 
     let mut did_fail = false;
@@ -132,8 +125,7 @@ fn test_feature_corpus_files() {
             let c_code = generate_result.unwrap().1;
             let language = get_test_language(language_name, c_code, &test_path);
             let test = parse_tests(&corpus_path).unwrap();
-            parser.set_language(language).unwrap();
-            did_fail |= run_mutation_tests(&mut parser, test);
+            did_fail |= run_mutation_tests(language, test);
         }
     }
 
@@ -142,7 +134,7 @@ fn test_feature_corpus_files() {
     }
 }
 
-fn run_mutation_tests(parser: &mut Parser, test: TestEntry) -> bool {
+fn run_mutation_tests(language: Language, test: TestEntry) -> bool {
     match test {
         TestEntry::Example {
             name,
@@ -157,23 +149,30 @@ fn run_mutation_tests(parser: &mut Parser, test: TestEntry) -> bool {
 
             eprintln!("  example: {:?}", name);
 
+            allocations::start_recording();
+            let mut log_session = None;
+            let mut parser = get_parser(&mut log_session, "log.html");
+            parser.set_language(language).unwrap();
             let tree = parser
                 .parse_utf8(&mut |byte_offset, _| &input[byte_offset..], None)
                 .unwrap();
             let actual = tree.root_node().to_sexp();
+            drop(tree);
+            drop(parser);
             if actual != output {
                 print_diff_key();
                 print_diff(&actual, &output);
                 println!("");
                 true
             } else {
+                allocations::stop_recording();
                 false
             }
         }
         TestEntry::Group { children, .. } => {
             let mut result = false;
             for child in children {
-                result |= run_mutation_tests(parser, child);
+                result |= run_mutation_tests(language, child);
             }
             result
         }

cli/src/tests/mod.rs

@@ -1,3 +1,4 @@
+mod allocations;
 mod corpuses;
 mod fixtures;
 mod parser_api;