From aeb17c73f646aa784dfe821be0231d1212df4555 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Mon, 6 May 2019 09:01:28 +1000
Subject: [PATCH] Add SszEncoder

---
 eth2/utils/ssz/examples/struct_definition.rs | 17 ++++------
 eth2/utils/ssz/src/encode.rs                 | 33 ++++++++++++++++++++
 eth2/utils/ssz/src/lib.rs                    |  2 +-
 3 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/eth2/utils/ssz/examples/struct_definition.rs b/eth2/utils/ssz/examples/struct_definition.rs
index f363030f5..eecf0ff14 100644
--- a/eth2/utils/ssz/examples/struct_definition.rs
+++ b/eth2/utils/ssz/examples/struct_definition.rs
@@ -1,4 +1,4 @@
-use ssz::{encode_length, Decodable, DecodeError, Encodable, SszDecoderBuilder, SszStream};
+use ssz::{Decodable, DecodeError, Encodable, SszDecoderBuilder, SszEncoder, SszStream};
 
 #[derive(Debug, PartialEq)]
 pub struct Foo {
@@ -17,18 +17,13 @@ impl Encodable for Foo {
             + <Vec<u16> as Encodable>::ssz_fixed_len()
             + <u16 as Encodable>::ssz_fixed_len();
 
-        let mut fixed = Vec::with_capacity(offset);
-        let mut variable = vec![];
+        let mut encoder = SszEncoder::container(offset);
 
-        if <u16 as Encodable>::is_ssz_fixed_len() {
-            self.a.ssz_append(&mut fixed);
-        } else {
-            fixed.append(encode_length())
-        }
+        encoder.append(&self.a);
+        encoder.append(&self.b);
+        encoder.append(&self.c);
 
-        if <Vec<u16> as Encodable>::is_ssz_fixed_len() {
-            self.a.ssz_append(&mut fixed);
-        }
+        buf.append(&mut encoder.drain());
     }
 }
 
diff --git a/eth2/utils/ssz/src/encode.rs b/eth2/utils/ssz/src/encode.rs
index dccc652a2..072d82262 100644
--- a/eth2/utils/ssz/src/encode.rs
+++ b/eth2/utils/ssz/src/encode.rs
@@ -25,6 +25,39 @@ pub trait Encodable {
     }
 }
 
+pub struct SszEncoder {
+    offset: usize,
+    fixed_bytes: Vec<u8>,
+    variable_bytes: Vec<u8>,
+}
+
+impl SszEncoder {
+    pub fn container(num_fixed_bytes: usize) -> Self {
+        Self {
+            offset: num_fixed_bytes,
+            fixed_bytes: vec![],
+            variable_bytes: vec![],
+        }
+    }
+
+    pub fn append<T: Encodable>(&mut self, item: &T) {
+        if T::is_ssz_fixed_len() {
+            item.ssz_append(&mut self.fixed_bytes);
+        } else {
+            self.fixed_bytes
+                .append(&mut encode_length(self.offset + self.variable_bytes.len()));
+
+            item.ssz_append(&mut self.variable_bytes);
+        }
+    }
+
+    pub fn drain(mut self) -> Vec<u8> {
+        self.fixed_bytes.append(&mut self.variable_bytes);
+
+        self.fixed_bytes
+    }
+}
+
 pub struct VariableLengths {
     pub fixed_bytes_position: usize,
     pub variable_bytes_length: usize,
diff --git a/eth2/utils/ssz/src/lib.rs b/eth2/utils/ssz/src/lib.rs
index 809daf49f..c91acb70b 100644
--- a/eth2/utils/ssz/src/lib.rs
+++ b/eth2/utils/ssz/src/lib.rs
@@ -17,7 +17,7 @@ mod decode;
 mod encode;
 
 pub use decode::{Decodable, DecodeError, SszDecoderBuilder};
-pub use encode::{Encodable, SszStream};
+pub use encode::{Encodable, SszEncoder, SszStream};
 
 pub const BYTES_PER_LENGTH_OFFSET: usize = 4;
 pub const MAX_LENGTH_VALUE: usize = 1 << (BYTES_PER_LENGTH_OFFSET * 8) - 1;