Skip to content

Commit 822b89e

Browse files
committed
Optimize alloc_from_iter
1 parent ce797de commit 822b89e

File tree

1 file changed

+24
-5
lines changed

1 file changed

+24
-5
lines changed

src/libarena/lib.rs

+24-5
Original file line numberDiff line numberDiff line change
@@ -486,9 +486,31 @@ impl DroplessArena {
486486
}
487487
}
488488

489+
#[inline]
490+
unsafe fn write_from_iter<T, I: Iterator<Item = T>>(
491+
&self,
492+
mut iter: I,
493+
len: usize,
494+
mem: *mut T,
495+
) -> &mut [T] {
496+
let mut i = 0;
497+
// Use a manual loop since LLVM manages to optimize it better for
498+
// slice iterators
499+
loop {
500+
let value = iter.next();
501+
if i >= len || value.is_none() {
502+
// We only return as many items as the iterator gave us, even
503+
// though it was supposed to give us `len`
504+
return slice::from_raw_parts_mut(mem, i);
505+
}
506+
ptr::write(mem.offset(i as isize), value.unwrap());
507+
i += 1;
508+
}
509+
}
510+
489511
#[inline]
490512
pub fn alloc_from_iter<T, I: IntoIterator<Item = T>>(&self, iter: I) -> &mut [T] {
491-
let mut iter = iter.into_iter();
513+
let iter = iter.into_iter();
492514
assert!(mem::size_of::<T>() != 0);
493515
assert!(!mem::needs_drop::<T>());
494516

@@ -505,10 +527,7 @@ impl DroplessArena {
505527
let size = len.checked_mul(mem::size_of::<T>()).unwrap();
506528
let mem = self.alloc_raw(size, mem::align_of::<T>()) as *mut _ as *mut T;
507529
unsafe {
508-
for i in 0..len {
509-
ptr::write(mem.offset(i as isize), iter.next().unwrap())
510-
}
511-
slice::from_raw_parts_mut(mem, len)
530+
self.write_from_iter(iter, len, mem)
512531
}
513532
}
514533
(_, _) => {

0 commit comments

Comments
 (0)