1use proc_macro2::TokenStream as TokenStream2;
2use proc_macro_error2::abort_call_site;
3use quote::quote;
4use syn::{DataEnum, Ident};
5
6use crate::construct;
7
8use super::EncodeData;
9
10pub(crate) fn encode(
11 ident: &Ident,
12 data: &DataEnum,
13 defmt_path: &syn::Path,
14) -> syn::Result<EncodeData> {
15 if data.variants.is_empty() {
16 return Ok(EncodeData {
17 stmts: vec![quote!(match *self {})],
18 format_tag: construct::interned_string("!", "derived", false, None, defmt_path),
19 where_predicates: vec![],
20 });
21 }
22
23 let mut format_string = String::new();
24 let mut where_predicates = vec![];
25
26 let mut match_arms = vec![];
27 let mut is_first_variant = true;
28 let discriminant_encoder = DiscriminantEncoder::new(data.variants.len());
29 let enum_ident = ident;
30 for (index, variant) in data.variants.iter().enumerate() {
31 let variant_ident = &variant.ident;
32
33 if is_first_variant {
34 is_first_variant = false;
35 } else {
36 format_string.push('|');
37 }
38 format_string.push_str(&variant_ident.to_string());
39
40 let mut field_patterns = vec![];
41 let (encode_fields_stmts, encode_field_where_predicates) = super::fields::codegen(
42 &variant.fields,
43 &mut format_string,
44 &mut field_patterns,
45 defmt_path,
46 )?;
47 where_predicates.extend(encode_field_where_predicates.into_iter());
48 let pattern = quote!( { #(#field_patterns),* } );
49
50 let encode_discriminant_stmt = discriminant_encoder.encode(index, defmt_path);
51
52 match_arms.push(quote!(
53 #enum_ident::#variant_ident #pattern => {
54 #encode_discriminant_stmt
55 #(#encode_fields_stmts;)*
56 }
57 ))
58 }
59
60 let format_tag = construct::interned_string(&format_string, "derived", false, None, defmt_path);
61 let stmts = vec![quote!(match self {
62 #(#match_arms)*
63 })];
64 where_predicates.dedup_by(|a, b| a == b);
65
66 Ok(EncodeData {
67 format_tag,
68 stmts,
69 where_predicates,
70 })
71}
72
73enum DiscriminantEncoder {
74 Nop,
75 U8,
76 U16,
77 U32,
78 U64,
79}
80
81impl DiscriminantEncoder {
82 fn new(number_of_variants: usize) -> Self {
83 if number_of_variants == 1 {
84 Self::Nop
85 } else if number_of_variants <= usize::from(u8::MAX) {
86 Self::U8
87 } else if number_of_variants <= usize::from(u16::MAX) {
88 Self::U16
89 } else if number_of_variants as u128 > u128::from(u64::MAX) {
90 // unreachable on existing hardware?
91 abort_call_site!(
92 "`#[derive(Format)]` does not support enums with more than {} variants",
93 number_of_variants
94 )
95 } else if number_of_variants as u64 <= u64::from(u32::MAX) {
96 Self::U32
97 } else {
98 Self::U64
99 }
100 }
101
102 // NOTE this assumes `index` < `number_of_variants` used to construct `self`
103 fn encode(&self, index: usize, defmt_path: &syn::Path) -> TokenStream2 {
104 match self {
105 // For single-variant enums, there is no need to encode the discriminant.
106 DiscriminantEncoder::Nop => quote!(),
107 DiscriminantEncoder::U8 => {
108 let index = index as u8;
109 quote!(#defmt_path::export::u8(&#index);)
110 }
111 DiscriminantEncoder::U16 => {
112 let index = index as u16;
113 quote!(#defmt_path::export::u16(&#index);)
114 }
115 DiscriminantEncoder::U32 => {
116 let index = index as u32;
117 quote!(#defmt_path::export::u32(&#index);)
118 }
119 DiscriminantEncoder::U64 => {
120 let index = index as u64;
121 quote!(#defmt_path::export::u64(&#index);)
122 }
123 }
124 }
125}
126