Add SszEncoder

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 encoder = SszEncoder::container(offset);
-        let mut variable = vec![];
 
-        if <u16 as Encodable>::is_ssz_fixed_len() {
+        encoder.append(&self.a);
-            self.a.ssz_append(&mut fixed);
+        encoder.append(&self.b);
-        } else {
+        encoder.append(&self.c);
-            fixed.append(encode_length())
-        }
 
-        if <Vec<u16> as Encodable>::is_ssz_fixed_len() {
+        buf.append(&mut encoder.drain());
-            self.a.ssz_append(&mut fixed);
-        }
     }
 }
 

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,

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;