rust/alloc/vec/spec_extend.rs

   1 // SPDX-License-Identifier: Apache-2.0 OR MIT
   2
   3 use crate::alloc::Allocator;
   4 use crate::collections::{TryReserveError, TryReserveErrorKind};
   5 use core::iter::TrustedLen;
   6 use core::ptr::{self};
   7 use core::slice::{self};
   8
   9 use super::{IntoIter, SetLenOnDrop, Vec};
  10
  11 // Specialization trait used for Vec::extend
  12 #[cfg(not(no_global_oom_handling))]
  13 pub(super) trait SpecExtend<T, I> {
  14     fn spec_extend(&mut self, iter: I);
  15 }
  16
  17 // Specialization trait used for Vec::try_extend
  18 pub(super) trait TrySpecExtend<T, I> {
  19     fn try_spec_extend(&mut self, iter: I) -> Result<(), TryReserveError>;
  20 }
  21
  22 #[cfg(not(no_global_oom_handling))]
  23 impl<T, I, A: Allocator> SpecExtend<T, I> for Vec<T, A>
  24 where
  25     I: Iterator<Item = T>,
  26 {
  27     default fn spec_extend(&mut self, iter: I) {
  28         self.extend_desugared(iter)
  29     }
  30 }
  31
  32 impl<T, I, A: Allocator> TrySpecExtend<T, I> for Vec<T, A>
  33 where
  34     I: Iterator<Item = T>,
  35 {
  36     default fn try_spec_extend(&mut self, iter: I) -> Result<(), TryReserveError> {
  37         self.try_extend_desugared(iter)
  38     }
  39 }
  40
  41 #[cfg(not(no_global_oom_handling))]
  42 impl<T, I, A: Allocator> SpecExtend<T, I> for Vec<T, A>
  43 where
  44     I: TrustedLen<Item = T>,
  45 {
  46     default fn spec_extend(&mut self, iterator: I) {
  47         // This is the case for a TrustedLen iterator.
  48         let (low, high) = iterator.size_hint();
  49         if let Some(additional) = high {
  50             debug_assert_eq!(
  51                 low,
  52                 additional,
  53                 "TrustedLen iterator's size hint is not exact: {:?}",
  54                 (low, high)
  55             );
  56             self.reserve(additional);
  57             unsafe {
  58                 let mut ptr = self.as_mut_ptr().add(self.len());
  59                 let mut local_len = SetLenOnDrop::new(&mut self.len);
  60                 iterator.for_each(move |element| {
  61                     ptr::write(ptr, element);
  62                     ptr = ptr.offset(1);
  63                     // Since the loop executes user code which can panic we have to bump the pointer
  64                     // after each step.
  65                     // NB can't overflow since we would have had to alloc the address space
  66                     local_len.increment_len(1);
  67                 });
  68             }
  69         } else {
  70             // Per TrustedLen contract a `None` upper bound means that the iterator length
  71             // truly exceeds usize::MAX, which would eventually lead to a capacity overflow anyway.
  72             // Since the other branch already panics eagerly (via `reserve()`) we do the same here.
  73             // This avoids additional codegen for a fallback code path which would eventually
  74             // panic anyway.
  75             panic!("capacity overflow");
  76         }
  77     }
  78 }
  79
  80 impl<T, I, A: Allocator> TrySpecExtend<T, I> for Vec<T, A>
  81 where
  82     I: TrustedLen<Item = T>,
  83 {
  84     default fn try_spec_extend(&mut self, iterator: I) -> Result<(), TryReserveError> {
  85         // This is the case for a TrustedLen iterator.
  86         let (low, high) = iterator.size_hint();
  87         if let Some(additional) = high {
  88             debug_assert_eq!(
  89                 low,
  90                 additional,
  91                 "TrustedLen iterator's size hint is not exact: {:?}",
  92                 (low, high)
  93             );
  94             self.try_reserve(additional)?;
  95             unsafe {
  96                 let mut ptr = self.as_mut_ptr().add(self.len());
  97                 let mut local_len = SetLenOnDrop::new(&mut self.len);
  98                 iterator.for_each(move |element| {
  99                     ptr::write(ptr, element);
 100                     ptr = ptr.offset(1);
 101                     // Since the loop executes user code which can panic we have to bump the pointer
 102                     // after each step.
 103                     // NB can't overflow since we would have had to alloc the address space
 104                     local_len.increment_len(1);
 105                 });
 106             }
 107             Ok(())
 108         } else {
 109             Err(TryReserveErrorKind::CapacityOverflow.into())
 110         }
 111     }
 112 }
 113
 114 #[cfg(not(no_global_oom_handling))]
 115 impl<T, A: Allocator> SpecExtend<T, IntoIter<T>> for Vec<T, A> {
 116     fn spec_extend(&mut self, mut iterator: IntoIter<T>) {
 117         unsafe {
 118             self.append_elements(iterator.as_slice() as _);
 119         }
 120         iterator.forget_remaining_elements();
 121     }
 122 }
 123
 124 impl<T, A: Allocator> TrySpecExtend<T, IntoIter<T>> for Vec<T, A> {
 125     fn try_spec_extend(&mut self, mut iterator: IntoIter<T>) -> Result<(), TryReserveError> {
 126         unsafe {
 127             self.try_append_elements(iterator.as_slice() as _)?;
 128         }
 129         iterator.forget_remaining_elements();
 130         Ok(())
 131     }
 132 }
 133
 134 #[cfg(not(no_global_oom_handling))]
 135 impl<'a, T: 'a, I, A: Allocator + 'a> SpecExtend<&'a T, I> for Vec<T, A>
 136 where
 137     I: Iterator<Item = &'a T>,
 138     T: Clone,
 139 {
 140     default fn spec_extend(&mut self, iterator: I) {
 141         self.spec_extend(iterator.cloned())
 142     }
 143 }
 144
 145 impl<'a, T: 'a, I, A: Allocator + 'a> TrySpecExtend<&'a T, I> for Vec<T, A>
 146 where
 147     I: Iterator<Item = &'a T>,
 148     T: Clone,
 149 {
 150     default fn try_spec_extend(&mut self, iterator: I) -> Result<(), TryReserveError> {
 151         self.try_spec_extend(iterator.cloned())
 152     }
 153 }
 154
 155 #[cfg(not(no_global_oom_handling))]
 156 impl<'a, T: 'a, A: Allocator + 'a> SpecExtend<&'a T, slice::Iter<'a, T>> for Vec<T, A>
 157 where
 158     T: Copy,
 159 {
 160     fn spec_extend(&mut self, iterator: slice::Iter<'a, T>) {
 161         let slice = iterator.as_slice();
 162         unsafe { self.append_elements(slice) };
 163     }
 164 }
 165
 166 impl<'a, T: 'a, A: Allocator + 'a> TrySpecExtend<&'a T, slice::Iter<'a, T>> for Vec<T, A>
 167 where
 168     T: Copy,
 169 {
 170     fn try_spec_extend(&mut self, iterator: slice::Iter<'a, T>) -> Result<(), TryReserveError> {
 171         let slice = iterator.as_slice();
 172         unsafe { self.try_append_elements(slice) }
 173     }
 174 }