ser: working Map serialization for flattened structs

src/ser/map.rs

@@ -0,0 +1,95 @@
+use std::io::{Seek, SeekFrom};
+
+use byteorder::WriteBytesExt;
+use failure::ResultExt;
+use serde::ser::{Serialize, SerializeMap};
+
+use error::{Error, KbinErrorKind};
+use node_types::StandardType;
+use ser::{Result, Serializer, TypeHint, WriteMode, ARRAY_MASK};
+
+pub struct Map<'a> {
+  ser: &'a mut Serializer,
+
+  current_node_index: u64,
+
+  key_node_type: StandardType,
+}
+
+impl<'a> Map<'a> {
+  pub fn new(ser: &'a mut Serializer) -> Result<Self> {
+    debug!("Map::new()");
+
+    ser.write_node(TypeHint::from_type(StandardType::NodeStart))?;
+    ser.write_identifier("something")?;
+
+    Ok(Self {
+      ser,
+      current_node_index: 0,
+      key_node_type: StandardType::NodeStart,
+    })
+  }
+}
+
+impl<'a> SerializeMap for Map<'a> {
+  type Ok = Option<TypeHint>;
+  type Error = Error;
+
+  fn serialize_key<T>(&mut self, key: &T) -> Result<()>
+    where T: ?Sized + Serialize
+  {
+    trace!("--> <Map as SerializeMap>::serialize_key()");
+
+    self.current_node_index = self.ser.node_buf.position();
+    self.ser.node_buf.write_u8(0).context(KbinErrorKind::DataWrite("placeholder"))?;
+
+    self.ser.write_mode = WriteMode::Identifier;
+    let hint = key.serialize(&mut *self.ser)?.ok_or(KbinErrorKind::MissingTypeHint)?;
+    debug!("<Map as SerializeMap>::serialize_key() => hint: {:?}", hint);
+
+    self.ser.write_mode = WriteMode::Single;
+
+    self.key_node_type = hint.node_type;
+    match hint.node_type {
+      StandardType::Attribute |
+      StandardType::String => Ok(()),
+      node_type => return Err(KbinErrorKind::TypeMismatch(*StandardType::String, *node_type).into()),
+    }
+  }
+
+  fn serialize_value<T>(&mut self, value: &T) -> Result<()>
+    where T: ?Sized + Serialize
+  {
+    trace!("--> <Map as SerializeMap>::serialize_value()");
+
+    let mut hint = value.serialize(&mut *self.ser)?.ok_or(KbinErrorKind::MissingTypeHint)?;
+    debug!("<Map as SerializeMap>::serialize_value() => hint: {:?}", hint);
+
+    // Attributes must have a string body
+    if self.key_node_type == StandardType::Attribute {
+      if hint.node_type == StandardType::String {
+        hint = TypeHint::from_type(StandardType::Attribute);
+      } else {
+        return Err(KbinErrorKind::TypeMismatch(*StandardType::Attribute, *hint.node_type).into());
+      }
+    }
+
+    let new_pos = self.ser.node_buf.position();
+    self.ser.node_buf.seek(SeekFrom::Start(self.current_node_index)).context(KbinErrorKind::Seek)?;
+    self.ser.write_node(hint)?;
+    self.ser.node_buf.seek(SeekFrom::Start(new_pos)).context(KbinErrorKind::Seek)?;
+
+    if self.key_node_type != StandardType::Attribute {
+      self.ser.node_buf.write_u8(StandardType::NodeEnd.id | ARRAY_MASK).context(KbinErrorKind::DataWrite("node end"))?;
+    }
+
+    Ok(())
+  }
+
+  fn end(self) -> Result<Self::Ok> {
+    trace!("<Map as SerializeMap>::end()");
+    self.ser.node_buf.write_u8(StandardType::NodeEnd.id | ARRAY_MASK).context(KbinErrorKind::DataWrite("node end"))?;
+
+    Ok(None)
+  }
+}

src/ser/mod.rs

@@ -9,13 +9,16 @@ use byte_buffer::ByteBufferWrite;
 use encoding_type::EncodingType;
 use node_types::StandardType;
 use error::{Error, KbinError, KbinErrorKind};
+use sixbit::pack_sixbit;
 use super::{ARRAY_MASK, SIGNATURE, SIG_COMPRESSED};
 
 mod custom;
+mod map;
 mod structure;
 mod tuple;
 
 use self::custom::Custom;
+use self::map::Map;
 use self::structure::Struct;
 use self::tuple::Tuple;
 
@@ -27,6 +30,7 @@ pub type Result<T> = StdResult<T, Error>;
 pub(crate) enum WriteMode {
   Single,
   Array,
+  Identifier,
 }
 
 pub struct Serializer {
@@ -97,6 +101,19 @@ impl Serializer {
 
     Ok(output)
   }
+
+  fn write_node(&mut self, hint: TypeHint) -> Result<()> {
+    let node_type = hint.node_type;
+    let array_mask = if hint.is_array { ARRAY_MASK } else { 0 };
+    self.node_buf.write_u8(node_type.id | array_mask).context(KbinErrorKind::DataWrite(node_type.name))?;
+
+    Ok(())
+  }
+
+  fn write_identifier(&mut self, key: &str) -> Result<()> {
+    pack_sixbit(&mut *self.node_buf, key)?;
+    Ok(())
+  }
 }
 
 // `straight_impl` passes a single element array to `write_aligned` where
@@ -105,7 +122,7 @@ impl Serializer {
 macro_rules! ser_type {
   (byte; $inner_type:ident, $method:ident, $standard_type:ident $($cast:tt)*) => {
     fn $method(self, value: $inner_type) -> Result<Self::Ok> {
-      debug!(concat!(stringify!($method), " => value: {}"), value);
+      trace!(concat!("Serializer::", stringify!($method), " => value: {}"), value);
 
       let node_type = StandardType::$standard_type;
       match self.write_mode {
@@ -115,7 +132,8 @@ macro_rules! ser_type {
         },
         WriteMode::Array => {
           self.data_buf.write_u8(value $($cast)*).context(KbinErrorKind::DataWrite(node_type.name))?;
-        }
+        },
+        WriteMode::Identifier => return Err(KbinErrorKind::InvalidState.into()),
       };
 
       Ok(Some(TypeHint::from_type(node_type)))
@@ -123,7 +141,7 @@ macro_rules! ser_type {
   };
   (large; $inner_type:ident, $method:ident, $write_method:ident, $standard_type:ident $($cast:tt)*) => {
     fn $method(self, value: $inner_type) -> Result<Self::Ok> {
-      debug!(concat!(stringify!($method), " => value: {}"), value);
+      trace!(concat!("Serializer::", stringify!($method), " => value: {}"), value);
 
       let node_type = StandardType::$standard_type;
       match self.write_mode {
@@ -134,7 +152,8 @@ macro_rules! ser_type {
         },
         WriteMode::Array => {
           self.data_buf.$write_method::<BigEndian>(value $($cast)*).context(KbinErrorKind::DataWrite(node_type.name))?;
-        }
+        },
+        WriteMode::Identifier => return Err(KbinErrorKind::InvalidState.into()),
       };
 
       Ok(Some(TypeHint::from_type(node_type)))
@@ -150,7 +169,7 @@ impl<'a> ser::Serializer for &'a mut Serializer {
   type SerializeTuple = Tuple<'a>;
   type SerializeTupleStruct = Custom<'a>;
   type SerializeTupleVariant = Impossible<Self::Ok, Self::Error>;
-  type SerializeMap = Self;
+  type SerializeMap = Map<'a>;
   type SerializeStruct = Struct<'a>;
   type SerializeStructVariant = Impossible<Self::Ok, Self::Error>;
 
@@ -171,23 +190,39 @@ impl<'a> ser::Serializer for &'a mut Serializer {
   ser_type!(large; f64, serialize_f64, write_f64, Double);
 
   fn serialize_char(self, value: char) -> Result<Self::Ok> {
-    debug!("serialize_char => value: {}", value);
+    trace!("Serializer::serialize_char(value: {})", value);
     self.data_buf.write_str(self.encoding, &value.to_string())?;
 
     Ok(Some(TypeHint::from_type(StandardType::String)))
   }
 
   fn serialize_str(self, value: &str) -> Result<Self::Ok> {
-    debug!("serialize_str => value: {}", value);
-    self.data_buf.write_str(self.encoding, value)?;
+    trace!("Serializer::serialize_str(value: {})", value);
 
-    Ok(Some(TypeHint::from_type(StandardType::String)))
+    let hint = if value.starts_with("attr_") {
+      let key = &value["attr_".len()..];
+      debug!("Serializer::serialize_str(key: {}) => writing as attribute", key);
+
+      self.write_identifier(key)?;
+      Some(TypeHint::from_type(StandardType::Attribute))
+    } else if self.write_mode == WriteMode::Identifier {
+      debug!("Serializer::serialize_str(identifier: {}) => writing as identifier", value);
+
+      self.write_identifier(value)?;
+      Some(TypeHint::from_type(StandardType::String))
+    } else {
+      self.data_buf.write_str(self.encoding, value)?;
+
+      Some(TypeHint::from_type(StandardType::String))
+    };
+
+    Ok(hint)
   }
 
   // Binary data is handled separately from other array types.
   // Binary data should also be the only element of its node.
   fn serialize_bytes(self, value: &[u8]) -> Result<Self::Ok> {
-    debug!("serialize_bytes => value: {:02x?}", value);
+    trace!("Serializer::serialize_bytes(value: {:02x?})", value);
     let node_type = StandardType::Binary;
     let size = (value.len() as u32) * (node_type.size as u32);
     self.data_buf.write_u32::<BigEndian>(size).context(KbinErrorKind::DataWrite("binary node size"))?;
@@ -206,9 +241,8 @@ impl<'a> ser::Serializer for &'a mut Serializer {
   fn serialize_some<T>(self, v: &T) -> Result<Self::Ok>
     where T: ?Sized + Serialize
   {
-    debug!("serialize_some");
-    let hint = v.serialize(&mut *self)?;
-    Ok(hint)
+    trace!("Serializer::serialize_some()");
+    v.serialize(&mut *self)
   }
 
   // TODO: Figure out a good way to serialize this
@@ -269,8 +303,8 @@ impl<'a> ser::Serializer for &'a mut Serializer {
   }
 
   fn serialize_map(self, len: Option<usize>) -> Result<Self::SerializeMap> {
-    debug!("serialize_map => len: {:?}", len);
-    Ok(self)
+    trace!("Serializer::serialize_map(len: {:?})", len);
+    Map::new(self)
   }
 
   fn serialize_struct(self, name: &'static str, len: usize) -> Result<Self::SerializeStruct> {
@@ -284,27 +318,3 @@ impl<'a> ser::Serializer for &'a mut Serializer {
     Err(Error::Message("struct variant not supported".to_string()))
   }
 }
-
-impl<'a> ser::SerializeMap for &'a mut Serializer {
-  type Ok = Option<TypeHint>;
-  type Error = Error;
-
-  fn serialize_key<T>(&mut self, _key: &T) -> Result<()>
-    where T: ?Sized + Serialize
-  {
-    debug!("SerializeMap: serialize_key");
-    unimplemented!();
-  }
-
-  fn serialize_value<T>(&mut self, _value: &T) -> Result<()>
-    where T: ?Sized + Serialize
-  {
-    debug!("SerializeMap: serialize_value");
-    unimplemented!();
-  }
-
-  fn end(self) -> Result<Self::Ok> {
-    debug!("SerializeMap: end");
-    Ok(None)
-  }
-}

src/ser/structure.rs

@@ -2,9 +2,9 @@ use byteorder::WriteBytesExt;
 use failure::ResultExt;
 use serde::ser::{Serialize, SerializeStruct};
 
+use error::{Error, KbinErrorKind};
 use node_types::StandardType;
-use error::KbinErrorKind;
-use ser::{Error, Result, Serializer, TypeHint, ARRAY_MASK};
+use ser::{Result, Serializer, TypeHint, ARRAY_MASK};
 use sixbit::pack_sixbit;
 
 pub struct Struct<'a> {
@@ -65,10 +65,24 @@ impl<'a> SerializeStruct for Struct<'a> {
 
       // Struct handler outputs the `NodeStart` event by itself. Avoid repeating it.
       if node_type != StandardType::NodeStart {
-        self.ser.node_buf.write_u8(node_type.id | array_mask).context(KbinErrorKind::DataWrite(node_type.name))?;
-        pack_sixbit(&mut *self.ser.node_buf, key)?;
+        // Serialize fields that start with "attr_" as Attribute nodes
+        let (node_type, key) = if key.starts_with("attr_") {
+          let key = &key["attr_".len()..];
+          debug!("SerializeStruct(name: {})::serialize_field(key: {}) => writing as attribute", self.name, key);
+
+          // Attribute nodes are always strings
+          if node_type != StandardType::String {
+            return Err(KbinErrorKind::TypeMismatch(*StandardType::String, *node_type).into());
+          }
+
+          (StandardType::Attribute, key)
+        } else {
+          (node_type, key)
+        };
+
+        self.ser.node_buf.write_u8(node_type.id | array_mask).context(KbinErrorKind::DataWrite(node_type.name))?;
+        self.ser.write_identifier(key)?;
 
-        // TODO: Make sure this does not prematurely end nodes
         if node_type != StandardType::Attribute {
           self.ser.node_buf.write_u8(StandardType::NodeEnd.id | ARRAY_MASK).context(KbinErrorKind::DataWrite("node end"))?;
         }