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
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
//! Main public allocation methods.
//!
//! All ultimately call into `alloc_layout` or `try_alloc_layout`, defined in `alloc_impl.rs`.
use std::{
alloc::Layout,
ptr::{self},
slice, str,
};
use super::{Arena, bumpalo_alloc::AllocErr, utils::oom};
impl<const MIN_ALIGN: usize> Arena<MIN_ALIGN> {
/// Allocate an object in this `Arena`. Returns an exclusive reference to it.
///
/// # Panics
///
/// Panics if reserving space for `T` fails.
///
/// # Example
///
/// ```
/// # use oxc_allocator::arena::Arena;
///
/// let arena = Arena::new();
/// let x = arena.alloc("hello");
/// assert_eq!(*x, "hello");
/// ```
#[inline(always)]
pub fn alloc<T>(&self, val: T) -> &mut T {
self.alloc_with(|| val)
}
/// Try to allocate an object in this `Arena`. Returns an exclusive reference to it, if the allocation succeeds.
///
/// # Errors
///
/// Errors if reserving space for `T` fails.
///
/// # Example
///
/// ```
/// # use oxc_allocator::arena::Arena;
///
/// let arena = Arena::new();
/// let x = arena.try_alloc("hello");
/// assert_eq!(x, Ok(&mut "hello"));
/// ```
#[inline(always)]
pub fn try_alloc<T>(&self, val: T) -> Result<&mut T, AllocErr> {
self.try_alloc_with(|| val)
}
/// Pre-allocate space for an object in this `Arena`, and initialize it using the closure.
/// Returns an exclusive reference to it.
///
/// See [The `_with` Method Suffix](#initializer-functions-the-_with-method-suffix) for a discussion
/// of the differences between the `_with` suffixed methods and those methods without it,
/// their performance characteristics, and when you might or might not choose a `_with` suffixed method.
///
/// # Panics
///
/// Panics if reserving space for `T` fails.
///
/// # Example
///
/// ```
/// # use oxc_allocator::arena::Arena;
///
/// let arena = Arena::new();
/// let x = arena.alloc_with(|| "hello");
/// assert_eq!(*x, "hello");
/// ```
#[expect(clippy::mut_from_ref)]
#[inline(always)]
pub fn alloc_with<F, T>(&self, f: F) -> &mut T
where
F: FnOnce() -> T,
{
#[inline(always)]
unsafe fn inner_writer<T, F>(ptr: *mut T, f: F)
where
F: FnOnce() -> T,
{
// This function is translated as:
// - Allocate space for a T on the stack.
// - Call `f()` with the return value being put onto this stack space.
// - memcpy from the stack to the heap.
//
// Ideally we want LLVM to always realize that doing a stack allocation is unnecessary and optimize
// the code so it writes directly into the heap instead. It seems we get it to realize this most
// consistently if we put this critical line into it's own function instead of inlining it into the
// surrounding code.
unsafe { ptr::write(ptr, f()) };
}
let layout = Layout::new::<T>();
unsafe {
let p = self.alloc_layout(layout);
let p = p.as_ptr().cast::<T>();
inner_writer(p, f);
&mut *p
}
}
/// Try to pre-allocate space for an object in this `Arena`, and initialize it using the closure.
/// Returns an exclusive reference to it, if the allocation succeeds.
///
/// See [The `_with` Method Suffix](#initializer-functions-the-_with-method-suffix) for a discussion
/// of the differences between the `_with` suffixed methods and those methods without it,
/// their performance characteristics, and when you might or might not choose a `_with` suffixed method.
///
/// # Errors
///
/// Errors if reserving space for `T` fails.
///
/// # Example
///
/// ```
/// # use oxc_allocator::arena::Arena;
///
/// let arena = Arena::new();
/// let x = arena.try_alloc_with(|| "hello");
/// assert_eq!(x, Ok(&mut "hello"));
/// ```
#[expect(clippy::mut_from_ref)]
#[inline(always)]
pub fn try_alloc_with<F, T>(&self, f: F) -> Result<&mut T, AllocErr>
where
F: FnOnce() -> T,
{
#[inline(always)]
unsafe fn inner_writer<T, F>(ptr: *mut T, f: F)
where
F: FnOnce() -> T,
{
// This function is translated as:
// - Allocate space for a T on the stack.
// - Call `f()` with the return value being put onto this stack space.
// - memcpy from the stack to the heap.
//
// Ideally we want LLVM to always realize that doing a stack allocation is unnecessary and optimize
// the code so it writes directly into the heap instead. It seems we get it to realize this most
// consistently if we put this critical line into it's own function instead of inlining it into the
// surrounding code.
unsafe { ptr::write(ptr, f()) };
}
// Self-contained: `p` is allocated for `T` and then a `T` is written.
let layout = Layout::new::<T>();
let p = self.try_alloc_layout(layout)?;
let p = p.as_ptr().cast::<T>();
unsafe {
inner_writer(p, f);
Ok(&mut *p)
}
}
/// `Copy` a slice into this `Arena` and return an exclusive reference to the copy.
///
/// # Panics
///
/// Panics if reserving space for the slice fails.
///
/// # Example
///
/// ```
/// # use oxc_allocator::arena::Arena;
///
/// let arena = Arena::new();
/// let x = arena.alloc_slice_copy(&[1, 2, 3]);
/// assert_eq!(x, &[1, 2, 3]);
/// ```
#[expect(clippy::mut_from_ref)]
#[inline(always)]
pub fn alloc_slice_copy<T>(&self, src: &[T]) -> &mut [T]
where
T: Copy,
{
let layout = Layout::for_value(src);
let dst = self.alloc_layout(layout).cast::<T>();
unsafe {
ptr::copy_nonoverlapping(src.as_ptr(), dst.as_ptr(), src.len());
slice::from_raw_parts_mut(dst.as_ptr(), src.len())
}
}
/// `Clone` a slice into this `Arena`. Returns an exclusive reference to the clone.
/// Prefer [`alloc_slice_copy`](#method.alloc_slice_copy) if `T` is `Copy`.
///
/// # Panics
///
/// Panics if reserving space for the slice fails.
///
/// # Example
///
/// ```
/// # use oxc_allocator::arena::Arena;
///
/// #[derive(Clone, Debug, Eq, PartialEq)]
/// struct Sheep {
/// name: String,
/// }
///
/// let originals = [
/// Sheep { name: "Alice".into() },
/// Sheep { name: "Bob".into() },
/// Sheep { name: "Cathy".into() },
/// ];
///
/// let arena = Arena::new();
/// let clones = arena.alloc_slice_clone(&originals);
/// assert_eq!(originals, clones);
/// ```
#[expect(clippy::mut_from_ref)]
#[inline(always)]
pub fn alloc_slice_clone<T>(&self, src: &[T]) -> &mut [T]
where
T: Clone,
{
let layout = Layout::for_value(src);
let dst = self.alloc_layout(layout).cast::<T>();
unsafe {
for (i, val) in src.iter().cloned().enumerate() {
ptr::write(dst.as_ptr().add(i), val);
}
slice::from_raw_parts_mut(dst.as_ptr(), src.len())
}
}
/// `Copy` a string slice into this `Arena`. Returns an exclusive reference to it.
///
/// # Panics
///
/// Panics if reserving space for the string fails.
///
/// # Example
///
/// ```
/// # use oxc_allocator::arena::Arena;
///
/// let arena = Arena::new();
/// let hello = arena.alloc_str("hello world");
/// assert_eq!("hello world", hello);
/// ```
#[expect(clippy::mut_from_ref)]
#[inline(always)]
pub fn alloc_str(&self, src: &str) -> &mut str {
let buffer = self.alloc_slice_copy(src.as_bytes());
unsafe {
// This is OK, because it already came in as str, so it is guaranteed to be UTF-8
str::from_utf8_unchecked_mut(buffer)
}
}
/// Allocate a new slice of size `len` into this `Arena`. Returns an exclusive reference to the slice.
///
/// The elements of the slice are initialized using the supplied closure.
/// The closure argument is the position in the slice.
///
/// # Panics
///
/// Panics if reserving space for the slice fails.
///
/// # Example
///
/// ```
/// # use oxc_allocator::arena::Arena;
///
/// let arena = Arena::new();
/// let x = arena.alloc_slice_fill_with(5, |i| 5 * (i + 1));
/// assert_eq!(x, &[5, 10, 15, 20, 25]);
/// ```
#[expect(clippy::mut_from_ref)]
#[inline(always)]
pub fn alloc_slice_fill_with<T, F>(&self, len: usize, mut f: F) -> &mut [T]
where
F: FnMut(usize) -> T,
{
let layout = Layout::array::<T>(len).unwrap_or_else(|_| oom());
let dst = self.alloc_layout(layout).cast::<T>();
unsafe {
for i in 0..len {
ptr::write(dst.as_ptr().add(i), f(i));
}
let result = slice::from_raw_parts_mut(dst.as_ptr(), len);
debug_assert_eq!(Layout::for_value(result), layout);
result
}
}
/// Allocate a new slice of size `len` into this `Arena`. Returns an exclusive reference to the copy.
///
/// All elements of the slice are initialized to `value`.
///
/// # Panics
///
/// Panics if reserving space for the slice fails.
///
/// # Example
///
/// ```
/// # use oxc_allocator::arena::Arena;
///
/// let arena = Arena::new();
/// let x = arena.alloc_slice_fill_copy(5, 42);
/// assert_eq!(x, &[42, 42, 42, 42, 42]);
/// ```
#[inline(always)]
pub fn alloc_slice_fill_copy<T: Copy>(&self, len: usize, value: T) -> &mut [T] {
self.alloc_slice_fill_with(len, |_| value)
}
/// Allocate a new slice of size `len` into this `Arena`. Return an exclusive reference to the clone.
///
/// All elements of the slice are initialized to `value.clone()`.
///
/// # Panics
///
/// Panics if reserving space for the slice fails.
///
/// # Example
///
/// ```
/// # use oxc_allocator::arena::Arena;
///
/// let arena = Arena::new();
/// let s: String = "Hello Arena!".to_string();
/// let x: &[String] = arena.alloc_slice_fill_clone(2, &s);
/// assert_eq!(x.len(), 2);
/// assert_eq!(&x[0], &s);
/// assert_eq!(&x[1], &s);
/// ```
#[inline(always)]
pub fn alloc_slice_fill_clone<T: Clone>(&self, len: usize, value: &T) -> &mut [T] {
self.alloc_slice_fill_with(len, |_| value.clone())
}
/// Allocate a new slice of size `len` into this `Arena`. Returns an exclusive reference to the slice.
///
/// The elements are initialized using the supplied iterator.
///
/// # Panics
///
/// Panics if reserving space for the slice fails, or if the supplied iterator returns fewer elements than
/// it promised.
///
/// # Example
///
/// ```
/// # use oxc_allocator::arena::Arena;
///
/// let arena = Arena::new();
/// let x: &[i32] = arena.alloc_slice_fill_iter([2, 3, 5].iter().cloned().map(|i| i * i));
/// assert_eq!(x, [4, 9, 25]);
/// ```
#[inline(always)]
pub fn alloc_slice_fill_iter<T, I>(&self, iter: I) -> &mut [T]
where
I: IntoIterator<Item = T>,
I::IntoIter: ExactSizeIterator,
{
let mut iter = iter.into_iter();
self.alloc_slice_fill_with(iter.len(), |_| {
iter.next().expect("Iterator supplied too few elements")
})
}
/// Allocate a new slice of size `len` into this `Arena`. Returns an exclusive reference to the slice.
///
/// All elements of the slice are initialized to [`T::default()`].
///
/// [`T::default()`]: https://doc.rust-lang.org/std/default/trait.Default.html#tymethod.default
///
/// # Panics
///
/// Panics if reserving space for the slice fails.
///
/// # Example
///
/// ```
/// # use oxc_allocator::arena::Arena;
///
/// let arena = Arena::new();
/// let x = arena.alloc_slice_fill_default::<u32>(5);
/// assert_eq!(x, &[0, 0, 0, 0, 0]);
/// ```
#[inline(always)]
pub fn alloc_slice_fill_default<T: Default>(&self, len: usize) -> &mut [T] {
self.alloc_slice_fill_with(len, |_| T::default())
}
}