From dbb776d44307fd2841746cf4cb4e283cd36aa2b5 Mon Sep 17 00:00:00 2001
From: zachs18 <8355914+zachs18@users.noreply.github.com>
Date: Thu, 29 Dec 2022 18:23:14 -0600
Subject: [PATCH] Allow byte literals as enum discriminants in
 CheckedBitPattern derive macro. (#155)

---
 derive/src/traits.rs  |  1 +
 derive/tests/basic.rs | 23 +++++++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/derive/src/traits.rs b/derive/src/traits.rs
index 2590267..0f4e0a9 100644
--- a/derive/src/traits.rs
+++ b/derive/src/traits.rs
@@ -766,6 +766,7 @@ fn parse_int_expr(expr: &Expr) -> Result<i64> {
       parse_int_expr(expr).map(|int| -int)
     }
     Expr::Lit(ExprLit { lit: Lit::Int(int), .. }) => int.base10_parse(),
+    Expr::Lit(ExprLit { lit: Lit::Byte(byte), .. }) => Ok(byte.value().into()),
     _ => bail!("Not an integer expression"),
   }
 }
diff --git a/derive/tests/basic.rs b/derive/tests/basic.rs
index 25d1781..755a667 100644
--- a/derive/tests/basic.rs
+++ b/derive/tests/basic.rs
@@ -128,6 +128,16 @@ enum CheckedBitPatternEnumNonContiguous {
   E = 56,
 }
 
+#[repr(u8)]
+#[derive(Debug, Clone, Copy, NoUninit, CheckedBitPattern, PartialEq, Eq)]
+enum CheckedBitPatternEnumByteLit {
+  A = b'A',
+  B = b'B',
+  C = b'C',
+  D = b'D',
+  E = b'E',
+}
+
 #[derive(Debug, Copy, Clone, NoUninit, CheckedBitPattern, PartialEq, Eq)]
 #[repr(C)]
 struct CheckedBitPatternStruct {
@@ -187,6 +197,19 @@ fn passes_cast_noncontiguous() {
   assert_eq!(*res, CheckedBitPatternEnumNonContiguous::E);
 }
 
+#[test]
+fn fails_cast_bytelit() {
+  let can_cast = CheckedBitPatternEnumByteLit::is_valid_bit_pattern(&b'a');
+  assert!(!can_cast);
+}
+
+#[test]
+fn passes_cast_bytelit() {
+  let res =
+    bytemuck::checked::cast_slice::<u8, CheckedBitPatternEnumByteLit>(b"CAB");
+  assert_eq!(res, [CheckedBitPatternEnumByteLit::C, CheckedBitPatternEnumByteLit::A, CheckedBitPatternEnumByteLit::B]);
+}
+
 #[test]
 fn fails_cast_struct() {
   let pod = [0u8, 24u8];