Saving

crates/dht/src/routing_table.rs

@@ -4,6 +4,154 @@ use std::{
     time::{Duration, Instant},
 };
 
+enum BucketTreeNode {
+    Leaf(Vec<RoutingTableNode>),
+    LeftRight(Box<BucketTree>, Box<BucketTree>),
+}
+
+pub struct BucketTree {
+    bits: u8,
+    start: Id20,
+    end_inclusive: Id20,
+    data: BucketTreeNode,
+}
+
+pub struct BucketTreeNodeIterator<'a> {
+    current: std::slice::Iter<'a, RoutingTableNode>,
+    queue: Vec<&'a BucketTree>,
+}
+
+impl<'a> BucketTreeNodeIterator<'a> {
+    fn new(mut tree: &'a BucketTree) -> Self {
+        let mut queue = Vec::new();
+        let current = loop {
+            match &tree.data {
+                BucketTreeNode::Leaf(nodes) => break nodes.iter(),
+                BucketTreeNode::LeftRight(left, right) => {
+                    queue.push(right.as_ref());
+                    tree = left.as_ref()
+                }
+            }
+        };
+        BucketTreeNodeIterator { current, queue }
+    }
+}
+
+impl<'a> Iterator for BucketTreeNodeIterator<'a> {
+    type Item = &'a RoutingTableNode;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if let Some(v) = self.current.next() {
+            return Some(v);
+        };
+
+        loop {
+            let tree = self.queue.pop()?;
+            match &tree.data {
+                BucketTreeNode::Leaf(nodes) => {
+                    self.current = nodes.iter();
+                    match self.current.next() {
+                        Some(v) => return Some(v),
+                        None => continue,
+                    }
+                }
+                BucketTreeNode::LeftRight(left, right) => {
+                    self.queue.push(right.as_ref());
+                    self.queue.push(left.as_ref());
+                    continue;
+                }
+            }
+        }
+    }
+}
+
+fn compute_split_start_end(
+    start: Id20,
+    end_inclusive: Id20,
+    bits: u8,
+) -> ((Id20, Id20), (Id20, Id20)) {
+    let changing_bit = 160 - bits;
+    let new_left_end = {
+        let mut c = end_inclusive;
+        c.set_bit(changing_bit, false);
+        c
+    };
+    let new_right_start = {
+        let mut c = start;
+        c.set_bit(changing_bit, true);
+        c
+    };
+    ((start, new_left_end), (new_right_start, end_inclusive))
+}
+
+impl BucketTree {
+    pub fn new() -> Self {
+        BucketTree {
+            bits: 160,
+            start: Id20([0u8; 20]),
+            end_inclusive: Id20([0xff; 20]),
+            data: BucketTreeNode::Leaf(Vec::new()),
+        }
+    }
+    pub fn iter(&self) -> BucketTreeNodeIterator<'_> {
+        BucketTreeNodeIterator::new(self)
+    }
+    pub fn add_node(&mut self, self_id: &Id20, id: Id20, addr: SocketAddr) {
+        let mut tree = self;
+        loop {
+            match &mut tree.data {
+                BucketTreeNode::Leaf(_) => {
+                    assert!(id >= tree.start && id <= tree.end_inclusive);
+                    tree.insert_into_leaf(self_id, id, addr)
+                }
+                BucketTreeNode::LeftRight(left, right) => {
+                    if id >= right.start {
+                        // Erase lifetime.
+                        // Safety: this is safe as it's a tree, not a DAG or Graph.
+                        tree = unsafe { &mut *(right.as_mut() as *mut _) };
+                        continue;
+                    }
+                    tree = unsafe { &mut *(left.as_mut() as *mut _) };
+                }
+            }
+        }
+    }
+    fn insert_into_leaf(&mut self, self_id: &Id20, id: Id20, addr: SocketAddr) {
+        let nodes = match &mut self.data {
+            BucketTreeNode::Leaf(nodes) => nodes,
+            BucketTreeNode::LeftRight(_, _) => unreachable!(),
+        };
+        // if already found, quit
+        if nodes.iter().find(|r| r.id == id).is_some() {
+            return;
+        }
+
+        if nodes.len() < 8 {
+            nodes.push(RoutingTableNode {
+                id,
+                addr,
+                last_request: None,
+                last_response: None,
+                outstanding_queries_in_a_row: 0,
+            });
+            return;
+        }
+
+        // if our id is not inside, don't bother.
+        if *self_id < self.start || *self_id > self.end_inclusive {
+            return;
+        }
+
+        todo!()
+    }
+}
+
+impl Default for BucketTree {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 use crate::id20::Id20;
 
 pub struct RoutingTableNode {
@@ -41,44 +189,59 @@ impl RoutingTableNode {
     }
 }
 
-struct Bucket {
-    bits: u8,
-    nodes: Vec<RoutingTableNode>,
-    end: Id20,
-}
-
 pub struct RoutingTable {
     id: Id20,
     size: usize,
-    buckets: BTreeMap<Id20, Bucket>,
+    buckets: BucketTree,
 }
 
 impl RoutingTable {
     pub fn new(id: Id20) -> Self {
-        let initial_bucket = Id20([0u8; 20]);
-        let mut buckets = BTreeMap::new();
-        buckets.insert(
-            initial_bucket,
-            Bucket {
-                bits: 160,
-                nodes: Vec::new(),
-            },
-        );
         Self {
             id,
-            buckets,
+            buckets: BucketTree::new(),
             size: 0,
         }
     }
     pub fn sorted_by_distance_from(&self, id: Id20) -> Vec<&RoutingTableNode> {
         let mut result = Vec::with_capacity(self.size);
-        for bucket in self.buckets.values() {
-            for node in bucket.nodes.iter() {
-                result.push(node);
-            }
+        for node in self.buckets.iter() {
+            result.push(node);
         }
         result.sort_by_key(|n| id.distance(&n.id));
         result
     }
-    pub fn add_node(&mut self, id: Id20, addr: SocketAddr) -> bool {}
+    pub fn add_node(&mut self, id: Id20, addr: SocketAddr) -> bool {
+        todo!()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{id20::Id20, routing_table::compute_split_start_end};
+
+    #[test]
+    fn compute_split_start_end_root() {
+        let start = Id20([0u8; 20]);
+        let end = Id20([0xffu8; 20]);
+        assert_eq!(
+            compute_split_start_end(start, end, 160),
+            (
+                (
+                    start,
+                    Id20([
+                        0x7f, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+                        0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff
+                    ])
+                ),
+                (
+                    Id20([
+                        0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+                        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00
+                    ]),
+                    end
+                )
+            )
+        )
+    }
 }