Skip to main content

bdecode/
token.rs

1use std::fmt;
2
3use super::BdecodeError;
4
5const OFFSET_MASK: u64 = 0xFFFF_FFF8_0000_0000;
6const NEXT_ITEM_MASK: u64 = 0x0000_0007_FFFF_FFC0;
7const HEADER_MASK: u64 = 0x0000_0000_0000_0038;
8const TYPE_MASK: u64 = 0x0000_0000_0000_0007;
9
10const OFFSET_OFFSET: u64 = 35;
11const NEXT_ITEM_OFFSET: u64 = 6;
12const HEADER_OFFSET: u64 = 3;
13const TYPE_OFFSET: u64 = 0;
14
15#[derive(Debug, Clone, Copy, Eq, PartialEq)]
16pub enum TokenType {
17    Dict = 1,
18    List = 2,
19    Str = 3,
20    Int = 4,
21    /// the node with type 'end' is a logical node, pointing to the end of
22    /// the bencoded buffer.
23    End = 5,
24}
25
26#[derive(Clone, Copy, Eq, PartialEq)]
27pub struct Token {
28    inner: u64,
29}
30
31impl Token {
32    pub const MAX_OFFSET: usize = (1 << 29) - 1;
33    pub const MAX_NEXT_ITEM: usize = (1 << 29) - 1;
34    pub const MAX_HEADER: usize = (1 << 3) - 1;
35
36    pub fn new(
37        offset: usize,
38        token_type: TokenType,
39        next_item: usize,
40        header: usize,
41    ) -> Result<Token, BdecodeError> {
42        if (offset > Self::MAX_OFFSET)
43            || (next_item > Self::MAX_NEXT_ITEM)
44            || (header > Self::MAX_HEADER)
45        {
46            return Err(BdecodeError::LimitExceeded);
47        }
48
49        let inner = ((offset as u64) << OFFSET_OFFSET)
50            | ((next_item as u64) << NEXT_ITEM_OFFSET)
51            | ((header as u64) << HEADER_OFFSET)
52            | ((token_type as u64) << TYPE_OFFSET);
53
54        Ok(Token { inner })
55    }
56
57    #[inline]
58    pub fn offset(&self) -> usize {
59        ((self.inner & OFFSET_MASK) >> OFFSET_OFFSET) as usize
60    }
61
62    /// If this node is a member of a list, `next_item` is the number of nodes
63    /// to jump forward in th node array to get to the next item in the list.
64    /// if it's a key in a dictionary, it's the number of step forwards to get
65    /// to its corresponding value. If it's a value in a dictionary, it's the
66    /// number of steps to the next key, or to the end node.
67    /// this is the _relative_ offset to the next node
68    #[inline]
69    pub fn next_item(&self) -> usize {
70        ((self.inner & NEXT_ITEM_MASK) >> NEXT_ITEM_OFFSET) as usize
71    }
72
73    #[inline]
74    pub fn set_next_item(&mut self, new_next_item: usize) -> Result<(), BdecodeError> {
75        if new_next_item > Self::MAX_NEXT_ITEM {
76            return Err(BdecodeError::LimitExceeded);
77        }
78        let inner_zeroed_ni = self.inner & (!NEXT_ITEM_MASK);
79        self.inner = inner_zeroed_ni | ((new_next_item as u64) << NEXT_ITEM_OFFSET);
80        Ok(())
81    }
82
83    /// this is the number of bytes to skip forward from the offset to get to the
84    /// first character of the string, if this is a string. This field is not
85    /// used for other types. Essentially this is the length of the length prefix
86    /// and the colon. Since a string always has at least one character of length
87    /// prefix and always a colon, those 2 characters are implied.
88    #[inline]
89    pub fn header(&self) -> usize {
90        ((self.inner & HEADER_MASK) >> HEADER_OFFSET) as usize
91    }
92
93    #[inline]
94    pub fn token_type(&self) -> TokenType {
95        let type_int = ((self.inner & TYPE_MASK) >> TYPE_OFFSET) as usize;
96        match type_int {
97            1 => TokenType::Dict,
98            2 => TokenType::List,
99            3 => TokenType::Str,
100            4 => TokenType::Int,
101            5 => TokenType::End,
102            _ => unreachable!(),
103        }
104    }
105
106    #[inline]
107    pub fn start_offset(&self) -> usize {
108        // Shouldn't this just be an if statement based on type? One where we
109        // conditionally plus 2?
110        assert_eq!(self.token_type(), TokenType::Str);
111        self.header() + 2
112    }
113}
114
115impl fmt::Debug for Token {
116    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
117        f.debug_struct("Token")
118            .field("offset", &self.offset())
119            .field("next_item", &self.next_item())
120            .field("header", &self.header())
121            .field("token_type", &self.token_type())
122            .finish()
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129    use std::mem;
130
131    #[test]
132    fn test_token_fields() {
133        let mut tok = Token::new(42, TokenType::Dict, 11, 7).unwrap();
134        assert_eq!(tok.offset(), 42);
135        assert_eq!(tok.token_type(), TokenType::Dict);
136        assert_eq!(tok.next_item(), 11);
137        assert_eq!(tok.header(), 7);
138
139        tok.set_next_item(29312).unwrap();
140        // After setting next item, the rest of the fields should stay the
141        // same.
142        assert_eq!(tok.offset(), 42);
143        assert_eq!(tok.token_type(), TokenType::Dict);
144        assert_eq!(tok.next_item(), 29312);
145        assert_eq!(tok.header(), 7);
146    }
147
148    #[test]
149    fn test_token_size() {
150        assert_eq!(mem::size_of::<Token>(), 8);
151    }
152}