library/alloc/src/vec/splice.rs

   1 use crate::alloc::{Allocator, Global};
   2 use core::ptr::{self};
   3 use core::slice::{self};
   4
   5 use super::{Drain, Vec};
   6
   7 /// A splicing iterator for `Vec`.
   8 ///
   9 /// This struct is created by [`Vec::splice()`].
  10 /// See its documentation for more.
  11 ///
  12 /// # Example
  13 ///
  14 /// ```
  15 /// let mut v = vec![0, 1, 2];
  16 /// let new = [7, 8];
  17 /// let iter: std::vec::Splice<_> = v.splice(1.., new);
  18 /// ```
  19 #[derive(Debug)]
  20 #[stable(feature = "vec_splice", since = "1.21.0")]
  21 pub struct Splice<
  22     'a,
  23     I: Iterator + 'a,
  24     #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator + 'a = Global,
  25 > {
  26     pub(super) drain: Drain<'a, I::Item, A>,
  27     pub(super) replace_with: I,
  28 }
  29
  30 #[stable(feature = "vec_splice", since = "1.21.0")]
  31 impl<I: Iterator, A: Allocator> Iterator for Splice<'_, I, A> {
  32     type Item = I::Item;
  33
  34     fn next(&mut self) -> Option<Self::Item> {
  35         self.drain.next()
  36     }
  37
  38     fn size_hint(&self) -> (usize, Option<usize>) {
  39         self.drain.size_hint()
  40     }
  41 }
  42
  43 #[stable(feature = "vec_splice", since = "1.21.0")]
  44 impl<I: Iterator, A: Allocator> DoubleEndedIterator for Splice<'_, I, A> {
  45     fn next_back(&mut self) -> Option<Self::Item> {
  46         self.drain.next_back()
  47     }
  48 }
  49
  50 #[stable(feature = "vec_splice", since = "1.21.0")]
  51 impl<I: Iterator, A: Allocator> ExactSizeIterator for Splice<'_, I, A> {}
  52
  53 #[stable(feature = "vec_splice", since = "1.21.0")]
  54 impl<I: Iterator, A: Allocator> Drop for Splice<'_, I, A> {
  55     fn drop(&mut self) {
  56         self.drain.by_ref().for_each(drop);
  57         // At this point draining is done and the only remaining tasks are splicing
  58         // and moving things into the final place.
  59         // Which means we can replace the slice::Iter with pointers that won't point to deallocated
  60         // memory, so that Drain::drop is still allowed to call iter.len(), otherwise it would break
  61         // the ptr.sub_ptr contract.
  62         self.drain.iter = (&[]).iter();
  63
  64         unsafe {
  65             if self.drain.tail_len == 0 {
  66                 self.drain.vec.as_mut().extend(self.replace_with.by_ref());
  67                 return;
  68             }
  69
  70             // First fill the range left by drain().
  71             if !self.drain.fill(&mut self.replace_with) {
  72                 return;
  73             }
  74
  75             // There may be more elements. Use the lower bound as an estimate.
  76             // FIXME: Is the upper bound a better guess? Or something else?
  77             let (lower_bound, _upper_bound) = self.replace_with.size_hint();
  78             if lower_bound > 0 {
  79                 self.drain.move_tail(lower_bound);
  80                 if !self.drain.fill(&mut self.replace_with) {
  81                     return;
  82                 }
  83             }
  84
  85             // Collect any remaining elements.
  86             // This is a zero-length vector which does not allocate if `lower_bound` was exact.
  87             let mut collected = self.replace_with.by_ref().collect::<Vec<I::Item>>().into_iter();
  88             // Now we have an exact count.
  89             if collected.len() > 0 {
  90                 self.drain.move_tail(collected.len());
  91                 let filled = self.drain.fill(&mut collected);
  92                 debug_assert!(filled);
  93                 debug_assert_eq!(collected.len(), 0);
  94             }
  95         }
  96         // Let `Drain::drop` move the tail back if necessary and restore `vec.len`.
  97     }
  98 }
  99
 100 /// Private helper methods for `Splice::drop`
 101 impl<T, A: Allocator> Drain<'_, T, A> {
 102     /// The range from `self.vec.len` to `self.tail_start` contains elements
 103     /// that have been moved out.
 104     /// Fill that range as much as possible with new elements from the `replace_with` iterator.
 105     /// Returns `true` if we filled the entire range. (`replace_with.next()` didn’t return `None`.)
 106     unsafe fn fill<I: Iterator<Item = T>>(&mut self, replace_with: &mut I) -> bool {
 107         let vec = unsafe { self.vec.as_mut() };
 108         let range_start = vec.len;
 109         let range_end = self.tail_start;
 110         let range_slice = unsafe {
 111             slice::from_raw_parts_mut(vec.as_mut_ptr().add(range_start), range_end - range_start)
 112         };
 113
 114         for place in range_slice {
 115             if let Some(new_item) = replace_with.next() {
 116                 unsafe { ptr::write(place, new_item) };
 117                 vec.len += 1;
 118             } else {
 119                 return false;
 120             }
 121         }
 122         true
 123     }
 124
 125     /// Makes room for inserting more elements before the tail.
 126     unsafe fn move_tail(&mut self, additional: usize) {
 127         let vec = unsafe { self.vec.as_mut() };
 128         let len = self.tail_start + self.tail_len;
 129         vec.buf.reserve(len, additional);
 130
 131         let new_tail_start = self.tail_start + additional;
 132         unsafe {
 133             let src = vec.as_ptr().add(self.tail_start);
 134             let dst = vec.as_mut_ptr().add(new_tail_start);
 135             ptr::copy(src, dst, self.tail_len);
 136         }
 137         self.tail_start = new_tail_start;
 138     }
 139 }