1 files changed, 17 insertions, 17 deletions

diff --git a/library/core/src/iter/adapters/take.rs b/library/core/src/iter/adapters/take.rs
index d947c7b0e..ce18bffe7 100644
--- a/library/core/src/iter/adapters/take.rs
+++ b/library/core/src/iter/adapters/take.rs
@@ -1,5 +1,6 @@
 use crate::cmp;
 use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedLen};
+use crate::num::NonZeroUsize;
 use crate::ops::{ControlFlow, Try};
 
 /// An iterator that only iterates over the first `n` iterations of `iter`.
@@ -121,18 +122,15 @@ where
 
     #[inline]
     #[rustc_inherit_overflow_checks]
-    fn advance_by(&mut self, n: usize) -> Result<(), usize> {
+    fn advance_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
         let min = self.n.min(n);
-        match self.iter.advance_by(min) {
-            Ok(_) => {
-                self.n -= min;
-                if min < n { Err(min) } else { Ok(()) }
-            }
-            ret @ Err(advanced) => {
-                self.n -= advanced;
-                ret
-            }
-        }
+        let rem = match self.iter.advance_by(min) {
+            Ok(()) => 0,
+            Err(rem) => rem.get(),
+        };
+        let advanced = min - rem;
+        self.n -= advanced;
+        NonZeroUsize::new(n - advanced).map_or(Ok(()), Err)
     }
 }
 
@@ -223,7 +221,7 @@ where
 
     #[inline]
     #[rustc_inherit_overflow_checks]
-    fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
+    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
         // The amount by which the inner iterator needs to be shortened for it to be
         // at most as long as the take() amount.
         let trim_inner = self.iter.len().saturating_sub(self.n);
@@ -232,12 +230,14 @@ where
         // about having to advance more than usize::MAX here.
         let advance_by = trim_inner.saturating_add(n);
 
-        let advanced = match self.iter.advance_back_by(advance_by) {
-            Ok(_) => advance_by - trim_inner,
-            Err(advanced) => advanced - trim_inner,
+        let remainder = match self.iter.advance_back_by(advance_by) {
+            Ok(()) => 0,
+            Err(rem) => rem.get(),
         };
-        self.n -= advanced;
-        return if advanced < n { Err(advanced) } else { Ok(()) };
+        let advanced_by_inner = advance_by - remainder;
+        let advanced_by = advanced_by_inner - trim_inner;
+        self.n -= advanced_by;
+        NonZeroUsize::new(n - advanced_by).map_or(Ok(()), Err)
     }
 }