Saving

crates/dht/src/id20.rs

@@ -67,6 +67,20 @@ impl Id20 {
         }
         Id20(xor)
     }
+    pub fn set_bit(&mut self, bit: u8, value: bool) {
+        let n = &mut self.0[(bit / 8) as usize];
+        if value {
+            *n |= 1 << (7 - bit % 8)
+        } else {
+            let mask = !(1 << (7 - bit % 8));
+            *n &= mask;
+        }
+    }
+    pub fn set_bits_range(&mut self, r: std::ops::Range<u8>, value: bool) {
+        for bit in r {
+            self.set_bit(bit, value)
+        }
+    }
 }
 
 impl Ord for Id20 {
@@ -87,3 +101,18 @@ impl PartialOrd<Id20> for Id20 {
         Some(self.cmp(other))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::Id20;
+
+    #[test]
+    fn test_set_bit_range() {
+        let mut id = Id20([0u8; 20]);
+        id.set_bits_range(9..17, true);
+        assert_eq!(
+            id,
+            Id20([0, 127, 128, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
+        )
+    }
+}

crates/dht/src/main.rs

@@ -130,12 +130,12 @@ impl DhtState {
                 let nodes = response
                     .nodes
                     .ok_or_else(|| anyhow::anyhow!("expected nodes for find_node requests"))?;
-                let pn = match (response.nodes, response.values) {
+                // let pn = match (response.nodes, response.values) {
-                    (Some(nodes), None) => PeersOrNodes::Nodes(nodes),
+                //     (Some(nodes), None) => PeersOrNodes::Nodes(nodes),
-                    (None, Some(peers)) => PeersOrNodes::Peers(peers),
+                //     (None, Some(peers)) => PeersOrNodes::Peers(peers),
-                    _ => anyhow::bail!("expected nodes or values to be set in find_peers response"),
+                //     _ => anyhow::bail!("expected nodes or values to be set in find_peers response"),
-                };
+                // };
-                self.on_found_peers_or_nodes(id, pn)
+                // self.on_found_peers_or_nodes(id, pn)
             }
         };
         Ok(())
@@ -222,31 +222,31 @@ impl DhtWorker {
             // bootstrap
             for addr in bootstrap_addrs {
                 for addr in tokio::net::lookup_host(addr).await.unwrap() {
-                    let msg = MessageKind::FindNodeRequest(FindNodeRequest {
+                    // let msg = MessageKind::FindNodeRequest(FindNodeRequest {
-                        id: self.peer_id,
+                    //     id: self.peer_id,
-                        target: self.peer_id,
+                    //     target: self.peer_id,
-                    });
+                    // });
-                    in_tx.send((msg, addr)).await.unwrap();
+                    // in_tx.send((msg, addr)).await.unwrap();
                 }
             }
         };
         let mut bootstrap_done = false;
 
-        let request_reader = async {
+        // let request_reader = async {
-            while let Some((request, sender)) = self.request_rx.recv().await {
+        //     while let Some((request, sender)) = self.request_rx.recv().await {
-                self.on_request(request, sender)
+        //         self.on_request(request, sender)
-            }
+        //     }
-        };
+        // };
 
-        tokio::select! {
+        // tokio::select! {
-            _ = framer => {
+        //     _ = framer => {
-                anyhow::bail!("framer quit")
+        //         anyhow::bail!("framer quit")
-            },
+        //     },
-            _ = bootstrap, if !bootstrap_done => {
+        //     _ = bootstrap, if !bootstrap_done => {
-                bootstrap_done = true
+        //         bootstrap_done = true
-            },
+        //     },
-            _ = request_reader => {}
+        //     _ = request_reader => {}
-        }
+        // }
 
         todo!()
     }

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 self.buckets.iter() {
-            for node in bucket.nodes.iter() {
+            result.push(node);
-                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
+                )
+            )
+        )
+    }
 }