Beacon state diffs!

consensus/ssz/src/decode/impls.rs

@@ -2,6 +2,8 @@ use super::*;
 use core::num::NonZeroUsize;
 use ethereum_types::{H160, H256, U128, U256};
 use smallvec::SmallVec;
+use std::collections::BTreeMap;
+use std::iter::{self, FromIterator};
 use std::sync::Arc;
 
 macro_rules! impl_decodable_for_uint {
@@ -380,14 +382,14 @@ macro_rules! impl_for_vec {
 
             fn from_ssz_bytes(bytes: &[u8]) -> Result<Self, DecodeError> {
                 if bytes.is_empty() {
-                    Ok(vec![].into())
+                    Ok(Self::from_iter(iter::empty()))
                 } else if T::is_ssz_fixed_len() {
                     bytes
                         .chunks(T::ssz_fixed_len())
-                        .map(|chunk| T::from_ssz_bytes(chunk))
+                        .map(T::from_ssz_bytes)
                         .collect()
                 } else {
-                    decode_list_of_variable_length_items(bytes, $max_len).map(|vec| vec.into())
+                    decode_list_of_variable_length_items(bytes, $max_len)
                 }
             }
         }
@@ -404,17 +406,40 @@ impl_for_vec!(SmallVec<[T; 6]>, Some(6));
 impl_for_vec!(SmallVec<[T; 7]>, Some(7));
 impl_for_vec!(SmallVec<[T; 8]>, Some(8));
 
+impl<K, V> Decode for BTreeMap<K, V>
+where
+    K: Decode + Ord,
+    V: Decode,
+{
+    fn is_ssz_fixed_len() -> bool {
+        false
+    }
+
+    fn from_ssz_bytes(bytes: &[u8]) -> Result<Self, DecodeError> {
+        if bytes.is_empty() {
+            Ok(Self::from_iter(iter::empty()))
+        } else if <(K, V)>::is_ssz_fixed_len() {
+            bytes
+                .chunks(<(K, V)>::ssz_fixed_len())
+                .map(<(K, V)>::from_ssz_bytes)
+                .collect()
+        } else {
+            decode_list_of_variable_length_items(bytes, None)
+        }
+    }
+}
+
 /// Decodes `bytes` as if it were a list of variable-length items.
 ///
-/// The `ssz::SszDecoder` can also perform this functionality, however it it significantly faster
-/// as it is optimized to read same-typed items whilst `ssz::SszDecoder` supports reading items of
-/// differing types.
-pub fn decode_list_of_variable_length_items<T: Decode>(
+/// The `ssz::SszDecoder` can also perform this functionality, however this function is
+/// significantly faster as it is optimized to read same-typed items whilst `ssz::SszDecoder`
+/// supports reading items of differing types.
+pub fn decode_list_of_variable_length_items<T: Decode, Container: FromIterator<T>>(
     bytes: &[u8],
     max_len: Option<usize>,
-) -> Result<Vec<T>, DecodeError> {
+) -> Result<Container, DecodeError> {
     if bytes.is_empty() {
-        return Ok(vec![]);
+        return Ok(Container::from_iter(iter::empty()));
     }
 
     let first_offset = read_offset(bytes)?;
@@ -433,35 +458,25 @@ pub fn decode_list_of_variable_length_items<T: Decode>(
         )));
     }
 
-    // Only initialize the vec with a capacity if a maximum length is provided.
-    //
-    // We assume that if a max length is provided then the application is able to handle an
-    // allocation of this size.
-    let mut values = if max_len.is_some() {
-        Vec::with_capacity(num_items)
-    } else {
-        vec![]
-    };
-
     let mut offset = first_offset;
-    for i in 1..=num_items {
-        let slice_option = if i == num_items {
-            bytes.get(offset..)
-        } else {
-            let start = offset;
+    (1..=num_items)
+        .map(|i| {
+            let slice_option = if i == num_items {
+                bytes.get(offset..)
+            } else {
+                let start = offset;
 
-            let next_offset = read_offset(&bytes[(i * BYTES_PER_LENGTH_OFFSET)..])?;
-            offset = sanitize_offset(next_offset, Some(offset), bytes.len(), Some(first_offset))?;
+                let next_offset = read_offset(&bytes[(i * BYTES_PER_LENGTH_OFFSET)..])?;
+                offset =
+                    sanitize_offset(next_offset, Some(offset), bytes.len(), Some(first_offset))?;
 
-            bytes.get(start..offset)
-        };
+                bytes.get(start..offset)
+            };
 
-        let slice = slice_option.ok_or(DecodeError::OutOfBoundsByte { i: offset })?;
-
-        values.push(T::from_ssz_bytes(slice)?);
-    }
-
-    Ok(values)
+            let slice = slice_option.ok_or(DecodeError::OutOfBoundsByte { i: offset })?;
+            T::from_ssz_bytes(slice)
+        })
+        .collect()
 }
 
 #[cfg(test)]

consensus/ssz/src/encode/impls.rs

@@ -2,6 +2,7 @@ use super::*;
 use core::num::NonZeroUsize;
 use ethereum_types::{H160, H256, U128, U256};
 use smallvec::SmallVec;
+use std::collections::BTreeMap;
 use std::sync::Arc;
 
 macro_rules! impl_encodable_for_uint {
@@ -220,6 +221,65 @@ impl<T: Encode> Encode for Arc<T> {
     }
 }
 
+// Encode transparently through references.
+impl<'a, T: Encode> Encode for &'a T {
+    fn is_ssz_fixed_len() -> bool {
+        T::is_ssz_fixed_len()
+    }
+
+    fn ssz_fixed_len() -> usize {
+        T::ssz_fixed_len()
+    }
+
+    fn ssz_append(&self, buf: &mut Vec<u8>) {
+        T::ssz_append(self, buf)
+    }
+
+    fn ssz_bytes_len(&self) -> usize {
+        T::ssz_bytes_len(self)
+    }
+}
+
+/// Compute the encoded length of a vector-like sequence of `T`.
+pub fn sequence_ssz_bytes_len<I, T>(iter: I) -> usize
+where
+    I: Iterator<Item = T> + ExactSizeIterator,
+    T: Encode,
+{
+    // Compute length before doing any iteration.
+    let length = iter.len();
+    if <T as Encode>::is_ssz_fixed_len() {
+        <T as Encode>::ssz_fixed_len() * length
+    } else {
+        let mut len = iter.map(|item| item.ssz_bytes_len()).sum();
+        len += BYTES_PER_LENGTH_OFFSET * length;
+        len
+    }
+}
+
+/// Encode a vector-like sequence of `T`.
+pub fn sequence_ssz_append<I, T>(iter: I, buf: &mut Vec<u8>)
+where
+    I: Iterator<Item = T> + ExactSizeIterator,
+    T: Encode,
+{
+    if T::is_ssz_fixed_len() {
+        buf.reserve(T::ssz_fixed_len() * iter.len());
+
+        for item in iter {
+            item.ssz_append(buf);
+        }
+    } else {
+        let mut encoder = SszEncoder::container(buf, iter.len() * BYTES_PER_LENGTH_OFFSET);
+
+        for item in iter {
+            encoder.append(&item);
+        }
+
+        encoder.finalize();
+    }
+}
+
 macro_rules! impl_for_vec {
     ($type: ty) => {
         impl<T: Encode> Encode for $type {
@@ -228,32 +288,11 @@ macro_rules! impl_for_vec {
             }
 
             fn ssz_bytes_len(&self) -> usize {
-                if <T as Encode>::is_ssz_fixed_len() {
-                    <T as Encode>::ssz_fixed_len() * self.len()
-                } else {
-                    let mut len = self.iter().map(|item| item.ssz_bytes_len()).sum();
-                    len += BYTES_PER_LENGTH_OFFSET * self.len();
-                    len
-                }
+                sequence_ssz_bytes_len(self.iter())
             }
 
             fn ssz_append(&self, buf: &mut Vec<u8>) {
-                if T::is_ssz_fixed_len() {
-                    buf.reserve(T::ssz_fixed_len() * self.len());
-
-                    for item in self {
-                        item.ssz_append(buf);
-                    }
-                } else {
-                    let mut encoder =
-                        SszEncoder::container(buf, self.len() * BYTES_PER_LENGTH_OFFSET);
-
-                    for item in self {
-                        encoder.append(item);
-                    }
-
-                    encoder.finalize();
-                }
+                sequence_ssz_append(self.iter(), buf)
             }
         }
     };
@@ -269,6 +308,24 @@ impl_for_vec!(SmallVec<[T; 6]>);
 impl_for_vec!(SmallVec<[T; 7]>);
 impl_for_vec!(SmallVec<[T; 8]>);
 
+impl<K, V> Encode for BTreeMap<K, V>
+where
+    K: Encode + Ord,
+    V: Encode,
+{
+    fn is_ssz_fixed_len() -> bool {
+        false
+    }
+
+    fn ssz_bytes_len(&self) -> usize {
+        sequence_ssz_bytes_len(self.iter())
+    }
+
+    fn ssz_append(&self, buf: &mut Vec<u8>) {
+        sequence_ssz_append(self.iter(), buf)
+    }
+}
+
 impl Encode for bool {
     fn is_ssz_fixed_len() -> bool {
         true

consensus/ssz/tests/tests.rs

@@ -4,6 +4,8 @@ use ssz_derive::{Decode, Encode};
 
 mod round_trip {
     use super::*;
+    use std::collections::BTreeMap;
+    use std::iter::FromIterator;
 
     fn round_trip<T: Encode + Decode + std::fmt::Debug + PartialEq>(items: Vec<T>) {
         for item in items {
@@ -321,6 +323,52 @@ mod round_trip {
 
         round_trip(vec);
     }
+
+    #[test]
+    fn btree_map_fixed() {
+        let data = vec![
+            BTreeMap::new(),
+            BTreeMap::from_iter(vec![(0u8, 0u16), (1, 2), (2, 4), (4, 6)]),
+        ];
+        round_trip(data);
+    }
+
+    #[test]
+    fn btree_map_variable_value() {
+        let data = vec![
+            BTreeMap::new(),
+            BTreeMap::from_iter(vec![
+                (
+                    0u64,
+                    ThreeVariableLen {
+                        a: 1,
+                        b: vec![3, 5, 7],
+                        c: vec![],
+                        d: vec![0, 0],
+                    },
+                ),
+                (
+                    1,
+                    ThreeVariableLen {
+                        a: 99,
+                        b: vec![1],
+                        c: vec![2, 3, 4, 5, 6, 7, 8, 9, 10],
+                        d: vec![4, 5, 6, 7, 8],
+                    },
+                ),
+                (
+                    2,
+                    ThreeVariableLen {
+                        a: 0,
+                        b: vec![],
+                        c: vec![],
+                        d: vec![],
+                    },
+                ),
+            ]),
+        ];
+        round_trip(data);
+    }
 }
 
 mod derive_macro {