diff --git a/library/core/src/iter/adapters/take.rs b/library/core/src/iter/adapters/take.rs index 2962e0104d11d..fcace2499f12f 100644 --- a/library/core/src/iter/adapters/take.rs +++ b/library/core/src/iter/adapters/take.rs @@ -1,5 +1,9 @@ use crate::cmp; -use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedLen}; +use crate::iter::adapters::zip::try_get_unchecked; +use crate::iter::{ + adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedLen, TrustedRandomAccess, + TrustedRandomAccessNoCoerce, +}; use crate::ops::{ControlFlow, Try}; /// An iterator that only iterates over the first `n` iterations of `iter`. @@ -56,20 +60,7 @@ where #[inline] fn size_hint(&self) -> (usize, Option) { - if self.n == 0 { - return (0, Some(0)); - } - - let (lower, upper) = self.iter.size_hint(); - - let lower = cmp::min(lower, self.n); - - let upper = match upper { - Some(x) if x < self.n => Some(x), - _ => Some(self.n), - }; - - (lower, upper) + TakeSizeHint::size_hint(self) } #[inline] @@ -127,6 +118,16 @@ where } } } + + #[inline] + unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> Self::Item + where + Self: TrustedRandomAccessNoCoerce, + { + // SAFETY: the caller must uphold the contract for + // `Iterator::__iterator_get_unchecked`. + unsafe { try_get_unchecked(&mut self.iter, idx) } + } } #[unstable(issue = "none", feature = "inplace_iteration")] @@ -242,3 +243,48 @@ impl FusedIterator for Take where I: FusedIterator {} #[unstable(feature = "trusted_len", issue = "37572")] unsafe impl TrustedLen for Take {} + +#[doc(hidden)] +#[unstable(feature = "trusted_random_access", issue = "none")] +unsafe impl TrustedRandomAccess for Take where I: TrustedRandomAccess {} + +#[doc(hidden)] +#[unstable(feature = "trusted_random_access", issue = "none")] +unsafe impl TrustedRandomAccessNoCoerce for Take +where + I: TrustedRandomAccessNoCoerce, +{ + const MAY_HAVE_SIDE_EFFECT: bool = I::MAY_HAVE_SIDE_EFFECT; +} + +trait TakeSizeHint: Iterator { + fn size_hint(&self) -> (usize, Option); +} + +impl TakeSizeHint for Take { + #[inline] + default fn size_hint(&self) -> (usize, Option) { + if self.n == 0 { + return (0, Some(0)); + } + + let (lower, upper) = self.iter.size_hint(); + + let lower = cmp::min(lower, self.n); + + let upper = match upper { + Some(x) if x < self.n => Some(x), + _ => Some(self.n), + }; + + (lower, upper) + } +} + +impl TakeSizeHint for Take { + #[inline] + fn size_hint(&self) -> (usize, Option) { + let size = cmp::min(self.iter.size_hint().0, self.n); + (size, Some(size)) + } +}