feat(token): Allow a byte to pass to tag

winnow-rs · Feb 6, 2024 · e61f3ee · e61f3ee
1 parent aa6e0c8
commit e61f3ee
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 0 deletions.
diff --git a/src/stream/mod.rs b/src/stream/mod.rs
@@ -454,6 +454,13 @@ impl<'a> SliceLen for &'a str {
     }
 }
 
+impl SliceLen for u8 {
+    #[inline]
+    fn slice_len(&self) -> usize {
+        1
+    }
+}
+
 impl SliceLen for char {
     #[inline(always)]
     fn slice_len(&self) -> usize {
@@ -2137,6 +2144,28 @@ impl<'a, 'b> Compare<AsciiCaseless<&'b str>> for &'a [u8] {
     }
 }
 
+impl<'a> Compare<u8> for &'a [u8] {
+    #[inline]
+    fn compare(&self, t: u8) -> CompareResult {
+        match self.first().copied() {
+            Some(c) if t == c => CompareResult::Ok(t.slice_len()),
+            Some(_) => CompareResult::Error,
+            None => CompareResult::Incomplete,
+        }
+    }
+}
+
+impl<'a> Compare<AsciiCaseless<u8>> for &'a [u8] {
+    #[inline]
+    fn compare(&self, t: AsciiCaseless<u8>) -> CompareResult {
+        match self.first() {
+            Some(c) if t.0.eq_ignore_ascii_case(c) => CompareResult::Ok(t.slice_len()),
+            Some(_) => CompareResult::Error,
+            None => CompareResult::Incomplete,
+        }
+    }
+}
+
 impl<'a> Compare<char> for &'a [u8] {
     #[inline(always)]
     fn compare(&self, t: char) -> CompareResult {
@@ -2165,6 +2194,20 @@ impl<'a, 'b> Compare<AsciiCaseless<&'b str>> for &'a str {
     }
 }
 
+impl<'a> Compare<u8> for &'a str {
+    #[inline(always)]
+    fn compare(&self, t: u8) -> CompareResult {
+        self.as_bytes().compare(t)
+    }
+}
+
+impl<'a> Compare<AsciiCaseless<u8>> for &'a str {
+    #[inline(always)]
+    fn compare(&self, t: AsciiCaseless<u8>) -> CompareResult {
+        self.as_bytes().compare(t)
+    }
+}
+
 impl<'a> Compare<char> for &'a str {
     #[inline(always)]
     fn compare(&self, t: char) -> CompareResult {

diff --git a/src/token/tests.rs b/src/token/tests.rs
@@ -204,6 +204,21 @@ fn complete_tag_char() {
     );
 }
 
+#[test]
+fn complete_tag_byte() {
+    fn test(i: &[u8]) -> IResult<&[u8], &[u8]> {
+        tag(b'B').parse_peek(i)
+    }
+    assert_eq!(test(&[0x42, 0x00][..]), Ok((&b"\x00"[..], &b"\x42"[..])));
+    assert_eq!(
+        test(&[b'A', b'\0'][..]),
+        Err(ErrMode::Backtrack(error_position!(
+            &&b"A\0"[..],
+            ErrorKind::Tag
+        )))
+    );
+}
+
 #[test]
 fn partial_any_str() {
     use super::any;