1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
//! Implements dynamically typed arenas, where any type of item can be allocated.


use std::marker::PhantomData;
use std::cell::RefCell;
use std::os::raw::c_void;
use std::{mem, ptr};

pub mod bytes;
use bytes::ByteArena;

/// Marker trait that indicates whether or a `DynamicArena` may be sent across threads
pub trait SendAbility: Sized {
    fn create_arena<'a>() -> DynamicArena<'a, Self>;
}
/// Marker type that indicates you expect everything in the `DynamicArena` to be `Send`
///
/// Although this prevents you from allocating non-`Send` types in the arena,
/// it allows you to `Send` the dynamic itself arena across threads.
/// We can't safely implement `Send` for `DynamicArena` without this bound,
/// since you could otherwise place a `Rc` in the arena, send it across threads,
/// and then proceed to drop the arena and mutate the reference count.
pub enum Sendable {}
impl SendAbility for Sendable {
    #[inline]
    fn create_arena<'a>() -> DynamicArena<'a, Self> {
        DynamicArena::new_send()
    }
}
/// Marker type that indiates everything in the `DynamicArena` isn't nesicarrily `Send`.
///
/// This prevents you from `Send`ing the arena itself across threads,
/// as described in the `Sendable` docs.
pub enum NonSend {}
impl SendAbility for NonSend {
    #[inline]
    fn create_arena<'a>() -> DynamicArena<'a, Self> {
        DynamicArena::new_bounded()
    }
}

struct DynamicArenaItem {
    drop: unsafe fn(*mut c_void),
    value: *mut c_void
}
impl Drop for DynamicArenaItem {
    #[inline]
    fn drop(&mut self) {
        unsafe {
            (self.drop)(self.value)
        }
    }
}
unsafe impl Send for DynamicArenaItem {}

/// An alias for an arena allocator which requires that everything is `Send + 'a`.
pub type DynamicSendArena<'a> = DynamicArena<'a, Sendable>;

/// An arena allocator where any type of object can be allocated.
///
/// Unlike typed arenas, any `Sized` object can be allocated here and they
/// don't all have to have the same statically known type in advance.
/// Usually you don't have to worry about the arena's lifetime,
/// since it should be static by default (see).
///
/// ## Performance
/// Although this is _slightly_ slower than a `typed_arena::Arena`,
/// it can be much more memory and time efficient than having a ton of seperate typed arenas.
/// The only point where dynamic dispatch actually gets involved is when the arena is dropped,
/// since we have to dynamically dispatch the drop functions instead of statically dispatching them.
///
/// ## Safety
/// In order to prevent use after free in a `DynamicArena`, all pointers in the allocated items
/// need to be valid for the lifetime `'a` to ensure all references outlive the arena itself.
/// Unfortunately, this statically prevents all self referential structs with `alloc`,
/// since they can't be known ahead of time to be safe and outlive the arena itself,
/// and we can't perform [dropchk](https://doc.rust-lang.org/nightly/nomicon/dropck.html)
/// on a dynamically typed arena (only statically typed ones).
///
/// The alternatives to this are `alloc_unchecked`, which bypasses the lifetime and safety,
/// and `alloc_copy`, which bypasses the lifetime by ensuring `T: Copy`.
/// This is safe, since a copyable item can never have a custom drop,
/// and the drop function could never possibly trigger use after free.
///
/// This means you can use self-referential structs with a `DynamicArena` as long as they implement `Copy`.
/// One way to make your types implement copy and support self-refrential structs,
/// is by replacing owned objects with their borrowed counterparts and then allocating them in the arena.
/// For example
/// ````
/// struct OwnedSelfReferential<'a> {
///    next: Option<&'a OwnedSelfReferential<'a>>,
///    text: String,
///    array: Vec<u32>
/// }
/// ````
/// can't be used with either `alloc` (since it's self referential),
/// nor `alloc_copy` (since `String` and `Vec` need to be dropped).
/// However by replacing `String` with `&'a str` and `&'a [u32]`,
/// we can make the structure `Copy` and enable use with `alloc_copy`.
///
/// Then, when someone needs to arena-allocate the struct they can use
/// the same arena to allocate the `String` and `Vec<u32>` first,
/// before they proceed to allocate the copyable struct.
pub struct DynamicArena<'a, S = NonSend> {
    /// The underlying arena, where we request that they allocate arbitrary bytes.
    handle: ByteArena,
    /// The list of untyped values we've allocated in the arena,
    /// and whose drop functions need to be invoked.
    ///
    /// The drop functions are dynamically dispatched,
    /// and each item could invoke completely different code for completely different types.
    /// This is only needed for types that need to be dropped (as determined by `mem::needs_drop`),
    /// and types that need need to be dropped don't need to be added.
    items: RefCell<Vec<DynamicArenaItem>>,
    /// This is the magic `PhantomData` combination to have proper lifetime invariance.
    ///
    /// Otherwise the lifetime would be 'variant',
    /// and borrow checking wouldn't properly enforce the bound for `alloc`.
    /// This is enforced by the `compile-fail/invalid-drop-counted.rs`
    marker: PhantomData<*mut &'a ()>,
    send: PhantomData<S>
}
impl DynamicArena<'static, NonSend> {
    /// Create a new dynamic arena whose allocated items must outlive the `'static` lifetime,
    /// and whose items aren't required to be `Send`.
    ///
    /// Usually this is what you want,
    /// since it's only a bound for the allocated items.
    #[inline]
    pub fn new() -> Self {
        DynamicArena::new_bounded()
    }
}
impl<'a, S> DynamicArena<'a, S>  {
    pub fn with_capacity(item_capacity: usize, byte_capacity: usize) -> Self {
        DynamicArena {
            handle: ByteArena::with_capacity(byte_capacity),
            items: RefCell::new(Vec::with_capacity(item_capacity)),
            marker: PhantomData,
            send: PhantomData
        }
    }
    /// Allocate the specified value in this arena,
    /// returning a reference which will be valid for the lifetime of the entire arena.
    ///
    /// The bound on the item requires that `T: Copy`
    /// to ensure there's no drop function that needs to be invoked.
    #[inline]
    pub fn alloc_copy<T: Copy + Send>(&self, value: T) -> &mut T {
        unsafe {
            self.alloc_unchecked(value)
        }
    }
    /// Allocate the specified value in this arena,
    /// without ensuring that its drop function would be safe to invoke.
    ///
    /// Additionally this function leaks the underlying memory,
    /// and never runs the destructor, so its ownership needs to be tracked seperately.
    #[inline]
    pub unsafe fn alloc_unchecked<T>(&self, value: T) -> &mut T {
        let ptr = self.handle.alloc_uninitialized(mem::size_of::<T>()) as *mut T;
        ptr::write(ptr, value);
        &mut *ptr
    }
    /// Dynamically drop the specified value,
    /// invoking the drop function when the arena is dropped.
    ///
    /// This is unsafe because the drop function isn't necessarily safe to invoke,
    /// and the memory isn't nessicarrily .
    /// Not only are you assuming that `ptr::drop_in_place` would be safe,
    /// you're also assuming that the drop function won't reference any dangling pointers,
    /// and that [dropchk](https://doc.rust-lang.org/nomicon/dropck.html) would be successful.
    ///
    /// Normally these invariants are statically checked by the `alloc` method,
    /// which ensures that the memory is owned and all pointers
    /// would be valid for the lifetime of the entire arena.
    #[inline]
    pub unsafe fn dynamic_drop<T>(&self, value: *mut T) {
        if mem::needs_drop::<T>() {
            self.items.borrow_mut().push(DynamicArenaItem {
                drop: mem::transmute::<unsafe fn(*mut T),
                    unsafe fn(*mut c_void)>(ptr::drop_in_place::<T>),
                value: value as *mut c_void
            })
        }
    }
}
impl<'a> DynamicArena<'a, Sendable> {
    /// Create a new empty arena, bounded by the inferred lifetime for this type `'a`
    ///
    /// Since this arena has been marked `Sendable`,
    /// all items in the arena need to implement `Send`.
    pub fn new_send() -> Self {
        DynamicArena {
            handle: ByteArena::new(),
            items: RefCell::new(Vec::new()),
            marker: PhantomData,
            send: PhantomData
        }
    }
    /// Allocate the specified value in this arena,
    /// returning a reference which will be valid for the lifetime of the entire arena.
    ///
    /// The bound on this item requires that `T: 'a`
    /// to ensure the drop function is safe to invoke.
    /// Additionally, since the arena is `Sendable`,
    /// the bound on the item also requires that `T: Send`.
    #[inline]
    pub fn alloc<T: Send + 'a>(&self, value: T) -> &mut T {
        unsafe {
            let target = self.alloc_unchecked(value);
            self.dynamic_drop(target);
            target
        }
    }
}
impl<'a> DynamicArena<'a, NonSend> {
    /// Create a new empty arena, bounded by the inferred lifetime for this type `'a`
    ///
    /// Since this arena has been marked `NonSend`,
    /// the items in the arena don't necessarily need to implement `Send`.
    pub fn new_bounded() -> Self {
        DynamicArena {
            handle: ByteArena::new(),
            items: RefCell::new(Vec::new()),
            marker: PhantomData,
            send: PhantomData
        }
    }
    /// Allocate the specified value in this arena,
    /// returning a reference which will be valid for the lifetime of the entire arena.
    ///
    /// The bound on this item requires that `T: 'a`
    /// to ensure the drop function is safe to invoke.
    #[inline]
    pub fn alloc<T: 'a>(&self, value: T) -> &mut T {
        unsafe {
            let target = self.alloc_unchecked(value);
            self.dynamic_drop(target);
            target
        }
    }
}
impl<'a, S: SendAbility> Default for DynamicArena<'a, S> {
    #[inline]
    fn default() -> Self {
        S::create_arena()
    }
}
unsafe impl<'a> Send for DynamicArena<'a, Sendable> {}

#[cfg(test)]
mod test {
    use super::*;
    use std::cell::Cell;

    const EXPECTED_DROP_COUNT: u32 = 4787;
    const EXPECTED_DEPTHS: &[u32] = &[5, 27, 43];

    struct DropCounted<'a>(&'a Cell<u32>);
    impl<'a> Drop for DropCounted<'a> {
        fn drop(&mut self) {
            let old_count = self.0.get();
            self.0.set(old_count + 1);
        }
    }
    #[derive(Copy, Clone)]
    pub struct SelfReferential<'a>(u32, Option<&'a SelfReferential<'a>>);
    impl<'a> SelfReferential<'a> {
        #[inline]
        pub fn with_depth(arena: &'a DynamicArena, depth: u32) -> &'a Self {
            arena.alloc_copy(match depth {
                0 => SelfReferential(depth, None),
                _ => SelfReferential(depth, Some(
                    SelfReferential::with_depth(arena, depth - 1)))
            })
        }
        #[inline]
        pub fn depth(&self) -> u32 {
            match self.1 {
                Some(inner) => inner.depth() + 1,
                None => 0
            }
        }
    }
    #[test]
    fn test_send() {
        let arena = DynamicArena::new_send();
        verify_copyable(do_copyable(&arena));
        ::std::thread::spawn(move || {
            verify_copyable(do_copyable(&arena));
        });
    }

    #[test]
    fn copyable() {
        let arena = DynamicArena::new();
        for _ in 0..5 {
            verify_copyable(do_copyable(&arena));
        }
    }
    #[test]
    fn self_referential() {
        let arena = DynamicArena::new();
        for _ in 0..5 {
            verify_self_referential(do_self_referential(&arena));
        }
    }
    #[test]
    fn drop_counted() {
        let cell = Box::new(Cell::new(0));
        let arena = DynamicArena::new_bounded();
        {
            do_drop_counted(&arena, &cell);
            assert_eq!(cell.get(), 0);
        }
        drop(arena);
        assert_eq!(cell.get(), EXPECTED_DROP_COUNT);
    }
    #[test]
    fn mixed() {
        let cell = Cell::new(0);
        let arena = DynamicArena::new_bounded();
        {
            do_drop_counted(&arena, &cell);
            for _ in 0..5 {
                verify_copyable(do_copyable(&arena));
                verify_self_referential(do_self_referential(&arena));
            }
            assert_eq!(cell.get(), 0);
        }
        drop(arena);
        assert_eq!(cell.get(), EXPECTED_DROP_COUNT);
    }
    fn do_copyable<'a, S>(arena: &'a DynamicArena<S>) -> Vec<&'a u32> {
        let mut results = Vec::new();
        for i in 0..10 {
            results.push(&*arena.alloc_copy(i * 3));
        }
        results
    }
    fn verify_copyable(results: Vec<&u32>) {
        for (actual, expected) in results.iter().zip(0..10) {
            assert_eq!(**actual, expected * 3);
        }
    }
    fn do_drop_counted<'a, 'd: 'a>(arena: &'a DynamicArena<'d>, counter: &'d Cell<u32>) {
        for _ in 0..EXPECTED_DROP_COUNT {
            arena.alloc(DropCounted(counter));
        }
    }
    fn do_self_referential<'a>(arena: &'a DynamicArena) -> Vec<&'a SelfReferential<'a>> {
        let mut result = Vec::new();
        for &depth in EXPECTED_DEPTHS {
            result.push(SelfReferential::with_depth(arena, depth));
        }
        result
    }
    fn verify_self_referential<'a>(results: Vec<&'a SelfReferential<'a>>) {
        for (&actual, &depth) in results.iter().zip(EXPECTED_DEPTHS.iter()) {
            assert_eq!(actual.depth(), depth);
        }
    }
}