kbucket/
tree.rs

1use std::net;
2
3type NodeLink = Option<Box<Node>>;
4
5/// Main bucket tree data structure
6#[derive(Debug)]
7pub struct BucketTree {
8    pub node_id: [u8; 64],
9    pub root: Node,
10}
11
12/// Nodes in the tree can either be branches or leafs
13#[derive(Debug)]
14pub enum Node {
15    BranchNode(BranchNode),
16    LeafNode(LeafNode),
17}
18
19/// Branch nodes store references to either other branches or bucket nodes
20#[derive(Debug)]
21pub struct BranchNode {
22    pub left: NodeLink,
23    pub right: NodeLink,
24}
25
26/// Leaf nodes store the actual buckets
27#[derive(Debug)]
28pub struct LeafNode {
29    pub can_split: bool,
30    pub contacts: Vec<NodeInfo>,
31}
32
33/// Used to represent either IPv4 addresses or IPv6 addresses
34#[derive(Debug)]
35pub enum Address {
36    V4(net::Ipv4Addr),
37    V6(net::Ipv6Addr),
38}
39
40/// Used to store a noes ID and it's associated data
41#[derive(Debug)]
42pub struct NodeInfo {
43    id: [u8; 64],
44    ip: Address,
45    port: u64,
46}
47
48impl BucketTree {
49    /// Create a new bucket tree
50    pub fn new(node_id: [u8; 64]) -> Self {
51        BucketTree {
52            node_id,
53            root: Node::LeafNode(LeafNode {
54                can_split: true,
55                contacts: Vec::new(),
56            }),
57        }
58    }
59
60    /// Function to add node to the tree
61    pub fn add(&mut self, node: &NodeInfo) {
62        let mut current_node = &mut self.root;
63        let mut bit_index = 0;
64
65        // Walk the tree using the provided node id
66        while let Node::BranchNode(x) = current_node {
67            if bit_value(&node.id, bit_index) {
68                current_node = x.right.as_mut().unwrap();
69            } else {
70                current_node = x.left.as_mut().unwrap();
71            }
72            bit_index += 1;
73            continue;
74        }
75
76        println!("{:?}", current_node);
77    }
78}
79
80impl LeafNode {
81    /// Create a new leaf node
82    #[allow(clippy::all)]
83    pub fn new() -> Option<Box<Node>> {
84        Some(Box::new(Node::LeafNode(LeafNode {
85            can_split: true,
86            contacts: Vec::new(),
87        })))
88    }
89}
90
91impl BranchNode {
92    /// Create a new branch node
93    #[allow(clippy::all)]
94    pub fn new() -> Option<Box<Node>> {
95        Some(Box::new(Node::BranchNode(BranchNode {
96            left: None,
97            right: None,
98        })))
99    }
100}
101
102impl NodeInfo {
103    pub fn new(id: [u8; 64], ip: Address, port: u64) -> Self {
104        NodeInfo { id, ip, port }
105    }
106}
107
108// Find out if the value at a specific index is a 0 or 1
109fn bit_value(key: &[u8; 64], bit_index: usize) -> bool {
110    let byte_index = bit_index >> 3;
111    let inner_bit_index = bit_index % 8;
112
113    key[byte_index] & (1 << (7 - inner_bit_index)) >= 1
114}
kbucket/tree.rs

kbucket/
tree.rs
