Repository: 3Hren/msgpack-rust
Pull request #264 (2021-01-11 08:49:06)

Serializing seqs and maps of unknown length (#196) 
This fixes the serialization of structs that contain fields tagged with #[serde(flatten)].
The basic idea is to write elements to another Serializer, count the elements that are
serialized, and in the end, write the counted length and the serialzed elements to the
main Serializer.

A drawback might be that such a struct cannot be written in compact representation as
serde wants to serialize it as a map.

kornelski: That looks good. 

The addition of io::Error to the Error enum is a breaking change (breaks `match` expressions without catch-all). So you could either remove the unknown length too, or maybe change the io::Error to some of the old cases instead to keep back-compat.

Patch:
diff --git a/rmp-serde/CHANGELOG.md b/rmp-serde/CHANGELOG.md
index 8d1f23d..1ecd20d 100644
--- a/rmp-serde/CHANGELOG.md
+++ b/rmp-serde/CHANGELOG.md
@@ -10,6 +10,7 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 - Add missing `Debug` implementations.
 - More `decode::Error` conversions.
 - Support for serializing and deserializing 128-bit values in serde.
+- Support for serializing sequences and maps with unknown length, that enables the use of `#[serde(flatten)]` attribute (#196).
 
 ### Changed:
 - (Breaking) Serialize newtype structs by serializing its inner type without wrapping into a tuple.
diff --git a/rmp-serde/src/encode.rs b/rmp-serde/src/encode.rs
index 1d5de54..12a159e 100644
--- a/rmp-serde/src/encode.rs
+++ b/rmp-serde/src/encode.rs
@@ -9,7 +9,7 @@ use serde::Serialize;
 use serde::ser::{SerializeMap, SerializeSeq, SerializeStruct, SerializeStructVariant,
                  SerializeTuple, SerializeTupleStruct, SerializeTupleVariant};
 
-use rmp::encode;
+use rmp::{Marker, encode};
 use rmp::encode::ValueWriteError;
 
 use crate::config::{
@@ -24,6 +24,7 @@ use crate::MSGPACK_EXT_STRUCT_NAME;
 pub enum Error {
     /// Failed to write a MessagePack value.
     InvalidValueWrite(ValueWriteError),
+    //TODO: This can be removed at some point
     /// Failed to serialize struct, sequence or map, because its length is unknown.
     UnknownLength,
     /// Invalid Data model, i.e. Serialize trait is not implmented correctly
@@ -192,6 +193,24 @@ impl<'a, W: Write + 'a, C> Serializer<W, C> {
     }
 }
 
+impl<'a, W: Write + 'a, C: SerializerConfig> Serializer<W, C> {
+    #[inline]
+    fn maybe_unknown_len_compound<F>(&'a mut self, len: Option<usize>, f: F) -> Result<MaybeUnknownLengthCompound<'a, W, C>, Error>
+    where F: Fn(&mut W, u32) -> Result<Marker, ValueWriteError>
+    {
+        Ok(MaybeUnknownLengthCompound {
+            compound: match len {
+                Some(len) => {
+                    f(&mut self.wr, len as u32)?;
+                    None
+                }
+                None => Some(UnknownLengthCompound::from(&*self))
+            },
+            se: self,
+        })
+    }
+}
+
 impl<W: Write, C> Serializer<W, C> {
     /// Consumes this serializer returning the new one, which will serialize structs as a map.
     ///
@@ -354,12 +373,14 @@ impl<'a, W: Write + 'a, C: SerializerConfig> SerializeTupleStruct for Compound<'
     }
 }
 
-impl<'a, W: Write + 'a, C: SerializerConfig> SerializeTupleVariant for Compound<'a, W, C> {
+impl<'a, W: Write + 'a, C: SerializerConfig> SerializeStruct for Compound<'a, W, C> {
     type Ok = ();
     type Error = Error;
 
-    fn serialize_field<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), Self::Error> {
-        value.serialize(&mut *self.se)
+    fn serialize_field<T: ?Sized + Serialize>(&mut self, key: &'static str, value: &T) ->
+        Result<(), Self::Error>
+    {
+        C::write_struct_field(&mut *self.se, key, value)
     }
 
     fn end(self) -> Result<Self::Ok, Self::Error> {
@@ -367,15 +388,11 @@ impl<'a, W: Write + 'a, C: SerializerConfig> SerializeTupleVariant for Compound<
     }
 }
 
-impl<'a, W: Write + 'a, C: SerializerConfig> SerializeMap for Compound<'a, W, C> {
+impl<'a, W: Write + 'a, C: SerializerConfig> SerializeTupleVariant for Compound<'a, W, C> {
     type Ok = ();
     type Error = Error;
 
-    fn serialize_key<T: ?Sized + Serialize>(&mut self, key: &T) -> Result<(), Self::Error> {
-        key.serialize(&mut *self.se)
-    }
-
-    fn serialize_value<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), Self::Error> {
+    fn serialize_field<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), Self::Error> {
         value.serialize(&mut *self.se)
     }
 
@@ -384,7 +401,7 @@ impl<'a, W: Write + 'a, C: SerializerConfig> SerializeMap for Compound<'a, W, C>
     }
 }
 
-impl<'a, W: Write + 'a, C: SerializerConfig> SerializeStruct for Compound<'a, W, C> {
+impl<'a, W: Write + 'a, C: SerializerConfig> SerializeStructVariant for Compound<'a, W, C> {
     type Ok = ();
     type Error = Error;
 
@@ -399,17 +416,86 @@ impl<'a, W: Write + 'a, C: SerializerConfig> SerializeStruct for Compound<'a, W,
     }
 }
 
-impl<'a, W: Write + 'a, C: SerializerConfig> SerializeStructVariant for Compound<'a, W, C> {
+/// Contains a `Serializer` for sequences and maps whose length is not yet known
+/// and a counter for the number of elements that are encoded by the `Serializer`.
+#[derive(Debug)]
+struct UnknownLengthCompound<C> {
+    se: Serializer<Vec<u8>, C>,
+    elem_count: u32
+}
+impl<W, C: SerializerConfig> From<&Serializer<W, C>> for UnknownLengthCompound<C> {
+    fn from(se: &Serializer<W, C>) -> Self {
+        Self {
+            se: Serializer { wr: Vec::with_capacity(128), config: se.config, depth: se.depth },
+            elem_count: 0
+        }
+    }
+}
+
+/// Contains a `Serializer` for encoding elements of sequences and maps.
+///
+/// # Note
+///
+/// If , for example, a field inside a struct is tagged with `#serde(flatten)` the total number of
+/// fields of this struct will be unknown to serde because flattened fields may have name clashes
+/// and then will be overwritten. So, serde wants to serialize the struct as a map with an unknown
+/// length.
+///
+/// For the described case a `UnknownLengthCompound` is used to encode the elements. On `end()`
+/// the counted length and the encoded elements will be written to the `Serializer`. A caveat is,
+/// that structs that contain flattened fields arem always written as a map, even when compact
+/// representaion is desired.
+///
+/// Otherwise, if the length is known, the elements will be encoded directly by the `Serializer`.
+#[derive(Debug)]
+pub struct MaybeUnknownLengthCompound<'a, W: 'a, C: 'a> {
+    se: &'a mut Serializer<W, C>,
+    compound: Option<UnknownLengthCompound<C>>,
+}
+
+impl<'a, W: Write + 'a, C: SerializerConfig> SerializeSeq for MaybeUnknownLengthCompound<'a, W, C> {
     type Ok = ();
     type Error = Error;
 
-    fn serialize_field<T: ?Sized + Serialize>(&mut self, key: &'static str, value: &T) ->
-        Result<(), Self::Error>
-    {
-        C::write_struct_field(&mut *self.se, key, value)
+    fn serialize_element<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), Self::Error> {
+        match self.compound.as_mut() {
+            None => value.serialize(&mut *self.se),
+            Some(buf) =>  {
+                value.serialize(&mut buf.se)?;
+                buf.elem_count += 1;
+                Ok(())
+            }
+        }
+    }
+
+    fn end(self) -> Result<Self::Ok, Self::Error> {
+        if let Some(compound) = self.compound {
+            encode::write_array_len(&mut self.se.wr, compound.elem_count)?;
+            self.se.wr.write_all(&compound.se.into_inner())
+                .map_err(ValueWriteError::InvalidDataWrite)?;
+        }
+        Ok(())
+    }
+}
+
+impl<'a, W: Write + 'a, C: SerializerConfig> SerializeMap for MaybeUnknownLengthCompound<'a, W, C> {
+    type Ok = ();
+    type Error = Error;
+
+    fn serialize_key<T: ?Sized + Serialize>(&mut self, key: &T) -> Result<(), Self::Error> {
+        <Self as SerializeSeq>::serialize_element(self, key)
+    }
+
+    fn serialize_value<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), Self::Error> {
+        <Self as SerializeSeq>::serialize_element(self, value)
     }
 
     fn end(self) -> Result<Self::Ok, Self::Error> {
+        if let Some(compound) = self.compound {
+            encode::write_map_len(&mut self.se.wr, compound.elem_count / 2)?;
+            self.se.wr.write_all(&compound.se.into_inner())
+                .map_err(ValueWriteError::InvalidDataWrite)?;
+        }
         Ok(())
     }
 }
@@ -422,11 +508,11 @@ where
     type Ok = ();
     type Error = Error;
 
-    type SerializeSeq = Compound<'a, W, C>;
+    type SerializeSeq = MaybeUnknownLengthCompound<'a, W, C>;
     type SerializeTuple = Compound<'a, W, C>;
     type SerializeTupleStruct = Compound<'a, W, C>;
     type SerializeTupleVariant = Compound<'a, W, C>;
-    type SerializeMap = Compound<'a, W, C>;
+    type SerializeMap = MaybeUnknownLengthCompound<'a, W, C>;
     type SerializeStruct = Compound<'a, W, C>;
     type SerializeStructVariant = Compound<'a, W, C>;
 
@@ -558,20 +644,16 @@ where
     }
 
     fn serialize_seq(self, len: Option<usize>) -> Result<Self::SerializeSeq, Error> {
-        let len = match len {
-            Some(len) => len,
-            None => return Err(Error::UnknownLength),
-        };
+        self.maybe_unknown_len_compound(len, |wr, len| encode::write_array_len(wr, len))
+    }
 
+    //TODO: normal compund
+    fn serialize_tuple(self, len: usize) -> Result<Self::SerializeTuple, Self::Error> {
         encode::write_array_len(&mut self.wr, len as u32)?;
 
         self.compound()
     }
 
-    fn serialize_tuple(self, len: usize) -> Result<Self::SerializeTuple, Self::Error> {
-        self.serialize_seq(Some(len))
-    }
-
     fn serialize_tuple_struct(self, _name: &'static str, len: usize) ->
         Result<Self::SerializeTupleStruct, Self::Error>
     {
@@ -590,13 +672,7 @@ where
     }
 
     fn serialize_map(self, len: Option<usize>) -> Result<Self::SerializeMap, Error> {
-        match len {
-            Some(len) => {
-                encode::write_map_len(&mut self.wr, len as u32)?;
-                self.compound()
-            }
-            None => Err(Error::UnknownLength),
-        }
+        self.maybe_unknown_len_compound(len, |wr, len| encode::write_map_len(wr, len))
     }
 
     fn serialize_struct(self, _name: &'static str, len: usize) ->


Test patch:
diff --git a/rmp-serde/tests/decode_derive.rs b/rmp-serde/tests/decode_derive.rs
index 3b8919f..6cbf6fd 100644
--- a/rmp-serde/tests/decode_derive.rs
+++ b/rmp-serde/tests/decode_derive.rs
@@ -383,6 +383,111 @@ fn pass_enum_with_one_arg() {
     assert_eq!(buf.len() as u64, de.get_ref().position())
 }
 
+#[test]
+fn pass_struct_with_flattened_map_field() {
+    use std::collections::BTreeMap;
+
+    // The encoded bytearray is: { "f1": 0, "f2": { "german": "Hallo Welt!" }, "english": "Hello World!" }.
+    let buf = [
+        0x83, 0xA2, 0x66, 0x31, 0x00, 0xA2, 0x66, 0x32, 0x81, 0xA6, 0x67, 0x65, 0x72, 0x6D, 0x61, 0x6E, 0xAB,
+        0x48, 0x61, 0x6C, 0x6C, 0x6F, 0x20, 0x57, 0x65, 0x6C, 0x74, 0x21, 0xA7, 0x65, 0x6E, 0x67, 0x6C, 0x69,
+        0x73, 0x68, 0xAC, 0x48, 0x65, 0x6C, 0x6C, 0x6F, 0x20, 0x57, 0x6F, 0x72, 0x6C, 0x64, 0x21,
+    ];
+    let cur = Cursor::new(&buf[..]);
+
+    #[derive(Debug, PartialEq, Deserialize)]
+    struct Struct {
+        f1: u32,
+        // not flattend!
+        f2: BTreeMap<String, String>,
+        #[serde(flatten)]
+        f3: BTreeMap<String, String>
+    }
+
+    let expected = Struct {
+        f1: 0,
+        f2: {
+            let mut map = BTreeMap::new();
+            map.insert("german".to_string(), "Hallo Welt!".to_string());
+            map
+        },
+        f3: {
+            let mut map = BTreeMap::new();
+            map.insert("english".to_string(), "Hello World!".to_string());
+            map
+        }
+    };
+
+    let mut de = Deserializer::new(cur);
+    let actual: Struct = Deserialize::deserialize(&mut de).unwrap();
+
+    assert_eq!(expected, actual);
+    assert_eq!(buf.len() as u64, de.get_ref().position());
+}
+
+#[test]
+fn pass_struct_with_flattened_struct_field() {
+    #[derive(Debug, PartialEq, Deserialize)]
+    struct Struct {
+        f1: u32,
+        // not flattend!
+        f2: InnerStruct,
+        #[serde(flatten)]
+        f3: InnerStruct
+    }
+
+    #[derive(Debug, PartialEq, Deserialize)]
+    struct InnerStruct {
+        f4: u32,
+        f5: u32
+    }
+
+    let expected = Struct {
+        f1: 0,
+        f2: InnerStruct {
+            f4: 8,
+            f5: 13
+        },
+        f3: InnerStruct {
+            f4: 21,
+            f5: 34
+        }
+    };
+
+    // struct-as-tuple
+    {
+        // The encoded bytearray is: { "f1": 0, "f2": [8, 13], "f4": 21, "f5": 34 }.
+        let buf = [
+            0x84, 0xA2, 0x66, 0x31, 0x00, 0xA2, 0x66, 0x32, 0x92, 0x08, 0x0D, 0xA2, 0x66, 0x34, 0x15, 0xA2, 0x66, 0x35,
+            0x22,
+        ];
+        let cur = Cursor::new(&buf[..]);
+
+        let mut de = Deserializer::new(cur);
+        let actual: Struct = Deserialize::deserialize(&mut de).unwrap();
+
+        assert_eq!(expected, actual);
+        assert_eq!(buf.len() as u64, de.get_ref().position());
+    }
+
+    // struct-as-map
+    {
+        // The encoded bytearray is: { "f1": 0, "f2": { "f4": 8, "f5": 13 }, "f4": 21, "f5": 34 }.
+        let buf = [
+            0x84, 0xA2, 0x66, 0x31, 0x00, 0xA2, 0x66, 0x32, 0x82, 0xA2, 0x66, 0x34, 0x08,
+            0xA2, 0x66, 0x35, 0x0D, 0xA2, 0x66, 0x34, 0x15, 0xA2, 0x66, 0x35, 0x22,
+        ];
+        let cur = Cursor::new(&buf[..]);
+
+        let mut de = Deserializer::new(cur);
+        let actual: Struct = Deserialize::deserialize(&mut de).unwrap();
+
+        assert_eq!(expected, actual);
+        assert_eq!(buf.len() as u64, de.get_ref().position());
+    }
+}
+
+
 #[test]
 fn pass_from_slice() {
     let buf = [0x93, 0xa4, 0x4a, 0x6f, 0x68, 0x6e, 0xa5, 0x53, 0x6d, 0x69, 0x74, 0x68, 0x2a];
diff --git a/rmp-serde/tests/encode_derive.rs b/rmp-serde/tests/encode_derive.rs
index 9dd4de9..985e5c8 100644
--- a/rmp-serde/tests/encode_derive.rs
+++ b/rmp-serde/tests/encode_derive.rs
@@ -166,6 +166,103 @@ fn serialize_struct_variant_as_map() {
     assert_eq!(vec![0x81, 0x00, 0x81, 0xa2, 0x66, 0x31, 0x2a], se.into_inner());
 }
 
+#[test]
+fn serialize_struct_with_flattened_map_field() {
+    use std::collections::BTreeMap;
+
+    #[derive(Serialize)]
+    struct Struct {
+        f1: u32,
+        // not flattend!
+        f2: BTreeMap<String, String>,
+        #[serde(flatten)]
+        f3: BTreeMap<String, String>
+    }
+
+    let mut se = Serializer::new(Vec::new());
+    Struct {
+        f1: 0,
+        f2: {
+            let mut map = BTreeMap::new();
+            map.insert("german".to_string(), "Hallo Welt!".to_string());
+            map
+        },
+        f3: {
+            let mut map = BTreeMap::new();
+            map.insert("english".to_string(), "Hello World!".to_string());
+            map
+        }
+    }
+    .serialize(&mut se).unwrap();
+
+    // Expect: { "f1": 0, "f2": { "german": "Hallo Welt!" }, "english": "Hello World!" }.
+    assert_eq!(
+        vec![
+            0x83, 0xA2, 0x66, 0x31, 0x00, 0xA2, 0x66, 0x32, 0x81, 0xA6, 0x67, 0x65, 0x72, 0x6D, 0x61, 0x6E, 0xAB,
+            0x48, 0x61, 0x6C, 0x6C, 0x6F, 0x20, 0x57, 0x65, 0x6C, 0x74, 0x21, 0xA7, 0x65, 0x6E, 0x67, 0x6C, 0x69,
+            0x73, 0x68, 0xAC, 0x48, 0x65, 0x6C, 0x6C, 0x6F, 0x20, 0x57, 0x6F, 0x72, 0x6C, 0x64, 0x21,
+        ],
+        se.into_inner()
+    );
+}
+
+#[test]
+fn serialize_struct_with_flattened_struct_field() {
+    #[derive(Serialize)]
+    struct Struct {
+        f1: u32,
+        // not flattend!
+        f2: InnerStruct,
+        #[serde(flatten)]
+        f3: InnerStruct
+    }
+
+    #[derive(Serialize)]
+    struct InnerStruct {
+        f4: u32,
+        f5: u32
+    }
+
+    let strct = Struct {
+        f1: 0,
+        f2: InnerStruct {
+            f4: 8,
+            f5: 13
+        },
+        f3: InnerStruct {
+            f4: 21,
+            f5: 34
+        }
+    };
+
+    // struct-as-tuple
+    {
+        let mut se = Serializer::new(Vec::new());
+        strct.serialize(&mut se).unwrap();
+
+        // Expect: { "f1": 0, "f2": [8, 13], "f4": 21, "f5": 34 }.
+        assert_eq!(vec![
+            0x84, 0xA2, 0x66, 0x31, 0x00, 0xA2, 0x66, 0x32, 0x92, 0x08, 0x0D, 0xA2, 0x66, 0x34, 0x15, 0xA2, 0x66, 0x35, 0x22,
+        ],
+        se.into_inner());
+    }
+
+    // struct-as-map
+    {
+        let mut se = Serializer::new(Vec::new()).with_struct_map();
+        strct.serialize(&mut se).unwrap();
+
+        // Expect: { "f1": 0, "f2": { "f4": 8, "f5": 13 }, "f4": 21, "f5": 34 }.
+        assert_eq!(
+            vec![
+                0x84, 0xA2, 0x66, 0x31, 0x00, 0xA2, 0x66, 0x32, 0x82, 0xA2, 0x66, 0x34, 0x08,
+                0xA2, 0x66, 0x35, 0x0D, 0xA2, 0x66, 0x34, 0x15, 0xA2, 0x66, 0x35, 0x22,
+            ],
+            se.into_inner()
+        );
+    }
+}
+
 #[test]
 fn pass_struct_as_map_using_ext() {
     #[derive(Serialize)]
diff --git a/rmp-serde/tests/round.rs b/rmp-serde/tests/round.rs
index 67cd263..dd7fdc4 100644
--- a/rmp-serde/tests/round.rs
+++ b/rmp-serde/tests/round.rs
@@ -332,6 +332,82 @@ fn round_trip_unit_struct_untagged_enum() {
     }
 }
 
+#[test]
+fn round_trip_struct_with_flattened_map_field() {
+    use std::collections::BTreeMap;
+
+    #[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
+    struct Struct {
+        f1: u32,
+        // not flattend!
+        f2: BTreeMap<String, String>,
+        #[serde(flatten)]
+        f3: BTreeMap<String, String>
+    }
+
+    let strct = Struct {
+        f1: 0,
+        f2: {
+            let mut map = BTreeMap::new();
+            map.insert("german".to_string(), "Hallo Welt!".to_string());
+            map
+        },
+        f3: {
+            let mut map = BTreeMap::new();
+            map.insert("english".to_string(), "Hello World!".to_string());
+            map
+        }
+    };
+
+    let serialized: Vec<u8> = rmps::to_vec(&strct).unwrap();
+    let deserialized: Struct = rmps::from_slice(&serialized).unwrap();
+    assert_eq!(deserialized, strct);
+}
+
+#[test]
+fn round_trip_struct_with_flattened_struct_field() {
+    #[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
+    struct Struct {
+        f1: u32,
+        // not flattend!
+        f2: InnerStruct,
+        #[serde(flatten)]
+        f3: InnerStruct
+    }
+
+    #[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
+    struct InnerStruct {
+        f4: u32,
+        f5: u32
+    }
+
+    let strct = Struct {
+        f1: 0,
+        f2: InnerStruct {
+            f4: 8,
+            f5: 13
+        },
+        f3: InnerStruct {
+            f4: 21,
+            f5: 34
+        }
+    };
+
+    // struct-as-tuple
+    {
+        let serialized: Vec<u8> = rmps::to_vec(&strct).unwrap();
+        let deserialized: Struct = rmps::from_slice(&serialized).unwrap();
+        assert_eq!(deserialized, strct);
+    }
+
+    // struct-as-map
+    {
+        let serialized: Vec<u8> = rmps::to_vec_named(&strct).unwrap();
+        let deserialized: Struct = rmps::from_slice(&serialized).unwrap();
+        assert_eq!(deserialized, strct);
+    }
+}
+
 // Checks whether deserialization and serialization can both work with enum variants as strings
 #[test]
 fn round_variant_string() {