Merge pull request #1 from gavofyork/rlp_encoding

rlp encoding

src/rlp.rs

@@ -2,8 +2,9 @@
 
 use std::fmt;
 use std::cell::Cell;
+use std::collections::LinkedList;
 use std::error::Error as StdError;
-use bytes::{FromBytes, FromBytesError};
+use bytes::{ToBytes, FromBytes, FromBytesError};
 
 /// rlp container
 #[derive(Debug)]
@@ -41,7 +42,7 @@ impl ItemInfo {
 pub enum DecoderError {
     FromBytesError(FromBytesError),
     RlpIsTooShort,
-    RlpExpectedToBeArray,
+    RlpExpectedToBeList,
     BadRlp,
 }
 impl StdError for DecoderError {
@@ -71,16 +72,16 @@ impl <'a>Rlp<'a> {
     ///
     /// paren container caches searched position
     pub fn at(&self, index: usize) -> Result<Rlp<'a>, DecoderError> {
-        if !self.is_array() {
+        if !self.is_list() {
-            return Err(DecoderError::RlpExpectedToBeArray);
+            return Err(DecoderError::RlpExpectedToBeList);
         }
 
         // move to cached position if it's index is less or equal to
-        // current search index, otherwise move to beginning of array
+        // current search index, otherwise move to beginning of list
         let c = self.cache.get();
         let (mut bytes, to_skip) = match c.index <= index {
             true => (try!(Rlp::consume(self.bytes, c.offset)), index - c.index),
-            false => (try!(self.consume_array_prefix()), index)
+            false => (try!(self.consume_list_prefix()), index)
         };
 
         // skip up to x items
@@ -94,8 +95,8 @@ impl <'a>Rlp<'a> {
         Ok(Rlp::new(&bytes[0..found.prefix_len + found.value_len]))
     }
 
-    /// returns true if rlp is an array
+    /// returns true if rlp is an list
-    pub fn is_array(&self) -> bool {
+    pub fn is_list(&self) -> bool {
         self.bytes.len() > 0 && self.bytes[0] >= 0xc0
     }
 
@@ -110,7 +111,7 @@ impl <'a>Rlp<'a> {
     }
 
     /// consumes first found prefix
-    fn consume_array_prefix(&self) -> Result<&'a [u8], DecoderError> {
+    fn consume_list_prefix(&self) -> Result<&'a [u8], DecoderError> {
         let item = try!(Rlp::item_info(self.bytes));
         let bytes = try!(Rlp::consume(self.bytes, item.prefix_len));
         Ok(bytes)
@@ -189,17 +190,246 @@ impl <'a> Iterator for RlpIterator<'a> {
     }
 }
 
+#[derive(Debug)]
+struct ListInfo {
+    position: usize,
+    current: usize,
+    max: usize
+}
+
+impl ListInfo {
+    fn new(position: usize, max: usize) -> ListInfo {
+        ListInfo { 
+            position: position,
+            current: 0,
+            max: max
+        }
+    }
+}
+
+/// container that should be used to encode rlp
+pub struct RlpStream {
+    unfinished_lists: LinkedList<ListInfo>,
+    encoder: BasicEncoder
+}
+
+impl RlpStream {
+    /// create new container for values appended one after another,
+    /// but not being part of the same list
+    pub fn new() -> RlpStream {
+        RlpStream {
+            unfinished_lists: LinkedList::new(),
+            encoder: BasicEncoder::new()
+        }
+    }
+
+    /// create new container for list of size `max_len`
+    pub fn new_list(len: usize) -> RlpStream {
+        let mut stream = RlpStream::new();
+        stream.append_list(len);
+        stream
+    }
+
+    /// apends value to the end of stream, chainable
+    pub fn append<'a, E>(&'a mut self, object: &E) -> &'a mut RlpStream where E: Encodable {
+        // encode given value and add it at the end of the stream
+        object.encode(&mut self.encoder);
+
+        // if list is finished, prepend the length
+        self.try_to_finish();
+
+        // return chainable self
+        self
+    }
+
+    /// declare appending the list of given size
+    pub fn append_list<'a>(&'a mut self, len: usize) -> &'a mut RlpStream {
+        // push new list
+        let position = self.encoder.bytes.len();
+        match len {
+            0 => {
+                // we may finish, if the appended list len is equal 0
+                self.encoder.insert_list_len_at_pos(0, position);
+                self.try_to_finish();
+            },
+            _ => self.unfinished_lists.push_back(ListInfo::new(position, len))
+        }
+
+        // return chainable self
+        self
+    }
+
+    /// return true if stream is ready
+    pub fn is_finished(&self) -> bool {
+        self.unfinished_lists.back().is_none()
+    }
+
+    /// streams out encoded bytes
+    pub fn out(self) -> Result<Vec<u8>, EncoderError> {
+        match self.is_finished() {
+            true => Ok(self.encoder.out()),
+            false => Err(EncoderError::StreamIsUnfinished)
+        }
+    }
+
+    /// try to finish lists
+    fn try_to_finish(&mut self) -> () {
+        let should_finish = match self.unfinished_lists.back_mut() {
+            None => false,
+            Some(ref mut x) => {
+                x.current += 1;
+                x.current == x.max
+            }
+        };
+
+        if should_finish {    
+            let x = self.unfinished_lists.pop_back().unwrap();
+            let len = self.encoder.bytes.len() - x.position;
+            self.encoder.insert_list_len_at_pos(len, x.position);
+            self.try_to_finish();
+        }
+    }
+}
+
+/// shortcut function to encode a `T: Encodable` into a Rlp `Vec<u8>`
+pub fn encode<E>(object: &E) -> Vec<u8> where E: Encodable {
+    let mut encoder = BasicEncoder::new();
+    object.encode(&mut encoder);
+    encoder.out()
+}
+
+#[derive(Debug)]
+pub enum EncoderError {
+    StreamIsUnfinished
+}
+
+impl StdError for EncoderError {
+    fn description(&self) -> &str { "encoder error" }
+}
+
+impl fmt::Display for EncoderError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        fmt::Debug::fmt(&self, f)
+    }
+}
+
+pub trait Encodable {
+    fn encode<E>(&self, encoder: &mut E) -> () where E: Encoder;
+}
+
+pub trait Encoder {
+    fn emit_value(&mut self, bytes: &[u8]) -> ();
+    fn emit_list<F>(&mut self, f: F) -> () where F: FnOnce(&mut Self) -> ();
+}
+
+impl <T> Encodable for T where T: ToBytes {
+    fn encode<E>(&self, encoder: &mut E) -> () where E: Encoder {
+        encoder.emit_value(&self.to_bytes())
+    }
+}
+
+impl <'a, T> Encodable for &'a [T] where T: Encodable + 'a {
+    fn encode<E>(&self, encoder: &mut E) -> () where E: Encoder {
+        encoder.emit_list(|e| {
+            // insert all list elements
+            for el in self.iter() {
+                el.encode(e);
+            }
+        })
+    }
+}
+
+impl <T> Encodable for Vec<T> where T: Encodable {
+    fn encode<E>(&self, encoder: &mut E) -> () where E: Encoder {
+        let r: &[T] = self.as_ref();
+        r.encode(encoder)
+    }
+}
+
+struct BasicEncoder {
+    bytes: Vec<u8>
+}
+
+impl BasicEncoder {
+    fn new() -> BasicEncoder {
+        BasicEncoder { bytes: vec![] }
+    }
+
+    /// inserts list prefix at given position
+    fn insert_list_len_at_pos(&mut self, len: usize, pos: usize) -> () {
+        // new bytes
+        let mut res: Vec<u8> = vec![];
+        {
+            let (before_slice, after_slice) = self.bytes.split_at(pos); 
+            res.extend(before_slice);
+
+            match len {
+                0...55 => res.push(0xc0u8 + len as u8),
+                _ => {
+                    res.push(0x7fu8 + len.to_bytes_len() as u8);
+                    res.extend(len.to_bytes());
+                }
+            };
+
+            res.extend(after_slice);
+        }
+        self.bytes = res;
+    }
+
+    /// get encoded value
+    fn out(self) -> Vec<u8> {
+        self.bytes
+    }
+}
+
+impl Encoder for BasicEncoder {
+    fn emit_value(&mut self, bytes: &[u8]) -> () {
+        match bytes.len() {
+            // just 0
+            0 => self.bytes.push(0x80u8),
+            // byte is its own encoding
+            1 if bytes[0] < 0x80 => self.bytes.extend(bytes),
+            // (prefix + length), followed by the string
+            len @ 1 ... 55 => {
+                self.bytes.push(0x80u8 + len as u8);
+                self.bytes.extend(bytes);
+            }
+            // (prefix + length of length), followed by the length, followd by the string
+            len => {
+                self.bytes.push(0xb7 + len.to_bytes_len() as u8);
+                self.bytes.extend(len.to_bytes());
+                self.bytes.extend(bytes);
+            }
+        }
+    }
+
+    fn emit_list<F>(&mut self, f: F) -> () where F: FnOnce(&mut Self) -> () {
+        // get len before inserting an list
+        let before_len = self.bytes.len();
+
+        // insert all list elements
+        f(self);
+
+        // get len after inserting an list
+        let after_len = self.bytes.len();
+
+        // diff is list len
+        let list_len = after_len - before_len;
+        self.insert_list_len_at_pos(list_len, before_len);
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use rlp;
-    use rlp::Rlp;
+    use rlp::{Rlp, RlpStream};
 
     #[test]
     fn rlp_at() {
         let data = vec![0xc8, 0x83, b'c', b'a', b't', 0x83, b'd', b'o', b'g'];
         {
             let rlp = Rlp::new(&data);
-            assert!(rlp.is_array());
+            assert!(rlp.is_list());
            
             let cat = rlp.at(0).unwrap();
             assert!(cat.is_value());
@@ -220,7 +450,7 @@ mod tests {
         let data = vec![0xc8, 0x83, b'c', b'a', b't', 0x83, b'd', b'o'];
         {
             let rlp = Rlp::new(&data);
-            assert!(rlp.is_array());
+            assert!(rlp.is_list());
 
             let cat_err = rlp.at(0).unwrap_err();
             assert_eq!(cat_err, rlp::DecoderError::RlpIsTooShort);
@@ -253,5 +483,131 @@ mod tests {
             assert_eq!(cat_again.bytes, &[0x83, b'c', b'a', b't']);
         }
     }
+
+    struct ETestPair<T>(T, Vec<u8>) where T: rlp::Encodable;
+
+    fn run_encode_tests<T>(tests: Vec<ETestPair<T>>) where T: rlp::Encodable {
+        for t in &tests {
+            let res = rlp::encode(&t.0);
+            assert_eq!(res, &t.1[..]);
+        }
+    }
+
+    #[test]
+    fn encode_u8() {
+        let tests = vec![
+            ETestPair(0u8, vec![0x80u8]),
+            ETestPair(15, vec![15]),
+            ETestPair(55, vec![55]),
+            ETestPair(56, vec![56]),
+            ETestPair(0x7f, vec![0x7f]),
+            ETestPair(0x80, vec![0x81, 0x80]),
+            ETestPair(0xff, vec![0x81, 0xff]),
+        ];
+        run_encode_tests(tests);
+    }
+
+    #[test]
+    fn encode_u16() {
+        let tests = vec![
+            ETestPair(0u16, vec![0x80u8]),
+            ETestPair(0x100, vec![0x82, 0x01, 0x00]),
+            ETestPair(0xffff, vec![0x82, 0xff, 0xff]),
+        ];
+        run_encode_tests(tests);
+    }
+
+    #[test]
+    fn encode_u32() {
+        let tests = vec![
+            ETestPair(0u32, vec![0x80u8]),
+            ETestPair(0x10000, vec![0x83, 0x01, 0x00, 0x00]),
+            ETestPair(0xffffff, vec![0x83, 0xff, 0xff, 0xff]),
+        ];
+        run_encode_tests(tests);
+    }
+
+    #[test]
+    fn encode_u64() {
+        let tests = vec![
+            ETestPair(0u64, vec![0x80u8]),
+            ETestPair(0x1000000, vec![0x84, 0x01, 0x00, 0x00, 0x00]),
+            ETestPair(0xFFFFFFFF, vec![0x84, 0xff, 0xff, 0xff, 0xff]),
+        ];
+        run_encode_tests(tests);
+    }
+
+    #[test]
+    fn encode_str() {
+        let tests = vec![
+            ETestPair("cat", vec![0x83, b'c', b'a', b't']),
+            ETestPair("dog", vec![0x83, b'd', b'o', b'g']),
+            ETestPair("Marek", vec![0x85, b'M', b'a', b'r', b'e', b'k']),
+            ETestPair("", vec![0x80]),
+            ETestPair("Lorem ipsum dolor sit amet, consectetur adipisicing elit",
+                     vec![0xb8, 0x38, b'L', b'o', b'r', b'e', b'm', b' ', b'i',
+                    b'p', b's', b'u', b'm', b' ', b'd', b'o', b'l', b'o', b'r',
+                    b' ', b's', b'i', b't', b' ', b'a', b'm', b'e', b't', b',',
+                    b' ', b'c', b'o', b'n', b's', b'e', b'c', b't', b'e', b't',
+                    b'u', b'r', b' ', b'a', b'd', b'i', b'p', b'i', b's', b'i',
+                    b'c', b'i', b'n', b'g', b' ', b'e', b'l', b'i', b't'])
+        ];
+        run_encode_tests(tests);
+    }
+
+    #[test]
+    fn encode_vector_u8() {
+        let tests = vec![
+            ETestPair(vec![], vec![0xc0]),
+            ETestPair(vec![15u8], vec![0xc1, 0x0f]),
+            ETestPair(vec![1, 2, 3, 7, 0xff], vec![0xc6, 1, 2, 3, 7, 0x81, 0xff]),
+        ];
+        run_encode_tests(tests);
+    }
+
+    #[test]
+    fn encode_vector_u64() {
+        let tests = vec![
+            ETestPair(vec![], vec![0xc0]),
+            ETestPair(vec![15u64], vec![0xc1, 0x0f]),
+            ETestPair(vec![1, 2, 3, 7, 0xff], vec![0xc6, 1, 2, 3, 7, 0x81, 0xff]),
+            ETestPair(vec![0xffffffff, 1, 2, 3, 7, 0xff], vec![0xcb, 0x84, 0xff, 0xff, 0xff, 0xff,  1, 2, 3, 7, 0x81, 0xff]),
+        ];
+        run_encode_tests(tests);
+    }
+
+    #[test]
+    fn encode_vector_str() {
+        let tests = vec![
+            ETestPair(vec!["cat", "dog"], vec![0xc8, 0x83, b'c', b'a', b't', 0x83, b'd', b'o', b'g'])
+        ];
+        run_encode_tests(tests);
+    }
+
+    #[test]
+    fn encode_vector_of_vectors_str() {
+        let tests = vec![
+            ETestPair(vec![vec!["cat"]], vec![0xc5, 0xc4, 0x83, b'c', b'a', b't'])
+        ];
+        run_encode_tests(tests);
+    }
+
+    #[test]
+    fn rlp_stream() {
+        let mut stream = RlpStream::new_list(2);
+        stream.append(&"cat").append(&"dog");
+        let out = stream.out().unwrap();
+        assert_eq!(out, vec![0xc8, 0x83, b'c', b'a', b't', 0x83, b'd', b'o', b'g']);
+    }
+
+    #[test]
+    fn rlp_stream_list() {
+        let mut stream = RlpStream::new_list(3);
+        stream.append_list(0);
+        stream.append_list(1).append(&vec![] as &Vec<u8>);
+        stream.append_list(2).append_list(0).append_list(1).append_list(0);
+        let out = stream.out().unwrap();
+        assert_eq!(out, vec![0xc7, 0xc0, 0xc1, 0xc0, 0xc3, 0xc0, 0xc1, 0xc0]);
+    }
 }