refactor: pattern match on discrimintant consts

Author: omfj

soavec_derive/src/soa_enum.rs

@@ -4,11 +4,67 @@
 
 use proc_macro2::TokenStream;
 use quote::quote;
-use syn::{Data, DeriveInput, Fields, Expr, Lit, spanned::Spanned};
+use syn::{Data, DeriveInput, Expr, Fields, spanned::Spanned};
+
+/// Parse the repr type from #[repr(...)] attribute
+fn parse_repr_type(input: &DeriveInput) -> TokenStream {
+    for attr in &input.attrs {
+        if attr.path().is_ident("repr") {
+            // Parse as a single type first (e.g., #[repr(u8)])
+            if let Ok(ty) = attr.parse_args::<syn::Type>() {
+                if is_valid_discriminant_type(&ty) {
+                    return quote! { #ty };
+                }
+            } else {
+                // Parse as list of types (for example #[repr(C, u8)])
+                if let Ok(meta_list) = attr.parse_args_with(
+                    syn::punctuated::Punctuated::<syn::Meta, syn::Token![,]>::parse_terminated,
+                ) {
+                    for meta in meta_list {
+                        if let syn::Meta::Path(path) = meta {
+                            let ty: syn::Type = syn::Type::Path(syn::TypePath {
+                                qself: None,
+                                path: path.clone(),
+                            });
+                            if is_valid_discriminant_type(&ty) {
+                                return quote! { #ty };
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    // Fallback if no valid repr type found
+    quote! { isize }
+}
+
+/// Check if a type is a valid discriminant type
+fn is_valid_discriminant_type(ty: &syn::Type) -> bool {
+    if let syn::Type::Path(type_path) = ty
+        && let Some(ident) = type_path.path.get_ident()
+    {
+        let ident_str = ident.to_string();
+        return matches!(ident_str.as_str(), |"u8"| "i8"
+            | "u16"
+            | "i16"
+            | "u32"
+            | "i32"
+            | "u64"
+            | "i64"
+            | "u128"
+            | "i128"
+            | "usize"
+            | "isize");
+    }
+    false
+}
 
 pub fn expand_data_enum(input: DeriveInput) -> syn::Result<TokenStream> {
     let enum_name = &input.ident;
     let generics = &input.generics;
+    let discriminant_type = parse_repr_type(&input);
 
     let variants = match &input.data {
         Data::Enum(data_enum) => &data_enum.variants,
@@ -34,11 +90,12 @@ pub fn expand_data_enum(input: DeriveInput) -> syn::Result<TokenStream> {
 
     // Information about the variants
     let mut variant_data = Vec::new();
-    let mut current_discriminant: Expr = syn::parse_quote!(0); 
+    let mut current_discriminant: Expr = syn::parse_quote!(0);
     for variant in variants.iter() {
         let variant_name = &variant.ident;
 
-        let discriminant_value = variant.discriminant
+        let discriminant_value = variant
+            .discriminant
             .as_ref()
             .map(|(_, expr)| expr.clone())
             .unwrap_or_else(|| current_discriminant.clone());
@@ -220,14 +277,28 @@ pub fn expand_data_enum(input: DeriveInput) -> syn::Result<TokenStream> {
         .collect();
 
     let tuple_repr = if union_names.is_empty() {
-        quote! { (u8,) }
+        quote! { (#discriminant_type,) }
     } else {
-        quote! { (u8, #(#union_names),*) }
+        quote! { (#discriminant_type, #(#union_names),*) }
     };
 
+    // Generate const items for discriminant values
+    // This is so that we can later pattern match on the discriminant values
+    let discriminant_consts: Vec<_> = variant_data
+        .iter()
+        .map(|(variant_name, disc_value, _, _)| {
+            let const_name =
+                quote::format_ident!("__DISC_{}", variant_name.to_string().to_uppercase());
+            quote! {
+                const #const_name: #discriminant_type = #disc_value;
+            }
+        })
+        .collect();
+
     // Generate into_tuple match arms
-    let into_tuple_arms: Vec<_> = variant_data.iter().map(|(variant_name, disc_value, field_names, _field_types)| {
+    let into_tuple_arms: Vec<_> = variant_data.iter().map(|(variant_name, _disc_value, field_names, _field_types)| {
         let variant_name_lower = variant_name.to_string().to_lowercase();
+        let const_name = quote::format_ident!("__DISC_{}", variant_name.to_string().to_uppercase());
 
         if field_names.is_empty() {
             // Unit variant
@@ -239,16 +310,15 @@ pub fn expand_data_enum(input: DeriveInput) -> syn::Result<TokenStream> {
 
             if union_constructions.is_empty() {
                 quote! {
-                    Self::#variant_name => (#disc_value,)
+                    Self::#variant_name => (#const_name,)
                 }
             } else {
                 quote! {
-                    Self::#variant_name => (#disc_value, #(#union_constructions),*)
+                    Self::#variant_name => (#const_name, #(#union_constructions),*)
                 }
             }
-        }
-        else {
-        // Variant with fields
+        } else {
+            // Variant with fields
             let union_constructions: Vec<_> = (0..max_fields).map(|idx| {
                 let union_name = quote::format_ident!("{}Field{}", enum_name, idx);
                 let union_field = quote::format_ident!("{}", variant_name_lower);
@@ -261,21 +331,23 @@ pub fn expand_data_enum(input: DeriveInput) -> syn::Result<TokenStream> {
             }).collect();
 
             quote! {
-                Self::#variant_name(#(#field_names),*) => (#disc_value, #(#union_constructions),*)
+                Self::#variant_name(#(#field_names),*) => (#const_name, #(#union_constructions),*)
             }
         }
     }).collect();
 
     // Generate from_tuple match arms
     let from_tuple_arms: Vec<_> = variant_data
         .iter()
-        .map(|(variant_name, disc_value, field_names, _field_types)| {
+        .map(|(variant_name, _disc_value, field_names, _field_types)| {
             let variant_name_lower = variant_name.to_string().to_lowercase();
+            let const_name =
+                quote::format_ident!("__DISC_{}", variant_name.to_string().to_uppercase());
 
             if field_names.is_empty() {
                 // Unit variant
                 quote! {
-                    #disc_value => Self::#variant_name
+                    #const_name => Self::#variant_name
                 }
             } else {
                 // Variant with fields - extract from unions
@@ -291,7 +363,7 @@ pub fn expand_data_enum(input: DeriveInput) -> syn::Result<TokenStream> {
                     .collect();
 
                 quote! {
-                    #disc_value => {
+                    #const_name => {
                         #(#field_extractions)*
                         Self::#variant_name(#(#field_names),*)
                     }
@@ -310,6 +382,8 @@ pub fn expand_data_enum(input: DeriveInput) -> syn::Result<TokenStream> {
         .collect();
 
     let expanded = quote! {
+        #(#discriminant_consts)*
+
         #(#union_types)*
 
         #[allow(dead_code)]
@@ -378,7 +452,7 @@ pub fn expand_data_enum(input: DeriveInput) -> syn::Result<TokenStream> {
                 let (__disc_ptr, #(#union_field_names),*) = value;
                 unsafe {
                     #ref_struct_name {
-                        discriminant: core::mem::transmute(__disc_ptr.as_ptr()),
+                        discriminant: &*(__disc_ptr.as_ptr() as *const _),
                         #(#ref_union_field_names: #union_field_names.as_ref()),*
                     }
                 }
@@ -391,7 +465,7 @@ pub fn expand_data_enum(input: DeriveInput) -> syn::Result<TokenStream> {
                 let (__disc_ptr, #(mut #union_field_names),*) = value;
                 unsafe {
                     #mut_struct_name {
-                        discriminant: core::mem::transmute(__disc_ptr.as_ptr()),
+                        discriminant: &*(__disc_ptr.as_ptr() as *const _),
                         #(#ref_union_field_names: #union_field_names.as_mut()),*
                     }
                 }
@@ -406,7 +480,7 @@ pub fn expand_data_enum(input: DeriveInput) -> syn::Result<TokenStream> {
                 let (__disc_ptr, #(#union_field_names),*) = value;
                 unsafe {
                     #slice_struct_name {
-                        discriminant: core::slice::from_raw_parts(core::mem::transmute(__disc_ptr.as_ptr()), __soa_len),
+                        discriminant: core::slice::from_raw_parts(__disc_ptr.as_ptr() as *const _, __soa_len),
                         #(#ref_union_field_names: core::slice::from_raw_parts(#union_field_names.as_ptr(), __soa_len)),*
                     }
                 }
@@ -421,7 +495,7 @@ pub fn expand_data_enum(input: DeriveInput) -> syn::Result<TokenStream> {
                 let (__disc_ptr, #(#union_field_names),*) = value;
                 unsafe {
                     #slice_mut_struct_name {
-                        discriminant: core::slice::from_raw_parts_mut(core::mem::transmute(__disc_ptr.as_ptr()), __soa_len),
+                        discriminant: core::slice::from_raw_parts_mut(__disc_ptr.as_ptr() as *mut _, __soa_len),
                         #(#ref_union_field_names: core::slice::from_raw_parts_mut(#union_field_names.as_ptr(), __soa_len)),*
                     }
                 }
@@ -439,6 +513,7 @@ mod tests {
     #[test]
     fn test_expand_derive_enum() {
         let input: DeriveInput = syn::parse_quote! {
+            #[repr(u8)]
             enum TestEnum {
                 Undefined,
                 Null,
@@ -449,7 +524,10 @@ mod tests {
 
         let result = expand_data_enum(input).unwrap().to_string();
 
-        // std::fs::write("test_expand_derive_enum_output.rs", &result).unwrap();
+        assert!(result.contains("const __DISC_UNDEFINED : u8 = 0 ;"));
+        assert!(result.contains("const __DISC_NULL : u8 = 0 + 1 ;"));
+        assert!(result.contains("const __DISC_BOOLEAN : u8 = 0 + 1 + 1 ;"));
+        assert!(result.contains("const __DISC_NUMBER : u8 = 0 + 1 + 1 + 1 ;"));
 
         assert!(result.contains("union TestEnumField0"));
         assert!(result.contains("undefined : ()"));
@@ -480,8 +558,6 @@ mod tests {
 
         let result = expand_data_enum(input).unwrap().to_string();
 
-        // std::fs::write("test_expand_derive_enum_explicit_discriminant.rs", &result).unwrap();
-
         assert!(result.contains("union ValueField0"));
         assert!(result.contains("undefined : ()"));
         assert!(result.contains("string : & 'static str"));
@@ -541,8 +617,6 @@ mod tests {
 
         let result = expand_data_enum(input).unwrap().to_string();
 
-        // std::fs::write("test_multiple_field_variants_output.rs", &result).unwrap();
-
         assert!(result.contains("union MultiFieldField0"));
         assert!(result.contains("union MultiFieldField1"));
         assert!(result.contains("union MultiFieldField2"));
@@ -553,4 +627,88 @@ mod tests {
         assert!(result.contains("String"));
         assert!(result.contains("i32"));
     }
+
+    #[test]
+    fn test_respects_repr() {
+        let input: DeriveInput = syn::parse_quote! {
+            #[repr(u16)]
+            enum CustomRepr {
+                A,
+                B,
+                C,
+            }
+        };
+
+        let result = expand_data_enum(input).unwrap().to_string();
+
+        assert!(result.contains("const __DISC_A : u16 = 0 ;"));
+        assert!(result.contains("const __DISC_B : u16 = 0 + 1 ;"));
+        assert!(result.contains("const __DISC_C : u16 = 0 + 1 + 1 ;"));
+    }
+
+    #[test]
+    fn test_respects_custom_discriminants() {
+        let input: DeriveInput = syn::parse_quote! {
+            enum CustomDiscriminants {
+                A = 10,
+                B = 20,
+                C, // 20 + 1
+            }
+        };
+
+        let result = expand_data_enum(input).unwrap().to_string();
+
+        assert!(result.contains("const __DISC_A : isize = 10 ;"));
+        assert!(result.contains("const __DISC_B : isize = 20 ;"));
+        assert!(result.contains("const __DISC_C : isize = 20 + 1 ;"));
+    }
+
+    #[test]
+    fn test_function_const_discriminants() {
+        let input: DeriveInput = syn::parse_quote! {
+            enum FuncConstDiscriminants {
+                A = MY_CONST,
+                B = my_function(),
+                C, // my_function() + 1
+            }
+        };
+
+        let result = expand_data_enum(input).unwrap().to_string();
+
+        assert!(result.contains("const __DISC_A : isize = MY_CONST ;"));
+        assert!(result.contains("const __DISC_B : isize = my_function () ;"));
+        assert!(result.contains("const __DISC_C : isize = my_function () + 1 ;"));
+    }
+
+    #[test]
+    fn test_does_not_use_c_as_type() {
+        let input: DeriveInput = syn::parse_quote! {
+            #[repr(C)] // should default to isize
+            enum NotC {
+                A,
+                B,
+            }
+        };
+
+        let result = expand_data_enum(input).unwrap().to_string();
+
+        assert!(result.contains("const __DISC_A : isize = 0 ;"));
+        assert!(result.contains("const __DISC_B : isize = 0 + 1 ;"));
+    }
+
+    #[test]
+    fn test_repr_c_u8_uses_u8() {
+        let input: DeriveInput = syn::parse_quote! {
+            #[repr(C, u8)]
+            enum CAndU8 {
+                A,
+                B,
+            }
+        };
+
+        let result = expand_data_enum(input).unwrap().to_string();
+
+        assert!(result.contains("const __DISC_A : u8 = 0 ;"));
+        assert!(result.contains("const __DISC_B : u8 = 0 + 1 ;"));
+    }
 }