writer: new Writeable trait to support writing from Element, NodeCollection, and Node

src/error.rs

@@ -5,6 +5,7 @@ use std::string::FromUtf8Error;
 
 use failure::{Backtrace, Compat, Context, Fail};
 use serde::{de, ser};
+use value::Value;
 
 use node_types::{KbinType, StandardType};
 
@@ -107,10 +108,13 @@ pub enum KbinErrorKind {
   #[fail(display = "Type mismatch, expected: {}, found: {}", _0, _1)]
   TypeMismatch(StandardType, StandardType),
 
+  #[fail(display = "Value mismatch, expected {}, but found {:?}", _0, _1)]
+  ValueTypeMismatch(StandardType, Value),
+
   #[fail(display = "Invalid input for boolean: {}", _0)]
   InvalidBooleanInput(u8),
 
-  #[fail(display = "Invalid node type for operation: {:?}", _0)]
+  #[fail(display = "Invalid node type {:?} for operation", _0)]
   InvalidNodeType(StandardType),
 
   #[fail(display = "Invalid state")]

src/lib.rs

@@ -48,7 +48,7 @@ pub use options::Options;
 pub use de::from_bytes;
 pub use ser::to_bytes;
 pub use value::Value;
-pub use writer::Writer;
+pub use writer::{Writer, Writeable};
 
 const SIGNATURE: u8 = 0xA0;
 
@@ -161,12 +161,16 @@ pub fn node_from_binary(input: Bytes) -> Result<(Node, EncodingType)> {
   Ok((node, encoding))
 }
 
-pub fn to_binary(input: &Element) -> Result<Vec<u8>> {
+pub fn to_binary<T>(input: &T) -> Result<Vec<u8>>
+  where T: Writeable<T>
+{
   let mut writer = Writer::new();
   writer.to_binary(input)
 }
 
-pub fn to_binary_with_options(options: Options, input: &Element) -> Result<Vec<u8>> {
+pub fn to_binary_with_options<T>(options: Options, input: &T) -> Result<Vec<u8>>
+  where T: Writeable<T>
+{
   let mut writer = Writer::with_options(options);
   writer.to_binary(input)
 }

src/value/mod.rs

@@ -309,7 +309,7 @@ macro_rules! tuple {
       Ok(value)
     }
 
-    fn to_bytes_inner(self, output: &mut Vec<u8>) -> Result<(), KbinError> {
+    fn to_bytes_inner(&self, output: &mut Vec<u8>) -> Result<(), KbinError> {
       debug!("Value::to_bytes_inner(self: {:?})", self);
 
       macro_rules! gen_error {
@@ -319,20 +319,20 @@ macro_rules! tuple {
       }
 
       match self {
-        Value::S8(n) => output.push(n as u8),
-        Value::U8(n) => output.push(n),
-        Value::S16(n) => output.write_i16::<BigEndian>(n).context(gen_error!(S16))?,
-        Value::U16(n) => output.write_u16::<BigEndian>(n).context(gen_error!(U16))?,
-        Value::S32(n) => output.write_i32::<BigEndian>(n).context(gen_error!(S32))?,
-        Value::U32(n) => output.write_u32::<BigEndian>(n).context(gen_error!(U32))?,
-        Value::S64(n) => output.write_i64::<BigEndian>(n).context(gen_error!(S64))?,
-        Value::U64(n) => output.write_u64::<BigEndian>(n).context(gen_error!(U64))?,
-        Value::Binary(data) => output.extend_from_slice(&data),
-        Value::Time(n) => output.write_u32::<BigEndian>(n).context(gen_error!(Time))?,
+        Value::S8(ref n) => output.push(*n as u8),
+        Value::U8(ref n) => output.push(*n),
+        Value::S16(ref n) => output.write_i16::<BigEndian>(*n).context(gen_error!(S16))?,
+        Value::U16(ref n) => output.write_u16::<BigEndian>(*n).context(gen_error!(U16))?,
+        Value::S32(ref n) => output.write_i32::<BigEndian>(*n).context(gen_error!(S32))?,
+        Value::U32(ref n) => output.write_u32::<BigEndian>(*n).context(gen_error!(U32))?,
+        Value::S64(ref n) => output.write_i64::<BigEndian>(*n).context(gen_error!(S64))?,
+        Value::U64(ref n) => output.write_u64::<BigEndian>(*n).context(gen_error!(U64))?,
+        Value::Binary(ref data) => output.extend_from_slice(data),
+        Value::Time(ref n) => output.write_u32::<BigEndian>(*n).context(gen_error!(Time))?,
         Value::Ip4(addr) => output.extend_from_slice(&addr.octets()),
-        Value::Float(n) => output.write_f32::<BigEndian>(n).context(gen_error!(Float))?,
-        Value::Double(n) => output.write_f64::<BigEndian>(n).context(gen_error!(Double))?,
-        Value::Boolean(v) => output.push(if v { 0x01 } else { 0x00 }),
+        Value::Float(ref n) => output.write_f32::<BigEndian>(*n).context(gen_error!(Float))?,
+        Value::Double(ref n) => output.write_f64::<BigEndian>(*n).context(gen_error!(Double))?,
+        Value::Boolean(ref v) => output.push(if *v { 0x01 } else { 0x00 }),
         Value::Array(_, values) => {
           for value in values {
             value.to_bytes_inner(output)?;
@@ -428,13 +428,18 @@ macro_rules! construct_types {
         ]
       }
 
-      pub fn to_bytes(self) -> Result<Vec<u8>, KbinError> {
+      pub fn to_bytes(&self) -> Result<Vec<u8>, KbinError> {
         let mut output = Vec::new();
         self.to_bytes_inner(&mut output)?;
 
         Ok(output)
       }
 
+      #[inline]
+      pub fn to_bytes_into(&self, output: &mut Vec<u8>) -> Result<(), KbinError> {
+        self.to_bytes_inner(output)
+      }
+
       pub fn standard_type(&self) -> StandardType {
         match *self {
           $(
@@ -448,17 +453,24 @@ macro_rules! construct_types {
         }
       }
 
-      pub fn as_i8(&self) -> Option<i8> {
+      pub fn as_i8(&self) -> Result<i8, KbinError> {
         match self {
-          Value::S8(ref n) => Some(*n),
-          _ => None,
+          Value::S8(ref n) => Ok(*n),
+          value => Err(KbinErrorKind::ValueTypeMismatch(StandardType::S8, value.clone()).into()),
         }
       }
 
-      pub fn as_u8(&self) -> Option<u8> {
+      pub fn as_u8(&self) -> Result<u8, KbinError> {
         match self {
-          Value::U8(ref n) => Some(*n),
-          _ => None,
+          Value::U8(ref n) => Ok(*n),
+          value => Err(KbinErrorKind::ValueTypeMismatch(StandardType::U8, value.clone()).into()),
+        }
+      }
+
+      pub fn as_attribute(self) -> Result<String, KbinError> {
+        match self {
+          Value::Attribute(s) => Ok(s),
+          value => Err(KbinErrorKind::ValueTypeMismatch(StandardType::Attribute, value).into()),
         }
       }
     }

src/writer.rs

@@ -4,35 +4,67 @@ use byteorder::{BigEndian, WriteBytesExt};
 use failure::ResultExt;
 use minidom::Element;
 use rustc_hex::FromHex;
-use byte_buffer::ByteBufferWrite;
-use node_types::StandardType;
-use sixbit::Sixbit;
 
+use byte_buffer::ByteBufferWrite;
 use compression::Compression;
 use error::{KbinErrorKind, Result};
+use node::{Node, NodeCollection};
+use node_types::StandardType;
 use options::Options;
+use sixbit::Sixbit;
 use value::Value;
 
 use super::{ARRAY_MASK, SIGNATURE};
 
-pub struct Writer {
-  options: Options,
+fn write_value(options: &Options, data_buf: &mut ByteBufferWrite, node_type: StandardType, is_array: bool, value: &Value) -> Result<()> {
+  match value {
+    Value::Binary(data) => {
+      trace!("data: 0x{:02x?}", data);
+
+      let size = (data.len() as u32) * (node_type.size as u32);
+      data_buf.write_u32::<BigEndian>(size).context(KbinErrorKind::DataWrite("binary node size"))?;
+      data_buf.write_all(&data).context(KbinErrorKind::DataWrite("binary"))?;
+      data_buf.realign_writes(None)?;
+    },
+    Value::String(text) => {
+      data_buf.write_str(options.encoding, &text)?;
+    },
+    Value::Array(node_type, values) => {
+      if !is_array {
+        return Err(KbinErrorKind::InvalidState.into());
+      }
+
+      let count = values.len();
+      let total_size = count * node_type.count * node_type.size;
+
+      let mut data = Vec::with_capacity(total_size);
+      for value in values {
+        value.to_bytes_into(&mut data)?;
+      }
+
+      data_buf.write_u32::<BigEndian>(total_size as u32).context(KbinErrorKind::DataWrite("node size"))?;
+      data_buf.write_all(&data).context(KbinErrorKind::DataWrite(node_type.name))?;
+      data_buf.realign_writes(None)?;
+    },
+    value => {
+      if is_array {
+        return Err(KbinErrorKind::InvalidState.into());
+      } else {
+        let data = value.to_bytes()?;
+        data_buf.write_aligned(*node_type, &data)?;
+      }
+    },
+  };
+
+  Ok(())
 }
 
-impl Writer {
-  pub fn new() -> Self {
-    Self {
-      options: Options::default(),
-    }
-  }
+pub trait Writeable<T> {
+  fn write_node(options: &Options, node_buf: &mut ByteBufferWrite, data_buf: &mut ByteBufferWrite, input: &T) -> Result<()>;
+}
 
-  pub fn with_options(options: Options) -> Self {
-    Self {
-      options,
-    }
-  }
-
-  fn write_node(&mut self, node_buf: &mut ByteBufferWrite, data_buf: &mut ByteBufferWrite, input: &Element) -> Result<()> {
+impl Writeable<Element> for Element {
+  fn write_node(options: &Options, node_buf: &mut ByteBufferWrite, data_buf: &mut ByteBufferWrite, input: &Element) -> Result<()> {
     let text = input.text();
     let node_type = match input.attr("__type") {
       Some(name) => StandardType::from_name(name),
@@ -66,10 +98,10 @@ impl Writer {
       count);
 
     node_buf.write_u8(node_type.id | array_mask).context(KbinErrorKind::DataWrite(node_type.name))?;
-    match self.options.compression {
+    match options.compression {
       Compression::Compressed => Sixbit::pack(&mut **node_buf, input.name())?,
       Compression::Uncompressed => {
-        let data = self.options.encoding.encode_bytes(input.name())?;
+        let data = options.encoding.encode_bytes(input.name())?;
         let len = (data.len() - 1) as u8;
         node_buf.write_u8(len | ARRAY_MASK).context(KbinErrorKind::DataWrite("node name length"))?;
         node_buf.write_all(&data).context(KbinErrorKind::DataWrite("node name bytes"))?;
@@ -89,7 +121,7 @@ impl Writer {
         data_buf.realign_writes(None)?;
       },
       StandardType::String => {
-        data_buf.write_str(self.options.encoding, &text)?;
+        data_buf.write_str(options.encoding, &text)?;
       },
 
       _ => {
@@ -117,14 +149,14 @@ impl Writer {
 
       trace!("write_node => attr: {}, value: {}", key, value);
 
-      data_buf.write_str(self.options.encoding, value)?;
+      data_buf.write_str(options.encoding, value)?;
 
       let node_type = StandardType::Attribute;
       node_buf.write_u8(node_type.id).context(KbinErrorKind::DataWrite(node_type.name))?;
-      match self.options.compression {
+      match options.compression {
         Compression::Compressed => Sixbit::pack(&mut **node_buf, key)?,
         Compression::Uncompressed => {
-          let data = self.options.encoding.encode_bytes(key)?;
+          let data = options.encoding.encode_bytes(key)?;
           let len = (data.len() - 1) as u8;
           node_buf.write_u8(len | ARRAY_MASK).context(KbinErrorKind::DataWrite("attribute name length"))?;
           node_buf.write_all(&data).context(KbinErrorKind::DataWrite("node name bytes"))?;
@@ -133,7 +165,7 @@ impl Writer {
     }
 
     for child in input.children() {
-      self.write_node(node_buf, data_buf, child)?;
+      Self::write_node(options, node_buf, data_buf, child)?;
     }
 
     // Always has the array bit set
@@ -141,8 +173,133 @@ impl Writer {
 
     Ok(())
   }
+}
 
-  pub fn to_binary(&mut self, input: &Element) -> Result<Vec<u8>> {
+impl Writeable<NodeCollection> for NodeCollection {
+  fn write_node(options: &Options, node_buf: &mut ByteBufferWrite, data_buf: &mut ByteBufferWrite, input: &NodeCollection) -> Result<()> {
+    let (node_type, is_array) = input.base().node_type_tuple();
+    let array_mask = if is_array { ARRAY_MASK } else { 0 };
+    let name = input.base().key()?.ok_or(KbinErrorKind::InvalidState)?;
+
+    node_buf.write_u8(node_type.id | array_mask).context(KbinErrorKind::DataWrite(node_type.name))?;
+    match options.compression {
+      Compression::Compressed => Sixbit::pack(&mut **node_buf, &name)?,
+      Compression::Uncompressed => {
+        let data = options.encoding.encode_bytes(&name)?;
+        let len = (data.len() - 1) as u8;
+        node_buf.write_u8(len | ARRAY_MASK).context(KbinErrorKind::DataWrite("node name length"))?;
+        node_buf.write_all(&data).context(KbinErrorKind::DataWrite("node name bytes"))?;
+      },
+    };
+
+    if node_type != StandardType::NodeStart {
+      let value = input.base().value()?;
+      write_value(options, data_buf, node_type, is_array, &value)?;
+    }
+
+    for attr in input.attributes() {
+      let key = attr.key()?.ok_or(KbinErrorKind::InvalidState)?;
+      let value = attr.value_bytes().ok_or(KbinErrorKind::InvalidState)?;
+
+      data_buf.buf_write(value)?;
+
+      node_buf.write_u8(StandardType::Attribute.id).context(KbinErrorKind::DataWrite(StandardType::Attribute.name))?;
+      match options.compression {
+        Compression::Compressed => Sixbit::pack(&mut **node_buf, &key)?,
+        Compression::Uncompressed => {
+          let data = options.encoding.encode_bytes(&key)?;
+          let len = (data.len() - 1) as u8;
+          node_buf.write_u8(len | ARRAY_MASK).context(KbinErrorKind::DataWrite("attribute name length"))?;
+          node_buf.write_all(&data).context(KbinErrorKind::DataWrite("node name bytes"))?;
+        },
+      };
+    }
+
+    for child in input.children() {
+      Self::write_node(options, node_buf, data_buf, child)?;
+    }
+
+    // Always has the array bit set
+    node_buf.write_u8(StandardType::NodeEnd.id | ARRAY_MASK).context(KbinErrorKind::DataWrite("node end"))?;
+
+    Ok(())
+  }
+}
+
+impl Writeable<Node> for Node {
+  fn write_node(options: &Options, node_buf: &mut ByteBufferWrite, data_buf: &mut ByteBufferWrite, input: &Node) -> Result<()> {
+    let (node_type, is_array) = match input.value() {
+      Some(Value::Array(node_type, _)) => (*node_type, true),
+      Some(ref value) => (value.standard_type(), false),
+      None => (StandardType::NodeStart, false),
+    };
+    let array_mask = if is_array { ARRAY_MASK } else { 0 };
+
+    node_buf.write_u8(node_type.id | array_mask).context(KbinErrorKind::DataWrite(node_type.name))?;
+    match options.compression {
+      Compression::Compressed => Sixbit::pack(&mut **node_buf, &input.key())?,
+      Compression::Uncompressed => {
+        let data = options.encoding.encode_bytes(&input.key())?;
+        let len = (data.len() - 1) as u8;
+        node_buf.write_u8(len | ARRAY_MASK).context(KbinErrorKind::DataWrite("node name length"))?;
+        node_buf.write_all(&data).context(KbinErrorKind::DataWrite("node name bytes"))?;
+      },
+    };
+
+    if let Some(value) = input.value() {
+      write_value(options, data_buf, node_type, is_array, value)?;
+    }
+
+    if let Some(attributes) = input.attributes() {
+      for (key, value) in attributes {
+        data_buf.write_str(options.encoding, value)?;
+
+        node_buf.write_u8(StandardType::Attribute.id).context(KbinErrorKind::DataWrite(StandardType::Attribute.name))?;
+        match options.compression {
+          Compression::Compressed => Sixbit::pack(&mut **node_buf, &key)?,
+          Compression::Uncompressed => {
+            let data = options.encoding.encode_bytes(&key)?;
+            let len = (data.len() - 1) as u8;
+            node_buf.write_u8(len | ARRAY_MASK).context(KbinErrorKind::DataWrite("attribute name length"))?;
+            node_buf.write_all(&data).context(KbinErrorKind::DataWrite("node name bytes"))?;
+          },
+        };
+      }
+    }
+
+    if let Some(children) = input.children() {
+      for child in children {
+        Self::write_node(options, node_buf, data_buf, child)?;
+      }
+    }
+
+    // Always has the array bit set
+    node_buf.write_u8(StandardType::NodeEnd.id | ARRAY_MASK).context(KbinErrorKind::DataWrite("node end"))?;
+
+    Ok(())
+  }
+}
+
+pub struct Writer {
+  options: Options,
+}
+
+impl Writer {
+  pub fn new() -> Self {
+    Self {
+      options: Options::default(),
+    }
+  }
+
+  pub fn with_options(options: Options) -> Self {
+    Self {
+      options,
+    }
+  }
+
+  pub fn to_binary<T>(&mut self, input: &T) -> Result<Vec<u8>>
+    where T: Writeable<T>
+  {
     let mut header = Cursor::new(Vec::with_capacity(8));
     header.write_u8(SIGNATURE).context(KbinErrorKind::HeaderWrite("signature"))?;
 
@@ -156,7 +313,7 @@ impl Writer {
     let mut node_buf = ByteBufferWrite::new(Vec::new());
     let mut data_buf = ByteBufferWrite::new(Vec::new());
 
-    self.write_node(&mut node_buf, &mut data_buf, input)?;
+    T::write_node(&self.options, &mut node_buf, &mut data_buf, input)?;
 
     node_buf.write_u8(StandardType::FileEnd.id | ARRAY_MASK).context(KbinErrorKind::DataWrite("file end"))?;
     node_buf.realign_writes(None)?;