library/alloc/src/vec/source_iter_marker.rs

   1 use core::iter::{InPlaceIterable, SourceIter};
   2 use core::mem::{self, ManuallyDrop};
   3 use core::ptr::{self};
   4
   5 use super::{AsIntoIter, InPlaceDrop, SpecFromIter, SpecFromIterNested, Vec};
   6
   7 /// Specialization marker for collecting an iterator pipeline into a Vec while reusing the
   8 /// source allocation, i.e. executing the pipeline in place.
   9 ///
  10 /// The SourceIter parent trait is necessary for the specializing function to access the allocation
  11 /// which is to be reused. But it is not sufficient for the specialization to be valid. See
  12 /// additional bounds on the impl.
  13 #[rustc_unsafe_specialization_marker]
  14 pub(super) trait SourceIterMarker: SourceIter<Source: AsIntoIter> {}
  15
  16 // The std-internal SourceIter/InPlaceIterable traits are only implemented by chains of
  17 // Adapter<Adapter<Adapter<IntoIter>>> (all owned by core/std). Additional bounds
  18 // on the adapter implementations (beyond `impl<I: Trait> Trait for Adapter<I>`) only depend on other
  19 // traits already marked as specialization traits (Copy, TrustedRandomAccess, FusedIterator).
  20 // I.e. the marker does not depend on lifetimes of user-supplied types. Modulo the Copy hole, which
  21 // several other specializations already depend on.
  22 impl<T> SourceIterMarker for T where T: SourceIter<Source: AsIntoIter> + InPlaceIterable {}
  23
  24 impl<T, I> SpecFromIter<T, I> for Vec<T>
  25 where
  26     I: Iterator<Item = T> + SourceIterMarker,
  27 {
  28     default fn from_iter(mut iterator: I) -> Self {
  29         // Additional requirements which cannot expressed via trait bounds. We rely on const eval
  30         // instead:
  31         // a) no ZSTs as there would be no allocation to reuse and pointer arithmetic would panic
  32         // b) size match as required by Alloc contract
  33         // c) alignments match as required by Alloc contract
  34         if mem::size_of::<T>() == 0
  35             || mem::size_of::<T>()
  36                 != mem::size_of::<<<I as SourceIter>::Source as AsIntoIter>::Item>()
  37             || mem::align_of::<T>()
  38                 != mem::align_of::<<<I as SourceIter>::Source as AsIntoIter>::Item>()
  39         {
  40             // fallback to more generic implementations
  41             return SpecFromIterNested::from_iter(iterator);
  42         }
  43
  44         let (src_buf, src_ptr, dst_buf, dst_end, cap) = unsafe {
  45             let inner = iterator.as_inner().as_into_iter();
  46             (
  47                 inner.buf.as_ptr(),
  48                 inner.ptr,
  49                 inner.buf.as_ptr() as *mut T,
  50                 inner.end as *const T,
  51                 inner.cap,
  52             )
  53         };
  54
  55         // use try-fold since
  56         // - it vectorizes better for some iterator adapters
  57         // - unlike most internal iteration methods, it only takes a &mut self
  58         // - it lets us thread the write pointer through its innards and get it back in the end
  59         let sink = InPlaceDrop { inner: dst_buf, dst: dst_buf };
  60         let sink = iterator
  61             .try_fold::<_, _, Result<_, !>>(sink, write_in_place_with_drop(dst_end))
  62             .unwrap();
  63         // iteration succeeded, don't drop head
  64         let dst = ManuallyDrop::new(sink).dst;
  65
  66         let src = unsafe { iterator.as_inner().as_into_iter() };
  67         // check if SourceIter contract was upheld
  68         // caveat: if they weren't we may not even make it to this point
  69         debug_assert_eq!(src_buf, src.buf.as_ptr());
  70         // check InPlaceIterable contract. This is only possible if the iterator advanced the
  71         // source pointer at all. If it uses unchecked access via TrustedRandomAccess
  72         // then the source pointer will stay in its initial position and we can't use it as reference
  73         if src.ptr != src_ptr {
  74             debug_assert!(
  75                 dst as *const _ <= src.ptr,
  76                 "InPlaceIterable contract violation, write pointer advanced beyond read pointer"
  77             );
  78         }
  79
  80         // drop any remaining values at the tail of the source
  81         // but prevent drop of the allocation itself once IntoIter goes out of scope
  82         // if the drop panics then we also leak any elements collected into dst_buf
  83         src.forget_allocation_drop_remaining();
  84
  85         let vec = unsafe {
  86             let len = dst.offset_from(dst_buf) as usize;
  87             Vec::from_raw_parts(dst_buf, len, cap)
  88         };
  89
  90         vec
  91     }
  92 }
  93
  94 fn write_in_place_with_drop<T>(
  95     src_end: *const T,
  96 ) -> impl FnMut(InPlaceDrop<T>, T) -> Result<InPlaceDrop<T>, !> {
  97     move |mut sink, item| {
  98         unsafe {
  99             // the InPlaceIterable contract cannot be verified precisely here since
 100             // try_fold has an exclusive reference to the source pointer
 101             // all we can do is check if it's still in range
 102             debug_assert!(sink.dst as *const _ <= src_end, "InPlaceIterable contract violation");
 103             ptr::write(sink.dst, item);
 104             sink.dst = sink.dst.add(1);
 105         }
 106         Ok(sink)
 107     }
 108 }