babbel_bencode 0.1.1

Fast, binary-safe BitTorrent Bencode parser, serializer, and DOM in pure Rust
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

//! Visitor pattern abstraction for serializing bencode data into target formats.
//! Enables Open/Closed Principle (OCP) compliance: new target formats implement `BencodeVisitor`
//! without altering core AST data structures or parser logic.

#[cfg(not(feature = "std"))]
use alloc::{string::String, vec::Vec};

use crate::nodes::node::Node;

/// Trait implemented by format serializers to receive traversal events.
pub trait BencodeVisitor {
    type Error;

    /// Visit an integer value.
    fn visit_integer(&mut self, value: i64) -> Result<(), Self::Error>;
    /// Visit a string value.
    fn visit_string(&mut self, value: &str) -> Result<(), Self::Error>;
    /// Visit the start of a list.
    fn visit_list_start(&mut self) -> Result<(), Self::Error>;
    /// Visit the end of a list.
    fn visit_list_end(&mut self) -> Result<(), Self::Error>;
    /// Visit the start of a dictionary.
    fn visit_dict_start(&mut self) -> Result<(), Self::Error>;
    /// Visit a dictionary key.
    fn visit_dict_key(&mut self, key: &str) -> Result<(), Self::Error>;
    /// Visit the end of a dictionary.
    fn visit_dict_end(&mut self) -> Result<(), Self::Error>;
    /// Visit an uninitialized/None value.
    fn visit_none(&mut self) -> Result<(), Self::Error>;
}

/// Trait implemented by data structures that can be traversed by a `BencodeVisitor`.
pub trait BencodeVisitable {
    /// Accepts a visitor and traverses the internal structure.
    fn accept<V: BencodeVisitor>(&self, visitor: &mut V) -> Result<(), V::Error>;
}

impl BencodeVisitable for Node {
    fn accept<V: BencodeVisitor>(&self, visitor: &mut V) -> Result<(), V::Error> {
        match self {
            Node::Integer(val) => visitor.visit_integer(*val),
            Node::Str(val) => visitor.visit_string(val),
            Node::List(items) => {
                visitor.visit_list_start()?;
                for item in items {
                    item.accept(visitor)?;
                }
                visitor.visit_list_end()
            }
            Node::Dictionary(items) => {
                visitor.visit_dict_start()?;
                if items.len() <= 16 {
                    let mut stack: [Option<(&String, &Node)>; 16] = [None; 16];
                    let mut count = 0;
                    for pair in items {
                        stack[count] = Some((pair.0, pair.1));
                        count += 1;
                    }
                    let slice = &mut stack[..count];
                    slice.sort_unstable_by(|a, b| a.unwrap().0.cmp(b.unwrap().0));
                    for item in slice.iter() {
                        let (key, value) = item.unwrap();
                        visitor.visit_dict_key(key)?;
                        value.accept(visitor)?;
                    }
                } else {
                    let mut sorted: Vec<_> = items.iter().collect();
                    sorted.sort_unstable_by(|a, b| a.0.cmp(b.0));
                    for (key, value) in sorted {
                        visitor.visit_dict_key(key)?;
                        value.accept(visitor)?;
                    }
                }
                visitor.visit_dict_end()
            }
            Node::None => visitor.visit_none(),
        }
    }
}