From 86098797034cbc7eb6db0cee54e17f8dcaedbc5d Mon Sep 17 00:00:00 2001
From: evuez <julien@mulga.net>
Date: Sat, 26 Nov 2022 15:38:06 -0500
Subject: Initial commit

---
 src/common/trie.rs | 151 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 151 insertions(+)
 create mode 100644 src/common/trie.rs

(limited to 'src/common/trie.rs')

diff --git a/src/common/trie.rs b/src/common/trie.rs
new file mode 100644
index 0000000..381d2b5
--- /dev/null
+++ b/src/common/trie.rs
@@ -0,0 +1,151 @@
+#![allow(dead_code)]
+use crate::common::ip::Block;
+use crate::common::ip::Meta;
+use crate::common::ip::BITMASK;
+
+pub struct Trie {
+    root: Node,
+}
+
+#[derive(Debug)]
+struct Node {
+    block: Option<Block>,
+    left: Option<Box<Node>>,
+    right: Option<Box<Node>>,
+}
+
+impl Node {
+    pub const NULL: Node = Self {
+        block: None,
+        left: None,
+        right: None,
+    };
+
+    fn new(block: Option<Block>) -> Self {
+        Node {
+            block,
+            left: None,
+            right: None,
+        }
+    }
+
+    fn is_leaf(&self) -> bool {
+        self.left.is_none() && self.right.is_none()
+    }
+}
+
+impl Trie {
+    pub fn new() -> Self {
+        Trie { root: Node::NULL }
+    }
+
+    pub fn insert(&mut self, block: Block, _meta: Meta) {
+        let mut node = &mut self.root;
+        let mut addr = block.addr;
+
+        for _ in 0..block.mlen {
+            if addr & BITMASK == 0 {
+                if node.left.is_none() {
+                    node.left = Some(Box::new(Node::NULL));
+                }
+                node = node.left.as_deref_mut().unwrap();
+            } else {
+                if node.right.is_none() {
+                    node.right = Some(Box::new(Node::NULL));
+                }
+                node = node.right.as_deref_mut().unwrap();
+            }
+            addr <<= 1;
+        }
+
+        node.block = Some(block);
+    }
+
+    pub fn find(&self, needle: &Block) -> Option<Block> {
+        self.find_block_node(needle)?.block
+    }
+
+    pub fn parent(&self, needle: &Block) -> Option<Block> {
+        let mut node = &self.root;
+        let mut addr = needle.addr;
+        let mut parent = None;
+
+        for _ in 0..needle.mlen {
+            if node.block.is_some() {
+                parent = node.block;
+            }
+
+            if addr & BITMASK == 0 {
+                if node.left.is_some() {
+                    node = node.left.as_deref().unwrap();
+                } else {
+                    break;
+                }
+            } else if node.right.is_some() {
+                node = node.right.as_deref().unwrap();
+            } else {
+                break;
+            }
+            addr <<= 1;
+        }
+
+        parent
+    }
+
+    pub fn children(&self, parent: &Block) -> Vec<Block> {
+        let root: &Node = match self.find_block_node(parent) {
+            Some(root) => root,
+            None => return vec![],
+        };
+
+        let mut children: Vec<Block> = Vec::new();
+        let mut stack: Vec<&Node> = vec![root];
+
+        while let Some(node) = stack.pop() {
+            if node.is_leaf() && node.block.is_some() && node.block.unwrap() != root.block.unwrap()
+            {
+                children.push(node.block.unwrap());
+                continue;
+            }
+
+            if node.left.is_some() {
+                stack.push(node.left.as_deref().unwrap());
+            }
+            if node.right.is_some() {
+                stack.push(node.right.as_deref().unwrap());
+            }
+        }
+
+        children
+    }
+
+    fn find_block_node(&self, needle: &Block) -> Option<&Node> {
+        let mut node = &self.root;
+        let mut addr = needle.addr;
+        let mut match_ = None;
+
+        for _ in 0..(needle.mlen + 1) {
+            if node.block.is_some() {
+                match_ = Some(node);
+            }
+
+            if addr & BITMASK == 0 {
+                if node.left.is_some() {
+                    node = node.left.as_deref().unwrap();
+                } else {
+                    break;
+                }
+            } else if node.right.is_some() {
+                node = node.right.as_deref().unwrap();
+            } else {
+                break;
+            }
+            addr <<= 1;
+        }
+
+        match match_?.block {
+            Some(block) if block == *needle => match_,
+            _ => None,
+        }
+    }
+}
-- 
cgit v1.2.3