diff --git a/compiler/rustc_arena/src/lib.rs b/compiler/rustc_arena/src/lib.rs index a0493056b816b..6789c844b9c12 100644 --- a/compiler/rustc_arena/src/lib.rs +++ b/compiler/rustc_arena/src/lib.rs @@ -19,16 +19,18 @@ #![cfg_attr(bootstrap, feature(min_const_generics))] #![feature(min_specialization)] #![cfg_attr(test, feature(test))] +#![feature(rustc_attrs)] use smallvec::SmallVec; use std::alloc::Layout; use std::cell::{Cell, RefCell}; use std::cmp; +use std::iter::Cloned; use std::marker::{PhantomData, Send}; use std::mem::{self, MaybeUninit}; use std::ptr; -use std::slice; +use std::slice::{self, Iter}; #[inline(never)] #[cold] @@ -36,6 +38,31 @@ pub fn cold_path R, R>(f: F) -> R { f() } +/// Move the contents of an iterator to a contiguous memory region. +/// +/// SAFETY: the caller must ensure that the destination memory is free and large enough to hold +/// the contents of the iterator. Must not be called for iterators that may allocate on the same +/// arena. +#[inline] +unsafe fn write_from_iter<'a, T, I: Iterator>( + mut iter: I, + len: usize, + mem: *mut T, +) -> &'a mut [T] { + let mut l = len; + for i in 0..len { + if let Some(value) = iter.next() { + std::ptr::write(mem.add(i), value); + } else { + l = i; + break; + } + } + // We only return as many items as the iterator gave us, even + // though it was supposed to give us `len` + return std::slice::from_raw_parts_mut(mem, l); +} + /// An arena that can hold objects of only one type. pub struct TypedArena { /// A pointer to the next object to be allocated. @@ -133,6 +160,10 @@ where } } +#[rustc_unsafe_specialization_marker] +trait Clonable: Clone {} +impl Clonable for T {} + impl IterExt for std::array::IntoIter { #[inline] fn alloc_from_iter(self, arena: &TypedArena) -> &mut [T] { @@ -167,6 +198,25 @@ impl IterExt for Vec { } } +impl<'a, T: 'a> IterExt for Cloned> +where + T: Clonable, +{ + #[inline] + fn alloc_from_iter(self, arena: &TypedArena) -> &mut [T] { + let len = self.len(); + if len == 0 { + return &mut []; + } + // Move the content to the arena by copying and then forgetting it + unsafe { + let len = self.len(); + let start_ptr = arena.alloc_raw_slice(len); + write_from_iter(self, len, start_ptr) + } + } +} + impl IterExt for SmallVec { #[inline] fn alloc_from_iter(mut self, arena: &TypedArena) -> &mut [A::Item] { @@ -489,28 +539,6 @@ impl DroplessArena { } } - #[inline] - unsafe fn write_from_iter>( - &self, - mut iter: I, - len: usize, - mem: *mut T, - ) -> &mut [T] { - let mut i = 0; - // Use a manual loop since LLVM manages to optimize it better for - // slice iterators - loop { - let value = iter.next(); - if i >= len || value.is_none() { - // We only return as many items as the iterator gave us, even - // though it was supposed to give us `len` - return slice::from_raw_parts_mut(mem, i); - } - ptr::write(mem.add(i), value.unwrap()); - i += 1; - } - } - #[inline] pub fn alloc_from_iter>(&self, iter: I) -> &mut [T] { let iter = iter.into_iter(); @@ -529,7 +557,7 @@ impl DroplessArena { } let mem = self.alloc_raw(Layout::array::(len).unwrap()) as *mut T; - unsafe { self.write_from_iter(iter, len, mem) } + unsafe { write_from_iter(iter, len, mem) } } (_, _) => { cold_path(move || -> &mut [T] {