Simplify allocation-recording in test suite using new ts_set_allocator API

This commit is contained in:
Max Brunsfeld 2021-12-30 16:09:07 -08:00
parent e01ea9ff51
commit 622359b400
19 changed files with 245 additions and 254 deletions

View file

@ -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
})
}

View file

@ -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;

View file

@ -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)

View file

@ -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 {}

View file

@ -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 {
free_ptr(self.ptr as *mut c_void);
}
unsafe { (FREE_FN)(self.ptr as *mut c_void) };
}
}