2 files changed, 211 insertions, 20 deletions

diff --git a/library/core/src/iter/adapters/take.rs b/library/core/src/iter/adapters/take.rs
index ce18bffe7..c1d8cc4ff 100644
--- a/library/core/src/iter/adapters/take.rs
+++ b/library/core/src/iter/adapters/take.rs
@@ -1,5 +1,7 @@
 use crate::cmp;
-use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedLen};
+use crate::iter::{
+    adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedLen, TrustedRandomAccess,
+};
 use crate::num::NonZeroUsize;
 use crate::ops::{ControlFlow, Try};
 
@@ -98,26 +100,18 @@ where
         }
     }
 
-    impl_fold_via_try_fold! { fold -> try_fold }
-
     #[inline]
-    fn for_each<F: FnMut(Self::Item)>(mut self, f: F) {
-        // The default implementation would use a unit accumulator, so we can
-        // avoid a stateful closure by folding over the remaining number
-        // of items we wish to return instead.
-        fn check<'a, Item>(
-            mut action: impl FnMut(Item) + 'a,
-        ) -> impl FnMut(usize, Item) -> Option<usize> + 'a {
-            move |more, x| {
-                action(x);
-                more.checked_sub(1)
-            }
-        }
+    fn fold<B, F>(self, init: B, f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        Self::spec_fold(self, init, f)
+    }
 
-        let remaining = self.n;
-        if remaining > 0 {
-            self.iter.try_fold(remaining - 1, check(f));
-        }
+    #[inline]
+    fn for_each<F: FnMut(Self::Item)>(self, f: F) {
+        Self::spec_for_each(self, f)
     }
 
     #[inline]
@@ -249,3 +243,72 @@ impl<I> FusedIterator for Take<I> where I: FusedIterator {}
 
 #[unstable(feature = "trusted_len", issue = "37572")]
 unsafe impl<I: TrustedLen> TrustedLen for Take<I> {}
+
+trait SpecTake: Iterator {
+    fn spec_fold<B, F>(self, init: B, f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B;
+
+    fn spec_for_each<F: FnMut(Self::Item)>(self, f: F);
+}
+
+impl<I: Iterator> SpecTake for Take<I> {
+    #[inline]
+    default fn spec_fold<B, F>(mut self, init: B, f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        use crate::ops::NeverShortCircuit;
+        self.try_fold(init, NeverShortCircuit::wrap_mut_2(f)).0
+    }
+
+    #[inline]
+    default fn spec_for_each<F: FnMut(Self::Item)>(mut self, f: F) {
+        // The default implementation would use a unit accumulator, so we can
+        // avoid a stateful closure by folding over the remaining number
+        // of items we wish to return instead.
+        fn check<'a, Item>(
+            mut action: impl FnMut(Item) + 'a,
+        ) -> impl FnMut(usize, Item) -> Option<usize> + 'a {
+            move |more, x| {
+                action(x);
+                more.checked_sub(1)
+            }
+        }
+
+        let remaining = self.n;
+        if remaining > 0 {
+            self.iter.try_fold(remaining - 1, check(f));
+        }
+    }
+}
+
+impl<I: Iterator + TrustedRandomAccess> SpecTake for Take<I> {
+    #[inline]
+    fn spec_fold<B, F>(mut self, init: B, mut f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        let mut acc = init;
+        let end = self.n.min(self.iter.size());
+        for i in 0..end {
+            // SAFETY: i < end <= self.iter.size() and we discard the iterator at the end
+            let val = unsafe { self.iter.__iterator_get_unchecked(i) };
+            acc = f(acc, val);
+        }
+        acc
+    }
+
+    #[inline]
+    fn spec_for_each<F: FnMut(Self::Item)>(mut self, mut f: F) {
+        let end = self.n.min(self.iter.size());
+        for i in 0..end {
+            // SAFETY: i < end <= self.iter.size() and we discard the iterator at the end
+            let val = unsafe { self.iter.__iterator_get_unchecked(i) };
+            f(val);
+        }
+    }
+}
diff --git a/library/core/src/iter/range.rs b/library/core/src/iter/range.rs
index 462f7170a..0e03d0c2d 100644
--- a/library/core/src/iter/range.rs
+++ b/library/core/src/iter/range.rs
@@ -1,5 +1,7 @@
+use crate::ascii::Char as AsciiChar;
 use crate::convert::TryFrom;
 use crate::mem;
+use crate::net::{Ipv4Addr, Ipv6Addr};
 use crate::num::NonZeroUsize;
 use crate::ops::{self, Try};
 
@@ -14,7 +16,7 @@ macro_rules! unsafe_impl_trusted_step {
         unsafe impl TrustedStep for $type {}
     )*};
 }
-unsafe_impl_trusted_step![char i8 i16 i32 i64 i128 isize u8 u16 u32 u64 u128 usize];
+unsafe_impl_trusted_step![AsciiChar char i8 i16 i32 i64 i128 isize u8 u16 u32 u64 u128 usize Ipv4Addr Ipv6Addr];
 
 /// Objects that have a notion of *successor* and *predecessor* operations.
 ///
@@ -484,6 +486,112 @@ impl Step for char {
     }
 }
 
+#[unstable(feature = "step_trait", reason = "recently redesigned", issue = "42168")]
+impl Step for AsciiChar {
+    #[inline]
+    fn steps_between(&start: &AsciiChar, &end: &AsciiChar) -> Option<usize> {
+        Step::steps_between(&start.to_u8(), &end.to_u8())
+    }
+
+    #[inline]
+    fn forward_checked(start: AsciiChar, count: usize) -> Option<AsciiChar> {
+        let end = Step::forward_checked(start.to_u8(), count)?;
+        AsciiChar::from_u8(end)
+    }
+
+    #[inline]
+    fn backward_checked(start: AsciiChar, count: usize) -> Option<AsciiChar> {
+        let end = Step::backward_checked(start.to_u8(), count)?;
+
+        // SAFETY: Values below that of a valid ASCII character are also valid ASCII
+        Some(unsafe { AsciiChar::from_u8_unchecked(end) })
+    }
+
+    #[inline]
+    unsafe fn forward_unchecked(start: AsciiChar, count: usize) -> AsciiChar {
+        // SAFETY: Caller asserts that result is a valid ASCII character,
+        // and therefore it is a valid u8.
+        let end = unsafe { Step::forward_unchecked(start.to_u8(), count) };
+
+        // SAFETY: Caller asserts that result is a valid ASCII character.
+        unsafe { AsciiChar::from_u8_unchecked(end) }
+    }
+
+    #[inline]
+    unsafe fn backward_unchecked(start: AsciiChar, count: usize) -> AsciiChar {
+        // SAFETY: Caller asserts that result is a valid ASCII character,
+        // and therefore it is a valid u8.
+        let end = unsafe { Step::backward_unchecked(start.to_u8(), count) };
+
+        // SAFETY: Caller asserts that result is a valid ASCII character.
+        unsafe { AsciiChar::from_u8_unchecked(end) }
+    }
+}
+
+#[unstable(feature = "step_trait", reason = "recently redesigned", issue = "42168")]
+impl Step for Ipv4Addr {
+    #[inline]
+    fn steps_between(&start: &Ipv4Addr, &end: &Ipv4Addr) -> Option<usize> {
+        u32::steps_between(&start.to_bits(), &end.to_bits())
+    }
+
+    #[inline]
+    fn forward_checked(start: Ipv4Addr, count: usize) -> Option<Ipv4Addr> {
+        u32::forward_checked(start.to_bits(), count).map(Ipv4Addr::from_bits)
+    }
+
+    #[inline]
+    fn backward_checked(start: Ipv4Addr, count: usize) -> Option<Ipv4Addr> {
+        u32::backward_checked(start.to_bits(), count).map(Ipv4Addr::from_bits)
+    }
+
+    #[inline]
+    unsafe fn forward_unchecked(start: Ipv4Addr, count: usize) -> Ipv4Addr {
+        // SAFETY: Since u32 and Ipv4Addr are losslessly convertible,
+        //   this is as safe as the u32 version.
+        Ipv4Addr::from_bits(unsafe { u32::forward_unchecked(start.to_bits(), count) })
+    }
+
+    #[inline]
+    unsafe fn backward_unchecked(start: Ipv4Addr, count: usize) -> Ipv4Addr {
+        // SAFETY: Since u32 and Ipv4Addr are losslessly convertible,
+        //   this is as safe as the u32 version.
+        Ipv4Addr::from_bits(unsafe { u32::backward_unchecked(start.to_bits(), count) })
+    }
+}
+
+#[unstable(feature = "step_trait", reason = "recently redesigned", issue = "42168")]
+impl Step for Ipv6Addr {
+    #[inline]
+    fn steps_between(&start: &Ipv6Addr, &end: &Ipv6Addr) -> Option<usize> {
+        u128::steps_between(&start.to_bits(), &end.to_bits())
+    }
+
+    #[inline]
+    fn forward_checked(start: Ipv6Addr, count: usize) -> Option<Ipv6Addr> {
+        u128::forward_checked(start.to_bits(), count).map(Ipv6Addr::from_bits)
+    }
+
+    #[inline]
+    fn backward_checked(start: Ipv6Addr, count: usize) -> Option<Ipv6Addr> {
+        u128::backward_checked(start.to_bits(), count).map(Ipv6Addr::from_bits)
+    }
+
+    #[inline]
+    unsafe fn forward_unchecked(start: Ipv6Addr, count: usize) -> Ipv6Addr {
+        // SAFETY: Since u128 and Ipv6Addr are losslessly convertible,
+        //   this is as safe as the u128 version.
+        Ipv6Addr::from_bits(unsafe { u128::forward_unchecked(start.to_bits(), count) })
+    }
+
+    #[inline]
+    unsafe fn backward_unchecked(start: Ipv6Addr, count: usize) -> Ipv6Addr {
+        // SAFETY: Since u128 and Ipv6Addr are losslessly convertible,
+        //   this is as safe as the u128 version.
+        Ipv6Addr::from_bits(unsafe { u128::backward_unchecked(start.to_bits(), count) })
+    }
+}
+
 macro_rules! range_exact_iter_impl {
     ($($t:ty)*) => ($(
         #[stable(feature = "rust1", since = "1.0.0")]
@@ -723,6 +831,15 @@ impl<A: Step> Iterator for ops::Range<A> {
     }
 
     #[inline]
+    fn count(self) -> usize {
+        if self.start < self.end {
+            Step::steps_between(&self.start, &self.end).expect("count overflowed usize")
+        } else {
+            0
+        }
+    }
+
+    #[inline]
     fn nth(&mut self, n: usize) -> Option<A> {
         self.spec_nth(n)
     }
@@ -1120,6 +1237,17 @@ impl<A: Step> Iterator for ops::RangeInclusive<A> {
     }
 
     #[inline]
+    fn count(self) -> usize {
+        if self.is_empty() {
+            return 0;
+        }
+
+        Step::steps_between(&self.start, &self.end)
+            .and_then(|steps| steps.checked_add(1))
+            .expect("count overflowed usize")
+    }
+
+    #[inline]
     fn nth(&mut self, n: usize) -> Option<A> {
         if self.is_empty() {
             return None;