diff --git a/ssz/src/encode.rs b/ssz/src/encode.rs index 80d07a328..b94c36ca5 100644 --- a/ssz/src/encode.rs +++ b/ssz/src/encode.rs @@ -40,6 +40,13 @@ impl SszStream { self.buffer.extend_from_slice(&vec); } + /// Append some ssz encoded bytes to the stream without calculating length + /// + /// The raw bytes will be concatenated to the stream. + pub fn append_encoded_raw(&mut self, vec: &Vec) { + self.buffer.extend_from_slice(&vec); + } + /// Append some vector (list) of encodable values to the stream. /// /// The length of the list will be concatenated to the stream, then diff --git a/ssz/src/impl_encode.rs b/ssz/src/impl_encode.rs index 88b4c1edc..3299cbffe 100644 --- a/ssz/src/impl_encode.rs +++ b/ssz/src/impl_encode.rs @@ -1,8 +1,11 @@ +extern crate bytes; + use super::{ Encodable, SszStream }; use super::ethereum_types::{ H256, U256 }; +use self::bytes::{ BytesMut, BufMut }; /* * Note: there is a "to_bytes" function for integers @@ -10,40 +13,39 @@ use super::ethereum_types::{ H256, U256 }; * use it instead. */ macro_rules! impl_encodable_for_uint { - ($type: ident) => { + ($type: ident, $bit_size: expr) => { impl Encodable for $type { fn ssz_append(&self, s: &mut SszStream) { - // Number of bits required to represent this integer. - // This could be optimised at the expense of complexity. - let num_bits = { - let mut n = *self; - let mut r: usize = 0; - while n > 0 { - n >>= 1; - r += 1; - } - if r == 0 { 1 } else { r } - }; - // Number of bytes required to represent this bit - let num_bytes = (num_bits + 8 - 1) / 8; - let mut ssz_val: Vec = Vec::with_capacity(num_bytes); - ssz_val.resize(num_bytes, 0); - for i in (0..num_bytes).rev() { - let offset = (num_bytes - i - 1) * 8; - ssz_val[i] = 0_u8 | (self >> offset) as u8 + // Ensure bit size is valid + assert!((0 < $bit_size) && + ($bit_size % 8 == 0) && + (2_u128.pow($bit_size) > *self as u128)); + + // Serialize to bytes + let mut buf = BytesMut::with_capacity($bit_size/8); + + // Match bit size with encoding + match $bit_size { + 8 => buf.put_u8(*self as u8), + 16 => buf.put_u16_be(*self as u16), + 32 => buf.put_u32_be(*self as u32), + 64 => buf.put_u64_be(*self as u64), + _ => { ; } } - s.append_encoded_val(&ssz_val); + + // Append bytes to the SszStream + s.append_encoded_raw(&mut buf.to_vec()); } } } } -impl_encodable_for_uint!(u8); -impl_encodable_for_uint!(u16); -impl_encodable_for_uint!(u32); -impl_encodable_for_uint!(u64); -impl_encodable_for_uint!(usize); +impl_encodable_for_uint!(u8, 8); +impl_encodable_for_uint!(u16, 16); +impl_encodable_for_uint!(u32, 32); +impl_encodable_for_uint!(u64, 64); +impl_encodable_for_uint!(usize, 64); impl Encodable for H256 { fn ssz_append(&self, s: &mut SszStream) { @@ -66,25 +68,25 @@ mod tests { #[test] fn test_ssz_encode_u8() { - let x: u16 = 0; + let x: u8 = 0; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 1, 0]); + assert_eq!(ssz.drain(), vec![0]); - let x: u16 = 1; + let x: u8 = 1; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 1, 1]); + assert_eq!(ssz.drain(), vec![1]); - let x: u16 = 100; + let x: u8 = 100; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 1, 100]); + assert_eq!(ssz.drain(), vec![100]); - let x: u16 = 255; + let x: u8 = 255; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 1, 255]); + assert_eq!(ssz.drain(), vec![255]); } #[test] @@ -92,22 +94,22 @@ mod tests { let x: u16 = 1; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 1, 1]); + assert_eq!(ssz.drain(), vec![0, 1]); let x: u16 = 100; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 1, 100]); + assert_eq!(ssz.drain(), vec![0, 100]); let x: u16 = 1 << 8; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 2, 1, 0]); + assert_eq!(ssz.drain(), vec![1, 0]); let x: u16 = 65535; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 2, 255, 255]); + assert_eq!(ssz.drain(), vec![255, 255]); } #[test] @@ -115,27 +117,27 @@ mod tests { let x: u32 = 1; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 1, 1]); + assert_eq!(ssz.drain(), vec![0, 0, 0, 1]); let x: u32 = 100; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 1, 100]); + assert_eq!(ssz.drain(), vec![0, 0, 0, 100]); let x: u32 = 1 << 16; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 3, 1, 0, 0]); + assert_eq!(ssz.drain(), vec![0, 1, 0, 0]); let x: u32 = 1 << 24; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 4, 1, 0, 0, 0]); + assert_eq!(ssz.drain(), vec![1, 0, 0, 0]); let x: u32 = !0; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 4, 255, 255, 255, 255]); + assert_eq!(ssz.drain(), vec![255, 255, 255, 255]); } #[test] @@ -143,22 +145,22 @@ mod tests { let x: u64 = 1; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 1, 1]); + assert_eq!(ssz.drain(), vec![0, 0, 0, 0, 0, 0, 0, 1]); let x: u64 = 100; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 1, 100]); + assert_eq!(ssz.drain(), vec![0, 0, 0, 0, 0, 0, 0, 100]); let x: u64 = 1 << 32; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 5, 1, 0, 0, 0, 0]); + assert_eq!(ssz.drain(), vec![0, 0, 0, 1, 0, 0, 0, 0]); let x: u64 = !0; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 8, 255, 255, 255, 255, 255, 255, 255, 255]); + assert_eq!(ssz.drain(), vec![255, 255, 255, 255, 255, 255, 255, 255]); } #[test] @@ -166,21 +168,21 @@ mod tests { let x: usize = 1; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 1, 1]); + assert_eq!(ssz.drain(), vec![0, 0, 0, 0, 0, 0, 0, 1]); let x: usize = 100; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 1, 100]); + assert_eq!(ssz.drain(), vec![0, 0, 0, 0, 0, 0, 0, 100]); let x: usize = 1 << 32; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 5, 1, 0, 0, 0, 0]); + assert_eq!(ssz.drain(), vec![0, 0, 0, 1, 0, 0, 0, 0]); let x: usize = !0; let mut ssz = SszStream::new(); ssz.append(&x); - assert_eq!(ssz.drain(), vec![0, 0, 8, 255, 255, 255, 255, 255, 255, 255, 255]); + assert_eq!(ssz.drain(), vec![255, 255, 255, 255, 255, 255, 255, 255]); } }