diff --git a/eth2/utils/ssz/src/decode/impls.rs b/eth2/utils/ssz/src/decode/impls.rs index 7f1ad26ee4..7dd7fac3c1 100644 --- a/eth2/utils/ssz/src/decode/impls.rs +++ b/eth2/utils/ssz/src/decode/impls.rs @@ -1,5 +1,6 @@ use super::*; use ethereum_types::{H256, U128, U256}; +use core::num::NonZeroUsize; macro_rules! impl_decodable_for_uint { ($type: ident, $bit_size: expr) => { @@ -62,6 +63,28 @@ impl Decode for bool { } } +impl Decode for NonZeroUsize { + fn is_ssz_fixed_len() -> bool { + ::is_ssz_fixed_len() + } + + fn ssz_fixed_len() -> usize { + ::ssz_fixed_len() + } + + fn from_ssz_bytes(bytes: &[u8]) -> Result { + let x = usize::from_ssz_bytes(bytes)?; + + if x == 0 { + Err(DecodeError::BytesInvalid("NonZeroUsize cannot be zero.".to_string())) + } else { + // `unwrap` is safe here as `NonZeroUsize::new()` succeeds if `x > 0` and this path + // never executes when `x == 0`. + Ok(NonZeroUsize::new(x).unwrap()) + } + } +} + /// The SSZ union type. impl Decode for Option { fn is_ssz_fixed_len() -> bool { diff --git a/eth2/utils/ssz/src/encode/impls.rs b/eth2/utils/ssz/src/encode/impls.rs index 8dbcf17074..04492a1f2d 100644 --- a/eth2/utils/ssz/src/encode/impls.rs +++ b/eth2/utils/ssz/src/encode/impls.rs @@ -1,4 +1,5 @@ use super::*; +use core::num::NonZeroUsize; use ethereum_types::{H256, U128, U256}; macro_rules! impl_encodable_for_uint { @@ -80,6 +81,20 @@ impl Encode for bool { } } +impl Encode for NonZeroUsize { + fn is_ssz_fixed_len() -> bool { + ::is_ssz_fixed_len() + } + + fn ssz_fixed_len() -> usize { + ::ssz_fixed_len() + } + + fn ssz_append(&self, buf: &mut Vec) { + self.get().ssz_append(buf) + } +} + impl Encode for H256 { fn is_ssz_fixed_len() -> bool { true