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]);

eth2/utils/ssz/tests/tests.rs

@@ -276,7 +276,7 @@ mod round_trip {
     fn offsets_decreasing() {
         let bytes = vec![
             //  1   2   3   4   5   6   7   8   9   10  11  12  13  14  15
-            //      | offset        | ofset         | offset        | variable
+            //      | offset        | offset        | offset        | variable
             01, 00, 14, 00, 00, 00, 15, 00, 00, 00, 14, 00, 00, 00, 00, 00,
         ];
 
@@ -285,4 +285,65 @@ mod round_trip {
             Err(DecodeError::OutOfBoundsByte { i: 14 })
         );
     }
+
+    #[derive(Debug, PartialEq, Encode, Decode)]
+    struct TwoVariableLenOptions {
+        a: u16,
+        b: Option<u16>,
+        c: Option<Vec<u16>>,
+        d: Option<Vec<u16>>,
+    }
+
+    #[test]
+    fn two_variable_len_options_encoding() {
+        let s = TwoVariableLenOptions {
+            a: 42,
+            b: None,
+            c: Some(vec![0]),
+            d: None,
+        };
+
+        let bytes = vec![
+            //  1   2   3   4   5   6   7   8   9   10  11  12  13  14  15  16  17  18  19  20  21
+            //      | option<u16>   | offset        | offset        | option<u16    | 1st list
+            42, 00, 14, 00, 00, 00, 18, 00, 00, 00, 24, 00, 00, 00, 00, 00, 00, 00, 01, 00, 00, 00,
+            //  23  24  25  26  27
+            //      | 2nd list
+            00, 00, 00, 00, 00, 00
+        ];
+
+        assert_eq!(s.as_ssz_bytes(), bytes);
+    }
+
+    #[test]
+    fn two_variable_len_options_round_trip() {
+        let vec: Vec<TwoVariableLenOptions> = vec![
+            TwoVariableLenOptions {
+                a: 42,
+                b: Some(12),
+                c: Some(vec![0]),
+                d: Some(vec![1]),
+            },
+            TwoVariableLenOptions {
+                a: 42,
+                b: Some(12),
+                c: Some(vec![0]),
+                d: None,
+            },
+            TwoVariableLenOptions {
+                a: 42,
+                b: None,
+                c: Some(vec![0]),
+                d: None,
+            },
+            TwoVariableLenOptions {
+                a: 42,
+                b: None,
+                c: None,
+                d: None,
+            },
+        ];
+
+        round_trip(vec);
+    }
 }