From bd021f80a324d03814fb064b076588f815444aa4 Mon Sep 17 00:00:00 2001 From: Simon Bernier St-Pierre Date: Sun, 11 Dec 2016 19:15:09 -0500 Subject: [PATCH] add utility methods to the bencode objects --- src/bencode/decode.rs | 16 ++-- src/bencode/encode.rs | 22 ++++-- src/bencode/mod.rs | 169 ++++++++++++++++++++++++++++++++++++++---- 3 files changed, 175 insertions(+), 32 deletions(-) diff --git a/src/bencode/decode.rs b/src/bencode/decode.rs index 8aeb852..c6134a0 100644 --- a/src/bencode/decode.rs +++ b/src/bencode/decode.rs @@ -2,7 +2,7 @@ use std::collections::BTreeMap; use std::num::ParseIntError; use std::str::{self, Utf8Error}; -use bencode::Object; +use bencode::{Bytes, Dict, List, Object}; use buffer::Buffer; #[derive(Debug)] @@ -46,7 +46,7 @@ fn decode_object(buf: &mut Buffer) -> DecodeResult { } list.push(decode_object(buf)?); } - Ok(Object::List(list)) + Ok(List::new(list)) } Some(b'd') => { buf.advance(1); @@ -58,9 +58,9 @@ fn decode_object(buf: &mut Buffer) -> DecodeResult { } let key = _decode_bytes(buf)?; let val = decode_object(buf)?; - dict.insert(key, val); + dict.insert(Bytes(key), val); } - Ok(Object::Dict(dict)) + Ok(Dict::new(dict)) } _ => Err(DecodeError), } @@ -84,7 +84,7 @@ fn _decode_int(buf: &mut Buffer, term: u8) -> DecodeResult { } fn decode_bytes(buf: &mut Buffer) -> DecodeResult { - _decode_bytes(buf).map(|bytes| Object::Bytes(bytes)) + _decode_bytes(buf).map(|bytes| Bytes::new(bytes)) } fn _decode_bytes(buf: &mut Buffer) -> Result, DecodeError> { @@ -116,7 +116,7 @@ fn test_int_neg() { fn test_bytes() { let mut buf = Buffer::new(b"5:hello"); - assert_eq!(decode_bytes(&mut buf).unwrap(), Object::Bytes(b"hello".to_vec())); + assert_eq!(decode_bytes(&mut buf).unwrap(), Bytes::new(b"hello".to_vec())); assert_eq!(buf.pos(), 7); } @@ -127,7 +127,7 @@ fn test_list() { let obj = decode_object(&mut buf).unwrap(); let list = obj.as_list().unwrap(); - assert_eq!(list, vec![Object::Int(1), Object::Int(2), Object::Int(3)]); + assert_eq!(list, &List(vec![Object::Int(1), Object::Int(2), Object::Int(3)])); assert_eq!(buf.pos(), 11); } @@ -138,6 +138,6 @@ fn test_dict() { let obj = decode_object(&mut buf).unwrap(); let dict = obj.as_dict().unwrap(); - assert_eq!(dict[&b"hello"[..]], Object::Int(1337)); + assert_eq!(dict.get_int("hello").unwrap(), 1337); assert_eq!(buf.pos(), 15); } diff --git a/src/bencode/encode.rs b/src/bencode/encode.rs index feb5635..a9592ee 100644 --- a/src/bencode/encode.rs +++ b/src/bencode/encode.rs @@ -1,7 +1,6 @@ use std::borrow::Borrow; -use std::collections::BTreeMap; -use bencode::Object; +use bencode::{Dict, Object}; pub fn encode(obj: B) -> Vec where B: Borrow { let mut buff = Vec::new(); @@ -35,9 +34,9 @@ fn encode_list(buff: &mut Vec, list: &[Object]) { buff.push(b'e'); } -fn encode_dict(buff: &mut Vec, dict: &BTreeMap, Object>) { +fn encode_dict(buff: &mut Vec, dict: &Dict) { buff.push(b'd'); - for (key, val) in dict { + for (key, val) in dict.iter() { encode_bytes(buff, key); encode_object(buff, val); } @@ -56,18 +55,25 @@ fn test_int_neg() { #[test] fn test_bytes() { - assert_eq!(encode(Object::Bytes(b"hello".to_vec())), b"5:hello"); + use bencode::Bytes; + + assert_eq!(encode(Bytes::new(b"hello".to_vec())), b"5:hello"); } #[test] fn test_list() { + use bencode::List; + let list = vec![Object::Int(1), Object::Int(2), Object::Int(3)]; - assert_eq!(encode(Object::List(list)), b"li1ei2ei3ee"); + assert_eq!(encode(List::new(list)), b"li1ei2ei3ee"); } #[test] fn test_dict() { + use std::collections::BTreeMap; + use bencode::{Bytes, Dict}; + let mut dict = BTreeMap::new(); - dict.insert(b"hello".to_vec(), Object::Int(1337)); - assert_eq!(encode(Object::Dict(dict)), b"d5:helloi1337ee") + dict.insert(Bytes(b"hello".to_vec()), Object::Int(1337)); + assert_eq!(encode(Dict::new(dict)), b"d5:helloi1337ee") } diff --git a/src/bencode/mod.rs b/src/bencode/mod.rs index 94493c2..2d8847b 100644 --- a/src/bencode/mod.rs +++ b/src/bencode/mod.rs @@ -1,46 +1,56 @@ mod decode; mod encode; +use std::borrow::Cow; use std::collections::BTreeMap; use std::fmt; +use std::ops::{Deref, DerefMut}; use std::str; pub use self::decode::{decode, DecodeError, DecodeResult}; pub use self::encode::encode; -#[derive(Eq, PartialEq)] +#[derive(Clone, Eq, PartialEq)] pub enum Object { Int(i64), - Bytes(Vec), - List(Vec), - Dict(BTreeMap, Object>), + Bytes(Bytes), + List(List), + Dict(Dict), } impl Object { - pub fn as_int(self) -> Option { - match self { + pub fn as_int(&self) -> Option { + match *self { Object::Int(num) => Some(num), _ => None, } } - pub fn as_bytes(self) -> Option> { - match self { - Object::Bytes(bytes) => Some(bytes), + pub fn as_bytes(&self) -> Option<&Bytes> { + match *self { + Object::Bytes(ref bytes) => Some(bytes), _ => None, } } - pub fn as_list(self) -> Option> { - match self { - Object::List(list) => Some(list), + pub fn as_str(&self) -> Option<&str> { + self.as_bytes().and_then(|b| str::from_utf8(b).ok()) + } + + pub fn as_string(&self) -> Option { + self.as_str().map(|s| s.to_string()) + } + + pub fn as_list(&self) -> Option<&List> { + match *self { + Object::List(ref list) => Some(list), _ => None, } } - pub fn as_dict(self) -> Option, Object>> { - match self { - Object::Dict(dict) => Some(dict), + pub fn as_dict(&self) -> Option<&Dict> { + match *self { + Object::Dict(ref dict) => Some(dict), _ => None, } } @@ -66,9 +76,136 @@ impl fmt::Debug for Object { } Object::Dict(ref dict) => { f.debug_map() - .entries(dict.iter().map(|(k, v)| (Object::Bytes(k.clone()), v))) + .entries(dict.iter().map(|(k, v)| (k.clone(), v))) .finish() } } } } + +#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)] +pub struct Bytes(Vec); + +impl Bytes { + pub fn new(bytes: Vec) -> Object { + Object::Bytes(Bytes(bytes)) + } + + pub fn str(&self) -> Option<&str> { + str::from_utf8(&self.0).ok() + } + + pub fn string(&self) -> Option { + self.str().map(|s| s.to_string()) + } +} + +impl Deref for Bytes { + type Target = Vec; + + fn deref(&self) -> &Vec { + &self.0 + } +} + +impl DerefMut for Bytes { + fn deref_mut(&mut self) -> &mut Vec { + &mut self.0 + } +} + +pub trait AsBytes { + fn as_bytes(&self) -> Cow; +} + +impl AsBytes for Bytes { + fn as_bytes(&self) -> Cow { + Cow::Borrowed(self) + } +} + +impl<'a> AsBytes for &'a str { + fn as_bytes(&self) -> Cow { + Cow::Owned(Bytes(str::as_bytes(self).to_vec())) + } +} + +impl<'a> AsBytes for &'a [u8] { + fn as_bytes(&self) -> Cow { + Cow::Owned(Bytes(self.to_vec())) + } +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct List(Vec); + +impl List { + pub fn new(list: Vec) -> Object { + Object::List(List(list)) + } +} + +impl Deref for List { + type Target = Vec; + + fn deref(&self) -> &Vec { + &self.0 + } +} + +impl DerefMut for List { + fn deref_mut(&mut self) -> &mut Vec { + &mut self.0 + } +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct Dict(BTreeMap); + +impl Dict { + pub fn new(dict: BTreeMap) -> Object { + Object::Dict(Dict(dict)) + } + + pub fn get_object(&self, key: A) -> Option<&Object> where A: AsBytes { + self.0.get(&key.as_bytes()) + } + + pub fn get_int(&self, key: A) -> Option where A: AsBytes { + self.0.get(&key.as_bytes()).and_then(|o| o.as_int()) + } + + pub fn get_bytes(&self, key: A) -> Option<&Bytes> where A: AsBytes { + self.0.get(&key.as_bytes()).and_then(|o| o.as_bytes()) + } + + pub fn get_str(&self, key: A) -> Option<&str> where A: AsBytes { + self.0.get(&key.as_bytes()).and_then(|o| o.as_bytes().and_then(|b| str::from_utf8(b).ok())) + } + + pub fn get_string(&self, key: A) -> Option where A: AsBytes { + self.get_str(key).map(|s| s.to_string()) + } + + pub fn get_list(&self, key: A) -> Option<&List> where A: AsBytes { + self.0.get(&key.as_bytes()).and_then(|o| o.as_list()) + } + + pub fn get_dict(&self, key: A) -> Option<&Dict> where A: AsBytes { + self.0.get(&key.as_bytes()).and_then(|o| o.as_dict()) + } +} + +impl Deref for Dict { + type Target = BTreeMap; + + fn deref(&self) -> &BTreeMap { + &self.0 + } +} + +impl DerefMut for Dict { + fn deref_mut(&mut self) -> &mut BTreeMap { + &mut self.0 + } +}