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 { + as Encodable>::ssz_fixed_len() + ::ssz_fixed_len(); - let mut fixed = Vec::with_capacity(offset); - let mut variable = vec![]; + let mut encoder = SszEncoder::container(offset); - if ::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 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, + variable_bytes: Vec, +} + +impl SszEncoder { + pub fn container(num_fixed_bytes: usize) -> Self { + Self { + offset: num_fixed_bytes, + fixed_bytes: vec![], + variable_bytes: vec![], + } + } + + pub fn append(&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 { + 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;