Handle raw variant names correctly

rkyv · Sep 8, 2024 · e309d01 · e309d01
1 parent a53150d
commit e309d01
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 6 deletions.
diff --git a/bytecheck_derive/src/lib.rs b/bytecheck_derive/src/lib.rs
@@ -23,7 +23,7 @@ use syn::{
 use crate::{
     attributes::{Attributes, FieldAttributes},
     repr::Repr,
-    util::iter_fields,
+    util::{iter_fields, strip_raw},
 };
 
 /// Derives `CheckBytes` for the labeled type.
@@ -321,8 +321,10 @@ fn derive_check_bytes(mut input: DeriveInput) -> Result<TokenStream, Error> {
 
             let variant_structs = data.variants.iter().map(|v| {
                 let variant = &v.ident;
-                let variant_name =
-                    Ident::new(&format!("Variant{variant}"), v.span());
+                let variant_name = Ident::new(
+                    &format!("Variant{}", strip_raw(variant)),
+                    v.span(),
+                );
                 match v.fields {
                     Fields::Named(ref fields) => {
                         let fields = fields.named.iter().map(|f| {
@@ -365,8 +367,10 @@ fn derive_check_bytes(mut input: DeriveInput) -> Result<TokenStream, Error> {
 
             let check_arms = data.variants.iter().map(|v| {
                 let variant = &v.ident;
-                let variant_name =
-                    Ident::new(&format!("Variant{variant}"), v.span());
+                let variant_name = Ident::new(
+                    &format!("Variant{}", strip_raw(variant)),
+                    v.span(),
+                );
                 match v.fields {
                     Fields::Named(ref fields) => {
                         let checks = fields.named.iter().map(|f| {

diff --git a/bytecheck_derive/src/util.rs b/bytecheck_derive/src/util.rs
@@ -1,7 +1,8 @@
 use core::iter::FlatMap;
 
 use syn::{
-    punctuated::Iter, Data, DataEnum, DataStruct, DataUnion, Field, Variant,
+    punctuated::Iter, Data, DataEnum, DataStruct, DataUnion, Field, Ident,
+    Variant,
 };
 
 type VariantFieldsFn = fn(&Variant) -> Iter<'_, Field>;
@@ -39,3 +40,11 @@ pub fn iter_fields(data: &Data) -> FieldsIter<'_> {
         }
     }
 }
+
+pub fn strip_raw(ident: &Ident) -> String {
+    let as_string = ident.to_string();
+    as_string
+        .strip_prefix("r#")
+        .map(ToString::to_string)
+        .unwrap_or(as_string)
+}