refactor: deduplicate allocation code in tests

This commit is contained in:
Amaan Qureshi 2025-09-18 19:45:21 -04:00 committed by Amaan Qureshi
parent f09dc3cf46
commit 074e991280
5 changed files with 34 additions and 134 deletions

View file

@ -163,7 +163,7 @@ pub fn fuzz_language_corpus(
println!(" {test_index}. {test_name}");
let passed = allocations::record(|| {
let passed = allocations::record_checked(|| {
let mut log_session = None;
let mut parser = get_parser(&mut log_session, "log.html");
parser.set_language(language).unwrap();
@ -207,7 +207,7 @@ pub fn fuzz_language_corpus(
for trial in 0..options.iterations {
let seed = start_seed + trial;
let passed = allocations::record(|| {
let passed = allocations::record_checked(|| {
let mut rand = Rand::new(seed);
let mut log_session = None;
let mut parser = get_parser(&mut log_session, "log.html");

View file

@ -40,7 +40,11 @@ extern "C" {
fn free(ptr: *mut c_void);
}
pub fn record<T>(f: impl FnOnce() -> T) -> Result<T, String> {
pub fn record<T>(f: impl FnOnce() -> T) -> T {
record_checked(f).unwrap()
}
pub fn record_checked<T>(f: impl FnOnce() -> T) -> Result<T, String> {
RECORDER.with(|recorder| {
recorder.enabled.store(true, SeqCst);
recorder.allocation_count.store(0, SeqCst);
@ -93,19 +97,34 @@ fn record_dealloc(ptr: *mut c_void) {
});
}
unsafe extern "C" fn ts_record_malloc(size: usize) -> *mut c_void {
/// # Safety
///
/// The caller must ensure that the returned pointer is eventually
/// freed by calling `ts_record_free`.
#[must_use]
pub 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 {
/// # Safety
///
/// The caller must ensure that the returned pointer is eventually
/// freed by calling `ts_record_free`.
#[must_use]
pub 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 {
/// # Safety
///
/// The caller must ensure that the returned pointer is eventually
/// freed by calling `ts_record_free`.
#[must_use]
pub unsafe extern "C" fn ts_record_realloc(ptr: *mut c_void, size: usize) -> *mut c_void {
let result = realloc(ptr, size);
if ptr.is_null() {
record_alloc(result);
@ -116,7 +135,11 @@ unsafe extern "C" fn ts_record_realloc(ptr: *mut c_void, size: usize) -> *mut c_
result
}
unsafe extern "C" fn ts_record_free(ptr: *mut c_void) {
/// # Safety
///
/// The caller must ensure that `ptr` was allocated by a previous call
/// to `ts_record_malloc`, `ts_record_calloc`, or `ts_record_realloc`.
pub unsafe extern "C" fn ts_record_free(ptr: *mut c_void) {
record_dealloc(ptr);
free(ptr);
}

View file

@ -216,8 +216,7 @@ pub fn test_language_corpus(
}
true
})
.unwrap();
});
if !passed {
failure_count += 1;
@ -312,7 +311,7 @@ pub fn test_language_corpus(
}
true
}).unwrap();
});
if !passed {
failure_count += 1;
@ -434,8 +433,7 @@ fn test_feature_corpus_files() {
println!();
false
}
})
.unwrap();
});
if !passed {
failure_count += 1;

View file

@ -1,4 +1,4 @@
pub mod allocations;
pub use crate::fuzz::allocations;
pub mod edits;
pub(super) mod fixtures;
pub(super) mod query_helpers;

View file

@ -1,121 +0,0 @@
use std::{
collections::HashMap,
os::raw::c_void,
sync::{
atomic::{AtomicBool, AtomicUsize, 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: AtomicUsize,
outstanding_allocations: Mutex<HashMap<Allocation, usize>>,
}
thread_local! {
static RECORDER: AllocationRecorder = AllocationRecorder::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<_>>()
});
assert!(
outstanding_allocation_indices.is_empty(),
"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 {
let result = realloc(ptr, size);
if ptr.is_null() {
record_alloc(result);
} else if !core::ptr::eq(ptr, result) {
record_dealloc(ptr);
record_alloc(result);
}
result
}
unsafe extern "C" fn ts_record_free(ptr: *mut c_void) {
record_dealloc(ptr);
free(ptr);
}