cubecl_ir/
arena.rs

1use alloc::vec::Vec;
2
3use bumpalo::Bump;
4use cubecl_macros_internal::TypeHash;
5
6/// Bump allocator that runs drop on non-trivially droppable values on reset.
7/// This is useful to prevent memory leaks without drastically affecting performance for the vast
8/// majority of values that do not need custom drops.
9///
10/// # SAFETY
11///
12/// Values inserted *must not* access external borrowed data in a custom `Drop` implementation.
13/// This shouldn't happen for normal types, but should be kept in mind regardless.
14/// Contrived Example:
15///
16/// ```no_run
17/// # use cubecl_ir::arena::DropBump;
18/// struct Wrapper<'a>(&'a mut String);
19///
20/// impl<'a> Drop for Wrapper<'a> {
21///     fn drop(&mut self) {
22///         self.0.push_str(" dropped"); // Accesses external borrow
23///     }
24/// }
25///
26/// let mut external = String::from("hello");
27///
28/// let mut pool = DropBump::new();
29/// let w = pool.alloc(Wrapper(&mut external));
30/// drop(external);
31/// pool.reset(); // Runs `drop_in_place::<Wrapper>`
32/// ```
33///
34/// So don't do that.
35///
36#[derive(Default, Debug, TypeHash)]
37pub struct DropBump {
38    bump: Bump,
39    drop_thunks: Vec<DropThunk>,
40}
41
42#[derive(Debug, TypeHash)]
43struct DropThunk {
44    func: fn(*mut ()),
45    ptr: *mut (),
46}
47
48impl DropBump {
49    pub fn new() -> Self {
50        Default::default()
51    }
52
53    pub fn reset(&mut self) {
54        for drop in self.drop_thunks.drain(..).rev() {
55            (drop.func)(drop.ptr)
56        }
57        self.bump.reset();
58    }
59
60    pub fn alloc<T>(&mut self, val: T) -> &mut T {
61        let ret = self.bump.alloc(val);
62
63        if core::mem::needs_drop::<T>() {
64            self.drop_thunks.push(DropThunk {
65                func: drop_glue::<T>,
66                ptr: ret as *mut T as *mut (),
67            });
68        }
69
70        ret
71    }
72}
73
74fn drop_glue<T>(ptr: *mut ()) {
75    unsafe { core::ptr::drop_in_place::<T>(ptr as *mut T) };
76}
cubecl_ir/arena.rs

cubecl_ir/
arena.rs
