Add support for SSZ union type via Option

eth2/utils/ssz/src/decode.rs

@@ -218,6 +218,12 @@ impl<'a> SszDecoder<'a> {
     }
 }
 
+/// Reads a `BYTES_PER_LENGTH_OFFSET`-byte union index from `bytes`, where `bytes.len() >=
+/// BYTES_PER_LENGTH_OFFSET`.
+pub fn read_union_index(bytes: &[u8]) -> Result<usize, DecodeError> {
+    read_offset(bytes)
+}
+
 /// Reads a `BYTES_PER_LENGTH_OFFSET`-byte length from `bytes`, where `bytes.len() >=
 /// BYTES_PER_LENGTH_OFFSET`.
 fn read_offset(bytes: &[u8]) -> Result<usize, DecodeError> {

eth2/utils/ssz/src/decode/impls.rs

@@ -62,6 +62,36 @@ impl Decode for bool {
     }
 }
 
+/// The SSZ union type.
+impl<T: Decode> Decode for Option<T> {
+    fn is_ssz_fixed_len() -> bool {
+        false
+    }
+
+    fn from_ssz_bytes(bytes: &[u8]) -> Result<Self, DecodeError> {
+        if bytes.len() < BYTES_PER_LENGTH_OFFSET {
+            return Err(DecodeError::InvalidByteLength {
+                len: bytes.len(),
+                expected: BYTES_PER_LENGTH_OFFSET,
+            });
+        }
+
+        let (index_bytes, value_bytes) = bytes.split_at(BYTES_PER_LENGTH_OFFSET);
+
+        let index = read_union_index(index_bytes)?;
+        if index == 0 {
+            Ok(None)
+        } else if index == 1 {
+            Ok(Some(T::from_ssz_bytes(value_bytes)?))
+        } else {
+            Err(DecodeError::BytesInvalid(format!(
+                "{} is not a valid union index for Option<T>",
+                index
+            )))
+        }
+    }
+}
+
 impl Decode for H256 {
     fn is_ssz_fixed_len() -> bool {
         true

eth2/utils/ssz/src/encode.rs

@@ -126,6 +126,13 @@ impl<'a> SszEncoder<'a> {
     }
 }
 
+/// Encode `index` as a little-endian byte vec of `BYTES_PER_LENGTH_OFFSET` length.
+///
+/// If `len` is larger than `2 ^ BYTES_PER_LENGTH_OFFSET`, a `debug_assert` is raised.
+pub fn encode_union_index(index: usize) -> Vec<u8> {
+    encode_length(index)
+}
+
 /// Encode `len` as a little-endian byte vec of `BYTES_PER_LENGTH_OFFSET` length.
 ///
 /// If `len` is larger than `2 ^ BYTES_PER_LENGTH_OFFSET`, a `debug_assert` is raised.

eth2/utils/ssz/src/encode/impls.rs

@@ -25,6 +25,23 @@ impl_encodable_for_uint!(u32, 32);
 impl_encodable_for_uint!(u64, 64);
 impl_encodable_for_uint!(usize, 64);
 
+/// The SSZ "union" type.
+impl<T: Encode> Encode for Option<T> {
+    fn is_ssz_fixed_len() -> bool {
+        false
+    }
+
+    fn ssz_append(&self, buf: &mut Vec<u8>) {
+        match self {
+            None => buf.append(&mut encode_union_index(0)),
+            Some(t) => {
+                buf.append(&mut encode_union_index(1));
+                t.ssz_append(buf);
+            }
+        }
+    }
+}
+
 impl<T: Encode> Encode for Vec<T> {
     fn is_ssz_fixed_len() -> bool {
         false
@@ -168,6 +185,25 @@ mod tests {
         );
     }
 
+    #[test]
+    fn ssz_encode_option_u16() {
+        assert_eq!(Some(65535_u16).as_ssz_bytes(), vec![1, 0, 0, 0, 255, 255]);
+
+        let none: Option<u16> = None;
+        assert_eq!(none.as_ssz_bytes(), vec![0, 0, 0, 0]);
+    }
+
+    #[test]
+    fn ssz_encode_option_vec_u16() {
+        assert_eq!(
+            Some(vec![0_u16, 1]).as_ssz_bytes(),
+            vec![1, 0, 0, 0, 0, 0, 1, 0]
+        );
+
+        let none: Option<Vec<u16>> = None;
+        assert_eq!(none.as_ssz_bytes(), vec![0, 0, 0, 0]);
+    }
+
     #[test]
     fn ssz_encode_u8() {
         assert_eq!(0_u8.as_ssz_bytes(), vec![0]);