1 use crate::alloc::Allocator; 2 use core::iter::TrustedLen; 3 use core::ptr::{self}; 4 use core::slice::{self}; 5 6 use super::{IntoIter, SetLenOnDrop, Vec}; 7 8 // Specialization trait used for Vec::extend 9 pub(super) trait SpecExtend<T, I> { 10 fn spec_extend(&mut self, iter: I); 11 } 12 13 impl<T, I, A: Allocator> SpecExtend<T, I> for Vec<T, A> 14 where 15 I: Iterator<Item = T>, 16 { 17 default fn spec_extend(&mut self, iter: I) { 18 self.extend_desugared(iter) 19 } 20 } 21 22 impl<T, I, A: Allocator> SpecExtend<T, I> for Vec<T, A> 23 where 24 I: TrustedLen<Item = T>, 25 { 26 default fn spec_extend(&mut self, iterator: I) { 27 // This is the case for a TrustedLen iterator. 28 let (low, high) = iterator.size_hint(); 29 if let Some(additional) = high { 30 debug_assert_eq!( 31 low, 32 additional, 33 "TrustedLen iterator's size hint is not exact: {:?}", 34 (low, high) 35 ); 36 self.reserve(additional); 37 unsafe { 38 let mut ptr = self.as_mut_ptr().add(self.len()); 39 let mut local_len = SetLenOnDrop::new(&mut self.len); 40 iterator.for_each(move |element| { 41 ptr::write(ptr, element); 42 ptr = ptr.offset(1); 43 // Since the loop executes user code which can panic we have to bump the pointer 44 // after each step. 45 // NB can't overflow since we would have had to alloc the address space 46 local_len.increment_len(1); 47 }); 48 } 49 } else { 50 // Per TrustedLen contract a `None` upper bound means that the iterator length 51 // truly exceeds usize::MAX, which would eventually lead to a capacity overflow anyway. 52 // Since the other branch already panics eagerly (via `reserve()`) we do the same here. 53 // This avoids additional codegen for a fallback code path which would eventually 54 // panic anyway. 55 panic!("capacity overflow"); 56 } 57 } 58 } 59 60 impl<T, A: Allocator> SpecExtend<T, IntoIter<T>> for Vec<T, A> { 61 fn spec_extend(&mut self, mut iterator: IntoIter<T>) { 62 unsafe { 63 self.append_elements(iterator.as_slice() as _); 64 } 65 iterator.forget_remaining_elements(); 66 } 67 } 68 69 impl<'a, T: 'a, I, A: Allocator + 'a> SpecExtend<&'a T, I> for Vec<T, A> 70 where 71 I: Iterator<Item = &'a T>, 72 T: Clone, 73 { 74 default fn spec_extend(&mut self, iterator: I) { 75 self.spec_extend(iterator.cloned()) 76 } 77 } 78 79 impl<'a, T: 'a, A: Allocator + 'a> SpecExtend<&'a T, slice::Iter<'a, T>> for Vec<T, A> 80 where 81 T: Copy, 82 { 83 fn spec_extend(&mut self, iterator: slice::Iter<'a, T>) { 84 let slice = iterator.as_slice(); 85 unsafe { self.append_elements(slice) }; 86 } 87 } 88