Rollup merge of #33849 - ranma42:escape-iters-count, r=alexcrichton

Implement `count` for `EscapeUnicode`

and cleanup the code for `count` for `EscapeDefault` (instead of repeating the `match` for `size_hint` and `count`).

This PR marks EscapeUnicode and EscapeDefault as ExactSizeIterator. The constraints for the trait implementations held even before this PR, but I am not sure if this is something we want to guarantee/expose (I would love feedback on this, especially on what would be the appropriate way to handle stabilisation, if needed).

Part of #24214, split from #31049.

The test for `count` was added in #33103.

2 files changed, 56 insertions, 28 deletions

diff --git a/src/libcore/char.rs b/src/libcore/char.rs
index 6a2331dddcf..d80b456181a 100644
--- a/src/libcore/char.rs
+++ b/src/libcore/char.rs
@@ -411,14 +411,17 @@ pub struct EscapeUnicode {
     hex_digit_idx: usize,
 }
 
+// The enum values are ordered so that their representation is the
+// same as the remaining length (besides the hexadecimal digits). This
+// likely makes `len()` a single load from memory) and inline-worth.
 #[derive(Clone, Debug)]
 enum EscapeUnicodeState {
-    Backslash,
-    Type,
-    LeftBrace,
-    Value,
-    RightBrace,
     Done,
+    RightBrace,
+    Value,
+    LeftBrace,
+    Type,
+    Backslash,
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -457,19 +460,17 @@ impl Iterator for EscapeUnicode {
         }
     }
 
+    #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
-        let n = match self.state {
-            EscapeUnicodeState::Backslash => 5,
-            EscapeUnicodeState::Type => 4,
-            EscapeUnicodeState::LeftBrace => 3,
-            EscapeUnicodeState::Value => 2,
-            EscapeUnicodeState::RightBrace => 1,
-            EscapeUnicodeState::Done => 0,
-        };
-        let n = n + self.hex_digit_idx;
+        let n = self.len();
         (n, Some(n))
     }
 
+    #[inline]
+    fn count(self) -> usize {
+        self.len()
+    }
+
     fn last(self) -> Option<char> {
         match self.state {
             EscapeUnicodeState::Done => None,
@@ -483,6 +484,22 @@ impl Iterator for EscapeUnicode {
     }
 }
 
+#[stable(feature = "exact_size_escape", since = "1.11.0")]
+impl ExactSizeIterator for EscapeUnicode {
+    #[inline]
+    fn len(&self) -> usize {
+        // The match is a single memory access with no branching
+        self.hex_digit_idx + match self.state {
+            EscapeUnicodeState::Done => 0,
+            EscapeUnicodeState::RightBrace => 1,
+            EscapeUnicodeState::Value => 2,
+            EscapeUnicodeState::LeftBrace => 3,
+            EscapeUnicodeState::Type => 4,
+            EscapeUnicodeState::Backslash => 5,
+        }
+    }
+}
+
 /// An iterator that yields the literal escape code of a `char`.
 ///
 /// This `struct` is created by the [`escape_default()`] method on [`char`]. See
@@ -498,9 +515,9 @@ pub struct EscapeDefault {
 
 #[derive(Clone, Debug)]
 enum EscapeDefaultState {
-    Backslash(char),
-    Char(char),
     Done,
+    Char(char),
+    Backslash(char),
     Unicode(EscapeUnicode),
 }
 
@@ -523,22 +540,15 @@ impl Iterator for EscapeDefault {
         }
     }
 
+    #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
-        match self.state {
-            EscapeDefaultState::Char(_) => (1, Some(1)),
-            EscapeDefaultState::Backslash(_) => (2, Some(2)),
-            EscapeDefaultState::Unicode(ref iter) => iter.size_hint(),
-            EscapeDefaultState::Done => (0, Some(0)),
-        }
+        let n = self.len();
+        (n, Some(n))
     }
 
+    #[inline]
     fn count(self) -> usize {
-        match self.state {
-            EscapeDefaultState::Char(_) => 1,
-            EscapeDefaultState::Unicode(iter) => iter.count(),
-            EscapeDefaultState::Done => 0,
-            EscapeDefaultState::Backslash(_) => 2,
-        }
+        self.len()
     }
 
     fn nth(&mut self, n: usize) -> Option<char> {
@@ -578,6 +588,18 @@ impl Iterator for EscapeDefault {
     }
 }
 
+#[stable(feature = "exact_size_escape", since = "1.11.0")]
+impl ExactSizeIterator for EscapeDefault {
+    fn len(&self) -> usize {
+        match self.state {
+            EscapeDefaultState::Done => 0,
+            EscapeDefaultState::Char(_) => 1,
+            EscapeDefaultState::Backslash(_) => 2,
+            EscapeDefaultState::Unicode(ref iter) => iter.len(),
+        }
+    }
+}
+
 /// An iterator over `u8` entries represending the UTF-8 encoding of a `char`
 /// value.
 ///
diff --git a/src/libcoretest/char.rs b/src/libcoretest/char.rs
index e959e71daf7..7da876b9459 100644
--- a/src/libcoretest/char.rs
+++ b/src/libcoretest/char.rs
@@ -276,6 +276,12 @@ fn eu_iterator_specializations() {
             // Check last
             assert_eq!(iter.clone().last(), Some('}'));
 
+            // Check len
+            assert_eq!(iter.len(), len - offset);
+
+            // Check size_hint (= len in ExactSizeIterator)
+            assert_eq!(iter.size_hint(), (iter.len(), Some(iter.len())));
+
             // Check counting
             assert_eq!(iter.clone().count(), len - offset);