Include string quotation marks in AtatLen (FactbirdHQ#137)

* Include string quotation marks in AtatLen * Formatting
tarfu · Jan 25, 2023 · a4608e9 · a4608e9
1 parent eae54e7
commit a4608e9
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 6 deletions.
diff --git a/atat/src/derive.rs b/atat/src/derive.rs
@@ -39,7 +39,7 @@ impl_length!(f32, 42);
 impl_length!(f64, 312);
 
 impl<const T: usize> AtatLen for String<T> {
-    const LEN: usize = T;
+    const LEN: usize = 1 + T + 1;
 }
 
 impl<T: AtatLen> AtatLen for Option<T> {
@@ -164,12 +164,15 @@ mod tests {
         assert_eq!(<SimpleEnum as AtatLen>::LEN, 3);
         assert_eq!(<SimpleEnumU32 as AtatLen>::LEN, 10);
         // (fields) + (n_fields - 1)
-        // (3 + 128 + 2 + 150 + 3 + 10 + 3 + (10*5)) + 7
+        // (3 + (1 + 128 + 1) + 2 + (1 + 150 + 1) + 3 + 10 + 3 + (10*5)) + 7
         assert_eq!(
             <LengthTester<'_> as AtatLen>::LEN,
-            (3 + 128 + 2 + 150 + 3 + 10 + 3) + 6
+            (3 + (1 + 128 + 1) + 2 + (1 + 150 + 1) + 3 + 10 + 3) + 6
+        );
+        assert_eq!(
+            <MixedEnum<'_> as AtatLen>::LEN,
+            (3 + 3 + (1 + 10 + 1) + 20 + 10) + 4
         );
-        assert_eq!(<MixedEnum<'_> as AtatLen>::LEN, (3 + 3 + 10 + 20 + 10) + 4);
     }
 
     #[test]

diff --git a/atat_derive/src/len.rs b/atat_derive/src/len.rs
@@ -1,6 +1,6 @@
 use crate::proc_macro::TokenStream;
 use quote::quote;
-use syn::{parse_macro_input, Ident};
+use syn::{parse_macro_input, Ident, Type};
 
 use crate::parse::{parse_field_attr, ArgAttributes, FieldAttributes, ParseInput, Variant};
 
@@ -13,7 +13,12 @@ pub fn struct_len(variants: Vec<Variant>, init_len: usize) -> proc_macro2::Token
     let mut struct_len = quote! { #init_len };
     for field in variants {
         let len = if let Some(ArgAttributes { len: Some(len), .. }) = field.attrs.at_arg {
-            quote! { #len }
+            let ty = field.ty.unwrap();
+            if is_ref_str(ty) {
+                quote! { 1 + #len + 1 }
+            } else {
+                quote! { #len }
+            }
         } else {
             let ty = field.ty.unwrap();
             quote! { <#ty as atat::AtatLen>::LEN }
@@ -25,6 +30,16 @@ pub fn struct_len(variants: Vec<Variant>, init_len: usize) -> proc_macro2::Token
     struct_len
 }
 
+fn is_ref_str(ty: Type) -> bool {
+    match ty {
+        Type::Reference(r) => match r.elem.as_ref() {
+            Type::Path(p) => p.path.segments.len() == 1 && p.path.segments[0].ident == "str",
+            _ => false,
+        },
+        _ => false,
+    }
+}
+
 /// Calculate the serialized length of an enum, as the longest of all variants
 ///
 /// Use `#[at_arg(len = xxx)]`, with a fallback to