tree_sitter_cli/fuzz/
allocations.rs

1use std::{
2    collections::HashMap,
3    os::raw::c_void,
4    sync::{
5        atomic::{AtomicBool, AtomicUsize, Ordering::SeqCst},
6        Mutex,
7    },
8};
9
10#[ctor::ctor]
11unsafe fn initialize_allocation_recording() {
12    tree_sitter::set_allocator(
13        Some(ts_record_malloc),
14        Some(ts_record_calloc),
15        Some(ts_record_realloc),
16        Some(ts_record_free),
17    );
18}
19
20#[derive(Debug, PartialEq, Eq, Hash)]
21struct Allocation(*const c_void);
22unsafe impl Send for Allocation {}
23unsafe impl Sync for Allocation {}
24
25#[derive(Default)]
26struct AllocationRecorder {
27    enabled: AtomicBool,
28    allocation_count: AtomicUsize,
29    outstanding_allocations: Mutex<HashMap<Allocation, usize>>,
30}
31
32thread_local! {
33    static RECORDER: AllocationRecorder = AllocationRecorder::default();
34}
35
36extern "C" {
37    fn malloc(size: usize) -> *mut c_void;
38    fn calloc(count: usize, size: usize) -> *mut c_void;
39    fn realloc(ptr: *mut c_void, size: usize) -> *mut c_void;
40    fn free(ptr: *mut c_void);
41}
42
43pub fn record<T>(f: impl FnOnce() -> T) -> Result<T, String> {
44    RECORDER.with(|recorder| {
45        recorder.enabled.store(true, SeqCst);
46        recorder.allocation_count.store(0, SeqCst);
47        recorder.outstanding_allocations.lock().unwrap().clear();
48    });
49
50    let value = f();
51
52    let outstanding_allocation_indices = RECORDER.with(|recorder| {
53        recorder.enabled.store(false, SeqCst);
54        recorder.allocation_count.store(0, SeqCst);
55        recorder
56            .outstanding_allocations
57            .lock()
58            .unwrap()
59            .drain()
60            .map(|e| e.1)
61            .collect::<Vec<_>>()
62    });
63    if !outstanding_allocation_indices.is_empty() {
64        return Err(format!(
65            "Leaked allocation indices: {outstanding_allocation_indices:?}",
66        ));
67    }
68    Ok(value)
69}
70
71fn record_alloc(ptr: *mut c_void) {
72    RECORDER.with(|recorder| {
73        if recorder.enabled.load(SeqCst) {
74            let count = recorder.allocation_count.fetch_add(1, SeqCst);
75            recorder
76                .outstanding_allocations
77                .lock()
78                .unwrap()
79                .insert(Allocation(ptr), count);
80        }
81    });
82}
83
84fn record_dealloc(ptr: *mut c_void) {
85    RECORDER.with(|recorder| {
86        if recorder.enabled.load(SeqCst) {
87            recorder
88                .outstanding_allocations
89                .lock()
90                .unwrap()
91                .remove(&Allocation(ptr));
92        }
93    });
94}
95
96unsafe extern "C" fn ts_record_malloc(size: usize) -> *mut c_void {
97    let result = malloc(size);
98    record_alloc(result);
99    result
100}
101
102unsafe extern "C" fn ts_record_calloc(count: usize, size: usize) -> *mut c_void {
103    let result = calloc(count, size);
104    record_alloc(result);
105    result
106}
107
108unsafe extern "C" fn ts_record_realloc(ptr: *mut c_void, size: usize) -> *mut c_void {
109    let result = realloc(ptr, size);
110    if ptr.is_null() {
111        record_alloc(result);
112    } else if ptr != result {
113        record_dealloc(ptr);
114        record_alloc(result);
115    }
116    result
117}
118
119unsafe extern "C" fn ts_record_free(ptr: *mut c_void) {
120    record_dealloc(ptr);
121    free(ptr);
122}
tree_sitter_cli/fuzz/allocations.rs

tree_sitter_cli/fuzz/
allocations.rs
