ez_bencoding/decode/
dict.rs

1use std::{borrow::Cow, collections::HashMap};
2
3use super::{BdecodeNode, token::BdecodeTokenType, IBdecodeNode};
4
5crate::collective_bdecode_node!(Dict);
6
7impl Dict {
8    /// 获取 dict 中指定索引的节点对(key, value)
9    pub fn item(&self, index: usize) -> (BdecodeNode, BdecodeNode) {
10        assert!(self.token_type() == BdecodeTokenType::Dict);
11
12        if index >= self.len() {
13            panic!("index out of range");
14        }
15
16        // get key node
17        let key_token_idx = self.item_indexes[index];
18        if key_token_idx as usize >= self.tokens.len() {
19            panic!("index out of range in tokens");
20        }
21        let key_node = BdecodeNode::new(key_token_idx, self.tokens(), self.buffer.clone());
22        let key_token = &self.tokens[key_token_idx as usize];
23        
24        // get value node
25        let val_token_idx = key_token_idx + key_token.next_item();
26        let val_node = BdecodeNode::new(val_token_idx, self.tokens(), self.buffer.clone());
27
28        (key_node, val_node)
29    }
30
31    /// 在 dict 中查找 key 对应的 value
32    pub fn find(&self, key: &[u8]) -> Option<BdecodeNode> {
33        assert!(self.token_type() == BdecodeTokenType::Dict);
34
35        for token_index in self.item_indexes.as_ref() {
36            let token = &self.tokens[*token_index as usize];
37            assert!(token.node_type() == BdecodeTokenType::Str);
38            let next_offset = self.tokens[(token_index + 1) as usize].offset() as usize;
39            let start = (token.offset() + token.header_size() as u32 + 1) as usize;
40
41            if &self.buffer[start..next_offset] == key {
42                let val_token_idx = *token_index + token.next_item();
43
44                return Some(BdecodeNode::new(val_token_idx, self.tokens(), self.buffer.clone()));
45            }
46        }
47
48        None
49    }
50
51    pub fn find_as_str(&self, key: &[u8]) -> Option<Cow<[u8]>> {
52        let node = self.find(key);
53
54        if let Some(node) = node {
55            let val = node.as_str();
56            let val_ptr = val.as_ref() as *const [u8];
57            let val_ref = unsafe { &*val_ptr };
58
59            let rst = Cow::Borrowed(val_ref);
60
61            return Some(rst);
62        }
63
64        None
65    }
66
67    pub fn find_as_int(&self, key: &[u8]) -> Option<i64> {
68        let node = self.find(key);
69
70        if let Some(node) = node {
71            return node.as_int().ok();
72        }
73
74        None
75    }
76
77    pub fn find_as_list(&self, key: &[u8]) -> Option<Vec<BdecodeNode>> {
78        let node = self.find(key);
79
80        if let Some(node) = node {
81            return if let BdecodeNode::List(node) = node {
82                let mut nodes = vec![];
83                for i in 0..node.len() {
84                    let node = node.item(i);
85                    nodes.push(node);
86                }
87
88                Some(nodes)
89            } else {
90                None
91            };
92        }
93
94        None
95    }
96
97    pub fn find_as_dict(&self, key: &[u8]) -> Option<HashMap<Cow<[u8]>, BdecodeNode>> {
98        let Some(node) = self.find(key) else {
99            return None;
100        };
101
102        let mut node_map = HashMap::new();
103        let BdecodeNode::Dict(node) = node else { return None };
104
105        for i in 0..node.len() {
106            let (key, value) = node.item(i);
107
108            let key_str = key.as_str();
109            let key_ptr = key_str.as_ref() as *const [u8];
110            let key_ref = unsafe { &*key_ptr };
111
112            let key = Cow::Borrowed(key_ref);
113
114            node_map.insert(key, value);
115        }
116
117        Some(node_map)
118    }
119
120    pub fn to_json(&self) -> String {
121        let mut sb = String::new();
122        let len = self.len();
123
124        for i in 0..len {
125            let (key, val) = self.item(i);
126            sb.push_str(&format!("{}: {}", key.to_json(), val.to_json()));
127
128            if i < len - 1 { 
129                sb.push_str(", "); 
130            }
131        }
132        
133        format!("{} {} {}", "{", sb, "}")
134    }
135}