93 lines
2.1 KiB
Rust
93 lines
2.1 KiB
Rust
use std::borrow::Cow;
|
|
|
|
use bencode::{Dict, Object};
|
|
|
|
/// Allows the encode function to take something that can be converted into an object,
|
|
/// or a reference to an object.
|
|
pub trait Encodable<'a> {
|
|
fn get_object(self) -> Cow<'a, Object>;
|
|
}
|
|
|
|
impl<'a, T> Encodable<'a> for T where T: Into<Object> {
|
|
fn get_object(self) -> Cow<'a, Object> {
|
|
Cow::Owned(self.into())
|
|
}
|
|
}
|
|
|
|
impl<'a> Encodable<'a> for &'a Object {
|
|
fn get_object(self) -> Cow<'a, Object> {
|
|
Cow::Borrowed(self)
|
|
}
|
|
}
|
|
|
|
pub fn encode<'a, O>(obj: O) -> Vec<u8> where O: Encodable<'a> {
|
|
let mut buff = Vec::new();
|
|
encode_object(&mut buff, &obj.get_object());
|
|
buff
|
|
}
|
|
|
|
fn encode_object(buff: &mut Vec<u8>, obj: &Object) {
|
|
match *obj {
|
|
Object::Int(num) => encode_int(buff, num),
|
|
Object::Bytes(ref bytes) => encode_bytes(buff, bytes),
|
|
Object::List(ref list) => encode_list(buff, list),
|
|
Object::Dict(ref dict) => encode_dict(buff, dict),
|
|
}
|
|
}
|
|
|
|
fn encode_int(buff: &mut Vec<u8>, num: i64) {
|
|
buff.extend_from_slice(format!("i{}e", num).as_bytes());
|
|
}
|
|
|
|
fn encode_bytes(buff: &mut Vec<u8>, bytes: &[u8]) {
|
|
buff.extend_from_slice(format!("{}:", bytes.len()).as_bytes());
|
|
buff.extend_from_slice(bytes);
|
|
}
|
|
|
|
fn encode_list(buff: &mut Vec<u8>, list: &[Object]) {
|
|
buff.push(b'l');
|
|
for obj in list {
|
|
encode_object(buff, obj);
|
|
}
|
|
buff.push(b'e');
|
|
}
|
|
|
|
fn encode_dict(buff: &mut Vec<u8>, dict: &Dict) {
|
|
buff.push(b'd');
|
|
for (key, val) in dict.iter() {
|
|
encode_bytes(buff, key);
|
|
encode_object(buff, val);
|
|
}
|
|
buff.push(b'e');
|
|
}
|
|
|
|
#[test]
|
|
fn test_int_pos() {
|
|
assert_eq!(encode(1337), b"i1337e");
|
|
}
|
|
|
|
#[test]
|
|
fn test_int_neg() {
|
|
assert_eq!(encode(-1337), b"i-1337e");
|
|
}
|
|
|
|
#[test]
|
|
fn test_bytes() {
|
|
assert_eq!(encode("hello"), b"5:hello");
|
|
}
|
|
|
|
#[test]
|
|
fn test_list() {
|
|
let list = vec![Object::Int(1), Object::Int(2), Object::Int(3)];
|
|
assert_eq!(encode(list), b"li1ei2ei3ee");
|
|
}
|
|
|
|
#[test]
|
|
fn test_dict() {
|
|
use bencode::Dict;
|
|
|
|
let mut dict = Dict::new();
|
|
dict.insert("hello", 1337);
|
|
assert_eq!(encode(dict), b"d5:helloi1337ee")
|
|
}
|