Skip to content

Commit 59ff059

Browse files
committed
Add ensure_capacity and rename min to len
1 parent e8b3aea commit 59ff059

File tree

1 file changed

+26
-15
lines changed

1 file changed

+26
-15
lines changed

src/libarena/lib.rs

+26-15
Original file line numberDiff line numberDiff line change
@@ -161,14 +161,21 @@ impl<T> TypedArena<T> {
161161
available_capacity_bytes >= at_least_bytes
162162
}
163163

164+
/// Ensures there's enough space in the current chunk to fit `len` objects.
165+
#[inline]
166+
fn ensure_capacity(&self, len: usize) {
167+
if !self.can_allocate(len) {
168+
self.grow(len);
169+
debug_assert!(self.can_allocate(len));
170+
}
171+
}
172+
164173
#[inline]
165174
unsafe fn alloc_raw_slice(&self, len: usize) -> *mut T {
166175
assert!(mem::size_of::<T>() != 0);
167176
assert!(len != 0);
168177

169-
if !self.can_allocate(len) {
170-
self.grow(len);
171-
}
178+
self.ensure_capacity(len);
172179

173180
let start_ptr = self.ptr.get();
174181
self.ptr.set(start_ptr.add(len));
@@ -203,19 +210,20 @@ impl<T> TypedArena<T> {
203210

204211
match size_hint {
205212
(min, Some(max)) if min == max => {
206-
if min == 0 {
213+
// We know the exact number of elements the iterator will produce here
214+
let len = min;
215+
216+
if len == 0 {
207217
return &mut [];
208218
}
209219

210-
if !self.can_allocate(min) {
211-
self.grow(min);
212-
}
220+
self.ensure_capacity(len);
213221

214222
let slice = self.ptr.get();
215223

216224
unsafe {
217225
let mut ptr = self.ptr.get();
218-
for _ in 0..min {
226+
for _ in 0..len {
219227
// Write into uninitialized memory.
220228
ptr::write(ptr, iter.next().unwrap());
221229
// Advance the pointer.
@@ -224,7 +232,7 @@ impl<T> TypedArena<T> {
224232
// we destroy the correct amount
225233
self.ptr.set(ptr);
226234
}
227-
slice::from_raw_parts_mut(slice, min)
235+
slice::from_raw_parts_mut(slice, len)
228236
}
229237
}
230238
_ => {
@@ -239,7 +247,7 @@ impl<T> TypedArena<T> {
239247
let len = vec.len();
240248
let start_ptr = self.alloc_raw_slice(len);
241249
vec.as_ptr().copy_to_nonoverlapping(start_ptr, len);
242-
vev.set_len(0);
250+
vec.set_len(0);
243251
slice::from_raw_parts_mut(start_ptr, len)
244252
}
245253
})
@@ -488,16 +496,19 @@ impl DroplessArena {
488496

489497
match size_hint {
490498
(min, Some(max)) if min == max => {
491-
if min == 0 {
499+
// We know the exact number of elements the iterator will produce here
500+
let len = min;
501+
502+
if len == 0 {
492503
return &mut []
493504
}
494-
let size = min.checked_mul(mem::size_of::<T>()).unwrap();
505+
let size = len.checked_mul(mem::size_of::<T>()).unwrap();
495506
let mem = self.alloc_raw(size, mem::align_of::<T>()) as *mut _ as *mut T;
496507
unsafe {
497-
for i in 0..min {
508+
for i in 0..len {
498509
ptr::write(mem.offset(i as isize), iter.next().unwrap())
499510
}
500-
slice::from_raw_parts_mut(mem, min)
511+
slice::from_raw_parts_mut(mem, len)
501512
}
502513
}
503514
(_, _) => {
@@ -515,7 +526,7 @@ impl DroplessArena {
515526
mem::align_of::<T>()
516527
) as *mut _ as *mut T;
517528
vec.as_ptr().copy_to_nonoverlapping(start_ptr, len);
518-
vev.set_len(0);
529+
vec.set_len(0);
519530
slice::from_raw_parts_mut(start_ptr, len)
520531
}
521532
})

0 commit comments

Comments
 (0)