tree_sitter_cli/fuzz/
allocations.rs

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

tree_sitter_cli/fuzz/
allocations.rs
