Make length prefix usable more that once

Make length prefix usable more that once

Author: oknozor

examples/get_enum_discriminant.rs

@@ -18,8 +18,8 @@ pub struct Goodbye {
     reason: String,
 }
 
-#[protocol(discriminant = "integer")]
 #[derive(Protocol, Clone, Debug, PartialEq)]
+#[protocol(discriminant = "integer")]
 #[repr(u16)]
 pub enum PacketKind {
     #[protocol(discriminator(0x00))]

protocol-derive/src/codegen/mod.rs

@@ -1,11 +1,13 @@
-pub mod enums;
-
-use crate::attr;
 use proc_macro2::TokenStream;
 use syn;
+use syn::Field;
+
+use crate::attr;
+
+pub mod enums;
 
 pub fn read_fields(fields: &syn::Fields)
-    -> TokenStream {
+                   -> TokenStream {
     match *fields {
         syn::Fields::Named(ref fields_named) => read_named_fields(fields_named),
         syn::Fields::Unnamed(ref fields_unnamed) => read_unnamed_fields(fields_unnamed),
@@ -14,7 +16,7 @@ pub fn read_fields(fields: &syn::Fields)
 }
 
 pub fn write_fields(fields: &syn::Fields)
-    -> TokenStream {
+                    -> TokenStream {
     match *fields {
         syn::Fields::Named(ref fields_named) => write_named_fields(fields_named),
         syn::Fields::Unnamed(ref fields_unnamed) => write_unnamed_fields(fields_unnamed),
@@ -32,9 +34,9 @@ fn read_named_fields(fields_named: &syn::FieldsNamed)
     let field_initializers: Vec<_> = fields_named.named.iter().map(|field| {
         let field_name = &field.ident;
         let field_ty = &field.ty;
-        // This field may store the length prefix of another field.
+        // This field may store the length prefix of one or more other field.
         let update_hints = update_hints_after_read(field, &fields_named.named);
-        let update_hints_fixed = update_hints__fixed_after_read(field, &fields_named.named);
+        let update_hints_fixed = update_hint_fixed_length(field, &fields_named.named);
 
         quote! {
             #field_name : {
@@ -53,26 +55,30 @@ fn read_named_fields(fields_named: &syn::FieldsNamed)
 fn update_hints_after_read<'a>(field: &'a syn::Field,
                                fields: impl IntoIterator<Item=&'a syn::Field> + Clone)
                                -> TokenStream {
-    if let Some((length_prefix_of, kind, prefix_subfield_names)) = length_prefix_of(field, fields.clone()) {
+    let hint_setters = length_prefix_of(field, fields.clone()).iter().map(|(length_prefix_of, kind, prefix_subfield_names)| {
         let kind = kind.path_expr();
+        quote! {
+            __hints.set_field_length(#length_prefix_of,
+                (parcel #(.#prefix_subfield_names)* ).clone() as usize,
+                #kind);
+        }
+    }).collect::<Vec<TokenStream>>();
 
+    if hint_setters.is_empty() {
+        quote! {}
+    } else {
         quote! {
             if let Ok(parcel) = res.as_ref() {
-                      __hints.set_field_length(#length_prefix_of,
-                                     (parcel #(.#prefix_subfield_names)* ).clone() as usize,
-                                     #kind);
+                #(#hint_setters);*
             }
         }
-    }else {
-        quote! { }
     }
 }
 
-fn update_hints__fixed_after_read<'a>(field: &'a syn::Field,
-                               fields: impl IntoIterator<Item=&'a syn::Field> + Clone)
-                               -> TokenStream {
-
-     if let Some(attr::Protocol::FixedLength(length)) = attr::protocol(&field.attrs) {
+fn update_hint_fixed_length<'a>(field: &'a syn::Field,
+                                fields: impl IntoIterator<Item=&'a syn::Field> + Clone)
+                                -> TokenStream {
+    if let Some(attr::Protocol::FixedLength(length)) = attr::protocol(&field.attrs) {
         let position = fields.clone().into_iter().position(|f| f == field).unwrap();
 
         quote! {
@@ -85,20 +91,26 @@ fn update_hints__fixed_after_read<'a>(field: &'a syn::Field,
 
 fn update_hints_after_write<'a>(field: &'a syn::Field,
                                 fields: impl IntoIterator<Item=&'a syn::Field> + Clone)
-    -> TokenStream {
-    if let Some((length_prefix_of, kind, prefix_subfield_names)) = length_prefix_of(field, fields.clone()) {
+                                -> TokenStream {
+    let hint_setters = length_prefix_of(field, fields.clone()).iter().map(|(length_prefix_of, kind, prefix_subfield_names)| {
         let field_name = &field.ident;
         let kind = kind.path_expr();
 
         quote! {
-            if let Ok(()) = res {
                 __hints.set_field_length(#length_prefix_of,
                                          (self.#field_name #(.#prefix_subfield_names)* ).clone() as usize,
                                          #kind);
-            }
         }
+    }).collect::<Vec<TokenStream>>();
+
+    if hint_setters.is_empty() {
+        quote! {}
     } else {
-        quote! { }
+        quote! {
+            if let Ok(()) = res {
+                #(#hint_setters);*
+            }
+        }
     }
 }
 
@@ -110,33 +122,33 @@ fn update_hints_after_write<'a>(field: &'a syn::Field,
 /// Returns the field index of the field whose length is specified.
 fn length_prefix_of<'a>(field: &'a syn::Field,
                         fields: impl IntoIterator<Item=&'a syn::Field> + Clone)
-    -> Option<(usize, attr::LengthPrefixKind, Vec<syn::Ident>)> {
+                        -> Vec<(usize, attr::LengthPrefixKind, Vec<syn::Ident>)> {
     let potential_prefix = field.ident.as_ref();
 
-    let prefix_of = fields.clone().into_iter().find(|potential_prefix_of| {
+    let prefixes_of: Vec<&'a Field> = fields.clone().into_iter().filter(|potential_prefix_of| {
         match attr::protocol(&potential_prefix_of.attrs) {
             Some(attr::Protocol::LengthPrefix { ref prefix_field_name, .. }) => {
                 if !fields.clone().into_iter().any(|f| f.ident.as_ref() == Some(prefix_field_name)) {
                     panic!("length prefix is invalid: there is no sibling field named '{}", prefix_field_name);
                 }
 
                 potential_prefix == Some(prefix_field_name)
-            },
+            }
             _ => false,
         }
-    });
-
-    if let Some(prefix_of) = prefix_of {
-        let prefix_of_index = fields.clone().into_iter().position(|f| f == prefix_of).unwrap();
-        match attr::protocol(&prefix_of.attrs).unwrap() {
-            attr::Protocol::LengthPrefix { kind, prefix_subfield_names, .. } => {
-                Some((prefix_of_index, kind.clone(), prefix_subfield_names))
-            },
-            _ => unreachable!(),
-        }
-    } else {
-        None
-    }
+    }).collect();
+
+
+    prefixes_of.iter()
+        .map(|prefix_of| {
+            let prefix_of_index = fields.clone().into_iter().position(|f| &f == prefix_of).unwrap();
+            match attr::protocol(&prefix_of.attrs).unwrap() {
+                attr::Protocol::LengthPrefix { kind, prefix_subfield_names, .. } => {
+                    (prefix_of_index, kind.clone(), prefix_subfield_names)
+                }
+                _ => unreachable!(),
+            }
+        }).collect()
 }
 
 fn write_named_fields(fields_named: &syn::FieldsNamed)
@@ -145,7 +157,7 @@ fn write_named_fields(fields_named: &syn::FieldsNamed)
         let field_name = &field.ident;
         // This field may store the length prefix of another field.
         let update_hints = update_hints_after_write(field, &fields_named.named);
-        let update_hints_fixed = update_hints__fixed_after_read(field, &fields_named.named);
+        let update_hints_fixed = update_hint_fixed_length(field, &fields_named.named);
 
         quote! {
             {
@@ -162,7 +174,7 @@ fn write_named_fields(fields_named: &syn::FieldsNamed)
 }
 
 fn read_unnamed_fields(fields_unnamed: &syn::FieldsUnnamed)
-    -> TokenStream {
+                       -> TokenStream {
     let field_initializers: Vec<_> = fields_unnamed.unnamed.iter().map(|field| {
         let field_ty = &field.ty;
 
@@ -179,7 +191,7 @@ fn read_unnamed_fields(fields_unnamed: &syn::FieldsUnnamed)
 }
 
 fn write_unnamed_fields(fields_unnamed: &syn::FieldsUnnamed)
-    -> TokenStream {
+                        -> TokenStream {
     let field_indices = (0..fields_unnamed.unnamed.len()).into_iter().map(syn::Index::from);
 
     let field_writers: Vec<_> = field_indices.map(|field_index| {

protocol/src/hint.rs

@@ -41,7 +41,10 @@ impl Hints {
     /// Gets the length of the field currently being
     /// read, if known.
     pub fn current_field_length(&self) -> Option<FieldLength> {
-        self.current_field_index.and_then(|index| self.known_field_lengths.get(&index)).cloned()
+        println!("Getting current field length");
+        let maybe_length = self.current_field_index.and_then(|index| self.known_field_lengths.get(&index)).cloned();
+        println!("Got {:?}", maybe_length);
+        maybe_length
     }
 }
 

protocol/src/types/vec.rs

@@ -49,6 +49,7 @@ mod std_vec {
         fn read_field(read: &mut dyn Read,
                       settings: &Settings,
                       hints: &mut hint::Hints) -> Result<Self, Error> {
+            println!("Vec hints {:?}", hints);
             util::read_list(read, settings, hints)
         }
 

protocol/src/util.rs

@@ -53,6 +53,7 @@ pub fn read_list<T>(read: &mut dyn Read,
                     hints: &mut hint::Hints)
     -> Result<Vec<T>, Error>
     where T: Parcel {
+    println!("read predixed list  with hints ->  {:?}", hints);
     self::read_list_ext::<SizeType, T>(read, settings, hints)
 }
 
@@ -102,6 +103,7 @@ pub fn read_list_ext<S,T>(read: &mut dyn Read,
                     Ok(items)
                 },
                 hint::LengthPrefixKind::Elements => {
+                    println!("Reading \"Element Kind\" with length : {:?}", length.length);
                     read_items(length.length, read, settings).map(|i| i.collect())
                 },
             }

tests/src/length_prefix.rs

@@ -29,13 +29,33 @@ pub struct WithElementsLength {
     pub data: Vec<u32>,
 }
 
+#[derive(Protocol, Debug, PartialEq, Eq)]
+pub struct WithTwoPrefix {
+    pub count: u32,
+    pub count_2: u32,
+    pub foo: bool,
+    #[protocol(length_prefix(elements(count)))]
+    pub data: Vec<u32>,
+    #[protocol(length_prefix(elements(count_2)))]
+    pub data_2: Vec<u32>,
+}
+
+#[derive(Protocol, Debug, PartialEq, Eq)]
+pub struct WithElementsLengthUsedTwice {
+    pub count: u32,
+    pub foo: bool,
+    #[protocol(length_prefix(elements(count)))]
+    pub data_1: Vec<u32>,
+    #[protocol(length_prefix(elements(count)))]
+    pub data_2: Vec<u32>,
+}
+
 #[derive(Protocol, Debug, PartialEq, Eq)]
 pub struct WithFixedLength {
     #[protocol(fixed_length(3))]
     pub data: Vec<u32>,
 }
 
-
 #[test]
 fn can_read_length_prefix_5_bytes_string() {
     assert_eq!(Foo {
@@ -70,6 +90,49 @@ fn can_read_length_prefix_3_elements() {
                              &Settings::default()).unwrap());
 }
 
+#[test]
+fn can_read_twice_the_same_prefix_length() {
+    assert_eq!(WithElementsLengthUsedTwice {
+        count: 3,
+        foo: true,
+        data_1: vec![1, 2, 3],
+        data_2: vec![4, 5, 6],
+    }, WithElementsLengthUsedTwice::from_raw_bytes(
+        &[0, 0, 0, 3, // disjoint length prefix
+            1, // boolean true
+            0, 0, 0, 1, // 1
+            0, 0, 0, 2, // 2
+            0, 0, 0, 3, // 3
+            0, 0, 0, 4, // 4
+            0, 0, 0, 5, // 5
+            0, 0, 0, 6  // 6
+        ],
+        &Settings::default()).unwrap());
+}
+
+#[test]
+fn can_read_two_prefix_length() {
+    assert_eq!(WithTwoPrefix {
+        count: 3,
+        count_2: 3,
+        foo: true,
+        data: vec![1, 2, 3],
+        data_2: vec![4, 5, 6],
+    }, WithTwoPrefix::from_raw_bytes(
+        &[
+            0, 0, 0, 3, // disjoint length prefix
+            0, 0, 0, 3, // disjoint length prefix
+            1, // boolean true
+            0, 0, 0, 1, // 1
+            0, 0, 0, 2, // 2
+            0, 0, 0, 3, // 3
+            0, 0, 0, 4, // 4
+            0, 0, 0, 5, // 5
+            0, 0, 0, 6  // 6
+        ],
+        &Settings::default()).unwrap());
+}
+
 #[test]
 fn can_read_fixed_length_prefix() {