1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139

//! A zero dependency Bencode parsing library.
//!
//! To get started try passing a bencoded byte array to `Benc::parse()`.
pub mod convert;
pub mod encode;
pub mod parse;
pub mod util;

use crate::encode::EncodeBencode;
use crate::parse::ParseResult;
use std::collections::HashMap;

/// The nested Bencode structure container.
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum Benc {
    Int(BInt),
    Str(BStr),
    Vec(BVec),
    Map(BMap),
}

impl Benc {
    /// parses the bencode object into a bencoded array of ascii bytes.
    pub fn bytes(&self) -> Vec<u8> {
        match self {
            Benc::Int(v) => (*v).bytes(),
            Benc::Str(v) => (*v).bytes(),
            Benc::Vec(v) => (*v).bytes(),
            Benc::Map(v) => (*v).bytes(),
        }
    }
    /// attempts to parse an array of ascii bytes into a bencode object.
    pub fn parse(bytes: &[u8]) -> ParseResult<Self> {
        match crate::util::parse_any(bytes) {
            Ok((val, _)) => Ok(val),
            Err(e) => Err(e),
        }
    }
}

/// A Bencode Integer, (format: "i[value]e")
pub type BInt = usize;
/// A Bencode ByteString, (format: "[length]:[bytes]")
pub type BStr = Vec<u8>;
/// A Bencode List of other Bencode types, (format: "l[bencode objects]e")
pub type BVec = Vec<Benc>;
/// A Bencode Dictionary of ByteString keys and any other bencode as values,
/// (format: "d[bencode items]e")
pub type BMap = BMapImpl;

/// Implementation: A HashMap wrapper that clones the contained Benc items when self.get() is called to avoid
/// causing a move on the un-Copy-able type Benc.
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct BMapImpl {
    // Saves the key and the index to the value in a lookup table, since usize is Copy-able.
    lookup: HashMap<BStr, usize>,
    // Saves the values here.
    values: Vec<Benc>,
}
impl BMap {
    pub fn get(&self, key: &BStr) -> Option<Benc> {
        let index: &usize = match self.lookup.get(key) {
            Some(index) => index,
            _ => return None,
        };
        Some(self.values.get(*index).unwrap().clone())
    }
    pub fn insert(&mut self, key: BStr, val: Benc) {
        let index = self.values.len();
        self.values.push(val);
        self.lookup.insert(key, index);
    }
    pub fn contains_key(&self, key: &BStr) -> bool {
        self.lookup.get(key).is_some()
    }
    pub fn new() -> Self {
        BMap {
            lookup: HashMap::new(),
            values: Vec::new(),
        }
    }
    pub fn len(&self) -> usize {
        self.values.len()
    }
    pub fn items(&self) -> Vec<(BStr, Benc)> {
        let mut keys: Vec<&BStr> = self.lookup.keys().collect();
        // We sort here to keep a stable order for the list. This is also part of the
        // Bencode spec that dictionaries are sorted alphabetically by their keys.
        keys.sort_by(|k1, k2| k1.get(0).unwrap().cmp(k2.get(0).unwrap()));

        keys.into_iter()
            .map(|k| {
                let i = self.lookup.get(k).unwrap();
                (k.clone(), self.values.get(*i).unwrap().clone())
            })
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_bmap_impl() {
        fn put(m: &mut BMap, k: &'static str, v: Benc) {
            m.insert(k.as_bytes().to_vec(), v)
        }
        fn get(m: &mut BMap, k: &'static str) -> Benc {
            m.get(&k.as_bytes().to_vec()).unwrap()
        }

        let mut bmap: BMap = BMap::new();
        assert_eq!(0, bmap.len());

        put(&mut bmap, "a", Benc::Int(1));
        put(&mut bmap, "b", Benc::Int(2));
        put(&mut bmap, "c", Benc::Int(3));
        put(&mut bmap, "d", Benc::Int(4));
        put(&mut bmap, "e", Benc::Int(5));
        put(&mut bmap, "f", Benc::Int(6));
        assert_eq!(6, bmap.len());

        assert_eq!(Benc::Int(1), get(&mut bmap, "a"));
        assert_eq!(Benc::Int(2), get(&mut bmap, "b"));
        assert_eq!(Benc::Int(3), get(&mut bmap, "c"));
        assert_eq!(Benc::Int(4), get(&mut bmap, "d"));
        assert_eq!(Benc::Int(5), get(&mut bmap, "e"));
        assert_eq!(Benc::Int(6), get(&mut bmap, "f"));
    }

    #[test]
    fn test_benc_parse_encode() {
        assert_eq!(
            b"i10e".to_vec(),
            Benc::parse(&b"i10e".to_vec()).unwrap().bytes()
        );
    }
}