Files
bt/src/bencode/encode.rs

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")
}