bolero_generator/driver/
cache.rs

1use alloc::{
2    boxed::Box,
3    collections::{btree_map::Entry as MapEntry, BTreeMap},
4    vec::Vec,
5};
6use core::{
7    any::{Any, TypeId},
8    fmt,
9    ops::Bound,
10};
11
12pub struct Cache {
13    max_total_entries: usize,
14    total_entries: usize,
15    max_entries: usize,
16    entries: BTreeMap<TypeId, Box<dyn Any>>,
17}
18
19impl Default for Cache {
20    fn default() -> Self {
21        Self {
22            max_total_entries: usize::MAX,
23            total_entries: 0,
24            max_entries: 1024,
25            entries: Default::default(),
26        }
27    }
28}
29
30impl fmt::Debug for Cache {
31    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
32        f.debug_struct("Cache")
33            .field("total_entries", &self.total_entries)
34            .finish()
35    }
36}
37
38impl core::panic::RefUnwindSafe for Cache {}
39
40impl Cache {
41    #[inline]
42    pub fn put<T: 'static>(&mut self, value: T) {
43        if self.total_entries >= self.max_total_entries {
44            return;
45        }
46
47        let id = TypeId::of::<Vec<T>>();
48
49        match self.entries.entry(id) {
50            MapEntry::Occupied(mut queue) => {
51                let queue = queue.get_mut();
52
53                let queue = unsafe { &mut *(queue as *mut Box<dyn Any> as *mut Box<Vec<T>>) };
54
55                if queue.len() >= self.max_entries {
56                    return;
57                }
58
59                queue.push(value);
60            }
61            MapEntry::Vacant(queue) => {
62                queue.insert(Box::new(alloc::vec![value]));
63            }
64        };
65
66        self.total_entries += 1;
67    }
68
69    #[inline]
70    pub fn get<T: 'static>(&mut self) -> Option<T> {
71        if self.total_entries == 0 {
72            return None;
73        }
74
75        let id = TypeId::of::<Vec<T>>();
76        let queue = self.entries.get_mut(&id)?;
77        let queue = unsafe { &mut *(queue as *mut Box<dyn Any> as *mut Box<Vec<T>>) };
78        let value = queue.pop()?;
79        self.total_entries -= 1;
80        Some(value)
81    }
82}
83
84pub struct Driver<'a, I: super::Driver> {
85    cache: &'a mut Cache,
86    inner: I,
87}
88
89impl<'a, I: super::Driver> Driver<'a, I> {
90    #[inline]
91    pub fn new(inner: I, cache: &'a mut Cache) -> Self {
92        Self { cache, inner }
93    }
94}
95
96impl<I: super::Driver> AsRef<I> for Driver<'_, I> {
97    #[inline]
98    fn as_ref(&self) -> &I {
99        &self.inner
100    }
101}
102
103impl<I: super::Driver> super::Driver for Driver<'_, I> {
104    #[inline(always)]
105    fn depth(&self) -> usize {
106        self.inner.depth()
107    }
108
109    #[inline(always)]
110    fn set_depth(&mut self, depth: usize) {
111        self.inner.set_depth(depth)
112    }
113
114    #[inline(always)]
115    fn max_depth(&self) -> usize {
116        self.inner.max_depth()
117    }
118
119    #[inline(always)]
120    fn gen_variant(&mut self, variants: usize, base_case: usize) -> Option<usize> {
121        self.inner.gen_variant(variants, base_case)
122    }
123
124    #[inline(always)]
125    fn gen_u8(&mut self, min: Bound<&u8>, max: Bound<&u8>) -> Option<u8> {
126        self.inner.gen_u8(min, max)
127    }
128
129    #[inline(always)]
130    fn gen_i8(&mut self, min: Bound<&i8>, max: Bound<&i8>) -> Option<i8> {
131        self.inner.gen_i8(min, max)
132    }
133
134    #[inline(always)]
135    fn gen_u16(&mut self, min: Bound<&u16>, max: Bound<&u16>) -> Option<u16> {
136        self.inner.gen_u16(min, max)
137    }
138
139    #[inline(always)]
140    fn gen_i16(&mut self, min: Bound<&i16>, max: Bound<&i16>) -> Option<i16> {
141        self.inner.gen_i16(min, max)
142    }
143
144    #[inline(always)]
145    fn gen_u32(&mut self, min: Bound<&u32>, max: Bound<&u32>) -> Option<u32> {
146        self.inner.gen_u32(min, max)
147    }
148
149    #[inline(always)]
150    fn gen_i32(&mut self, min: Bound<&i32>, max: Bound<&i32>) -> Option<i32> {
151        self.inner.gen_i32(min, max)
152    }
153
154    #[inline(always)]
155    fn gen_u64(&mut self, min: Bound<&u64>, max: Bound<&u64>) -> Option<u64> {
156        self.inner.gen_u64(min, max)
157    }
158
159    #[inline(always)]
160    fn gen_i64(&mut self, min: Bound<&i64>, max: Bound<&i64>) -> Option<i64> {
161        self.inner.gen_i64(min, max)
162    }
163
164    #[inline(always)]
165    fn gen_u128(&mut self, min: Bound<&u128>, max: Bound<&u128>) -> Option<u128> {
166        self.inner.gen_u128(min, max)
167    }
168
169    #[inline(always)]
170    fn gen_i128(&mut self, min: Bound<&i128>, max: Bound<&i128>) -> Option<i128> {
171        self.inner.gen_i128(min, max)
172    }
173
174    #[inline(always)]
175    fn gen_usize(&mut self, min: Bound<&usize>, max: Bound<&usize>) -> Option<usize> {
176        self.inner.gen_usize(min, max)
177    }
178
179    #[inline(always)]
180    fn gen_isize(&mut self, min: Bound<&isize>, max: Bound<&isize>) -> Option<isize> {
181        self.inner.gen_isize(min, max)
182    }
183
184    #[inline(always)]
185    fn gen_f32(&mut self, min: Bound<&f32>, max: Bound<&f32>) -> Option<f32> {
186        self.inner.gen_f32(min, max)
187    }
188
189    #[inline(always)]
190    fn gen_f64(&mut self, min: Bound<&f64>, max: Bound<&f64>) -> Option<f64> {
191        self.inner.gen_f64(min, max)
192    }
193
194    #[inline(always)]
195    fn gen_char(&mut self, min: Bound<&char>, max: Bound<&char>) -> Option<char> {
196        self.inner.gen_char(min, max)
197    }
198
199    #[inline(always)]
200    fn gen_bool(&mut self, probability: Option<f32>) -> Option<bool> {
201        self.inner.gen_bool(probability)
202    }
203
204    #[inline(always)]
205    fn gen_from_bytes<Hint, Gen, T>(&mut self, hint: Hint, produce: Gen) -> Option<T>
206    where
207        Hint: FnOnce() -> (usize, Option<usize>),
208        Gen: FnMut(&[u8]) -> Option<(usize, T)>,
209    {
210        self.inner.gen_from_bytes(hint, produce)
211    }
212
213    #[inline(always)]
214    fn cache_put<T: 'static>(&mut self, value: T) {
215        self.cache.put(value);
216    }
217
218    #[inline(always)]
219    fn cache_get<T: 'static>(&mut self) -> Option<T> {
220        self.cache.get()
221    }
222}
223
224#[cfg(test)]
225mod tests {
226    use super::*;
227
228    #[test]
229    fn cache_test() {
230        let mut cache = Cache::default();
231
232        cache.put(123usize);
233        assert_eq!(cache.get(), Some(123usize));
234        assert_eq!(cache.get::<usize>(), None);
235    }
236}