diff --git a/compiler/rustc_arena/src/lib.rs b/compiler/rustc_arena/src/lib.rs index 166f7f53c41ac..703c212cd27d8 100644 --- a/compiler/rustc_arena/src/lib.rs +++ b/compiler/rustc_arena/src/lib.rs @@ -14,6 +14,8 @@ #![feature(dropck_eyepatch)] #![feature(new_uninit)] #![feature(maybe_uninit_slice)] +#![feature(min_specialization)] +#![feature(trusted_len)] #![cfg_attr(test, feature(test))] use rustc_data_structures::cold_path; @@ -22,6 +24,7 @@ use smallvec::SmallVec; use std::alloc::Layout; use std::cell::{Cell, RefCell}; use std::cmp; +use std::iter::TrustedLen; use std::marker::{PhantomData, Send}; use std::mem::{self, MaybeUninit}; use std::ptr; @@ -109,6 +112,53 @@ impl Default for TypedArena { } } +trait IterExt { + fn write_to_arena(iter: I, arena: &TypedArena) -> &mut [T]; +} + +impl IterExt for I +where + I: Iterator, +{ + #[inline] + default fn write_to_arena(iter: I, arena: &TypedArena) -> &mut [T] { + arena.alloc_from_iter_gen(iter) + } +} + +impl IterExt for I +where + I: Iterator + TrustedLen, +{ + #[inline] + fn write_to_arena(mut iter: I, arena: &TypedArena) -> &mut [T] { + let size_hint = iter.size_hint(); + + match size_hint { + (0, Some(_)) => &mut [], + (len, Some(_)) => { + // no need to check min == max because of TrustedLen + + // SAFETY: TrustedLen implementors must ensure they return exactly as many + // elements as they tell via `size_hint()`. + unsafe { + // We know the exact number of elements the iterator will produce here + let start_addr = arena.alloc_raw_slice(len); + let mut mem = start_addr; + + while let Some(value) = iter.next() { + ptr::write(mem, value); + mem = mem.add(1); + } + + return slice::from_raw_parts_mut(start_addr, len); + } + } + _ => panic!("Trying to allocate from an iterator with more than usize::MAX elements"), + } + } +} + impl TypedArena { /// Allocates an object in the `TypedArena`, returning a reference to it. #[inline] @@ -183,10 +233,10 @@ impl TypedArena { } } + /// General case implementation for `alloc_from_iter()`. #[inline] - pub fn alloc_from_iter>(&self, iter: I) -> &mut [T] { - assert!(mem::size_of::() != 0); - let mut vec: SmallVec<[_; 8]> = iter.into_iter().collect(); + fn alloc_from_iter_gen>(&self, iter: I) -> &mut [T] { + let mut vec: SmallVec<[_; 8]> = iter.collect(); if vec.is_empty() { return &mut []; } @@ -201,6 +251,14 @@ impl TypedArena { } } + #[inline] + pub fn alloc_from_iter>(&self, iter: I) -> &mut [T] { + let iter = iter.into_iter(); + assert!(mem::size_of::() != 0); + + I::IntoIter::write_to_arena(iter, self) + } + /// Grows the arena. #[inline(never)] #[cold]