Allow byte literals as enum discriminants in CheckedBitPattern derive…

… macro. (#155)
Lokathor · Dec 30, 2022 · dbb776d · dbb776d
1 parent a758c09
commit dbb776d
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git 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
@@ -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];