Update ssz implementation of int types

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<u8>) {
+        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

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;
+                // 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),
+                    _ => { ; }
                 }
-                    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<u8> = 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
-                }
-                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]);
     }
 }