add visualizer.md for live btree state

Signed-off-by: Arshdeep54 <balarsh535@gmail.com>

Author: Arshdeep54

src/btree/README.md

@@ -1,5 +1,14 @@
 # Btree Implementation
 
+# Usage 
+
+- Run `cargo run` 
+- Start filling up the btree with `BTREE insert 4 heykey4` and so on.. 
+- Search using `BTREE search 4` 
+- Delete using `BTREE delete 4`
+
+### See the live visualization of the Btree in `/tests/visualizer.md` (Use CTRL+SHIFT+V for rendering markdown)
+
 # Resouces
 
 - https://www.dataquest.io/blog/b-tree-data-structure/

src/btree/mod.rs

@@ -6,7 +6,9 @@ use std::io::{self, Result};
 mod metadata;
 mod node;
 mod paging;
-pub const DEGREE: i32 = 3;
+pub mod utils;
+
+pub const DEGREE: i32 = 2;
 pub const MIN_ITEMS: i32 = DEGREE - 1;
 pub const MAX_ITEMS: i32 = DEGREE * 2;
 

src/btree/node.rs

@@ -1,7 +1,7 @@
 use std::{fmt, io::Result};
 
 use super::{
-    Item, MAX_ITEMS, MIN_ITEMS,
+    Item, MAX_ITEMS,
     paging::{Page, PageID, Pager},
 };
 
@@ -75,6 +75,8 @@ impl Node {
         if pos > self.num_items || pos < 0 {
             return;
         }
+
+        // Find the correct position to maintain sorted order
         let mut insert_pos = pos as usize;
         while insert_pos < self.items.len() && self.items[insert_pos].key < item.key {
             insert_pos += 1;
@@ -96,11 +98,11 @@ impl Node {
         let new_id = pager.allocate_page()?;
         let mut new_node = Node::new(new_id);
 
-        let mid = MIN_ITEMS;
+        let mid = self.num_items / 2;
         let mid_item = self.items[mid as usize].clone();
 
-        new_node.items = self.items[mid as usize..].to_vec();
-        new_node.num_items = self.num_items - mid;
+        new_node.items = self.items[mid as usize + 1..].to_vec();
+        new_node.num_items = self.num_items - mid - 1;
 
         if !self.is_leaf() {
             new_node.children = self.children[mid as usize + 1..].to_vec();
@@ -130,20 +132,31 @@ impl Node {
             return;
         }
 
-        if self.children[pos as usize].num_items >= MAX_ITEMS {
-            let (mid_item, new_node) = self.children[pos as usize].split(pager).unwrap();
-            self.insert_item_at(pos, mid_item);
-            self.insert_child_at(pos + 1, new_node);
+        let child_pos = if pos >= self.num_items {
+            (self.num_items) as usize
+        } else {
+            pos as usize
+        };
+
+        if self.children[child_pos].num_items >= MAX_ITEMS {
+            let (mid_item, new_node) = self.children[child_pos].split(pager).unwrap();
+            let mid_key = mid_item.key;
+            self.insert_item_at(child_pos as i32, mid_item);
+            self.insert_child_at(child_pos as i32 + 1, new_node);
 
-            let condition = item.key - self.items[pos as usize].key;
-            if condition < 0 {
-            } else if condition > 0 {
-                pos += 1;
+            if item.key > mid_key {
+                pos = child_pos as i32 + 1;
             } else {
-                return;
+                pos = child_pos as i32;
             }
         }
-        self.children[pos as usize].insert(item, pager);
+
+        let child_pos = if pos >= self.num_items {
+            (self.num_items) as usize
+        } else {
+            pos as usize
+        };
+        self.children[child_pos].insert(item, pager);
     }
 }
 

src/btree/tests/visualizer.md

@@ -0,0 +1,26 @@
+# B-tree Visualization
+
+Current state of the B-tree:
+
+```mermaid
+graph TD
+    n8["Node 8<br>51:he555"]
+    n3["Node 3<br>4:he555 | 6:h555"]
+    n1["Node 1<br>1:he555 | 2:he555 | 3:he555"]
+    n2["Node 2<br>5:he555"]
+    n9["Node 9<br>7:h555 | 8:he555 | 9:h555"]
+    n7["Node 7<br>5224:he555 | 5855:he555 | 58555:he555"]
+    n4["Node 4<br>58:he555 | 528:he555"]
+    n5["Node 5<br>5555:he555"]
+    n6["Node 6<br>9888:h555 | 55555:he555"]
+    n10["Node 10<br>58558:he555 | 58588:he555"]
+    n8 --> n3
+    n3 --> n1
+    n3 --> n2
+    n3 --> n9
+    n8 --> n7
+    n7 --> n4
+    n7 --> n5
+    n7 --> n6
+    n7 --> n10
+```

src/btree/utils.rs

@@ -0,0 +1,85 @@
+use std::fs::{self, write};
+use std::io;
+use std::path::Path;
+
+use super::{Btree, Node};
+
+pub struct Visualizer {
+    output_path: String,
+}
+
+impl Visualizer {
+    pub fn new(path: &str) -> Self {
+        if let Some(parent) = Path::new(path).parent() {
+            if !parent.exists() {
+                fs::create_dir_all(parent).expect("Failed to create visualization directory");
+            }
+        }
+
+        Visualizer {
+            output_path: path.to_string(),
+        }
+    }
+
+    pub fn update(&self, btree: &Btree) -> io::Result<()> {
+        let mermaid = self.generate_mermaid(btree);
+        let content = format!(
+            "# B-tree Visualization\n\n\
+            Current state of the B-tree:\n\n\
+            ```mermaid\n\
+            graph TD\n\
+            {mermaid}\n\
+            ```\n"
+        );
+        write(&self.output_path, content)
+    }
+
+    fn generate_mermaid(&self, btree: &Btree) -> String {
+        match &btree.root {
+            None => String::from("empty[Empty Tree]"),
+            Some(root) => {
+                let mut nodes = Vec::new();
+                let mut edges = Vec::new();
+                self.generate_node_diagram(&mut nodes, &mut edges, root, None);
+                format!("{}\n{}", nodes.join("\n"), edges.join("\n"))
+            }
+        }
+    }
+
+    fn generate_node_diagram(
+        &self,
+        nodes: &mut Vec<String>,
+        edges: &mut Vec<String>,
+        node: &Node,
+        parent_id: Option<u32>,
+    ) {
+        let items: Vec<String> = node
+            .items
+            .iter()
+            .map(|item| format!("{}:{}", item.key, self.truncate_value(&item.val, 5)))
+            .collect();
+
+        nodes.push(format!(
+            "    n{}[\"Node {}<br>{}\"]",
+            node.id,
+            node.id,
+            items.join(" | ")
+        ));
+
+        if let Some(parent) = parent_id {
+            edges.push(format!("    n{} --> n{}", parent, node.id));
+        }
+
+        for child in &node.children {
+            self.generate_node_diagram(nodes, edges, child, Some(node.id));
+        }
+    }
+
+    fn truncate_value(&self, value: &str, max_len: usize) -> String {
+        if value.len() <= max_len {
+            value.to_string()
+        } else {
+            format!("{}...", &value[..max_len])
+        }
+    }
+}

src/parsing/mod.rs

@@ -1,4 +1,8 @@
-use crate::{IndexSession, btree::Item};
+use crate::{
+    IndexSession,
+    btree::{Item, utils::Visualizer},
+};
+use std::path::PathBuf;
 
 struct Command<KeyType, ValType> {
     index_type: String,
@@ -12,30 +16,57 @@ where
     KeyType: std::str::FromStr,
     ValType: std::str::FromStr,
 {
-    fn new(command: &str) -> Self {
-        let tokens: Vec<&str> = command.split(" ").collect();
+    fn new(command: &str) -> Option<Self> {
+        let tokens: Vec<&str> = command.split_whitespace().collect();
+        if tokens.is_empty() {
+            return None;
+        }
+
+        let index_type = tokens[0].to_string();
+        if tokens.len() < 2 {
+            return None;
+        }
+
+        let index_function = tokens[1].to_string();
         let key = if tokens.len() > 2 {
-            Some(tokens[2].parse::<KeyType>().ok().unwrap())
+            tokens[2].parse::<KeyType>().ok()
         } else {
             None
         };
         let value = if tokens.len() > 3 {
-            Some(tokens[3].parse::<ValType>().ok().unwrap())
+            tokens[3].parse::<ValType>().ok()
         } else {
             None
         };
-        Command {
-            index_type: tokens[0].to_string(),
-            index_function: tokens[1].to_string(),
+
+        Some(Command {
+            index_type,
+            index_function,
             key,
             value,
-        }
+        })
     }
 }
+
 pub fn parse_command(index_session: &mut IndexSession, command: &str) {
     let trimmed_command = command.trim();
+    if trimmed_command.is_empty() {
+        return;
+    }
     println!("Command: {trimmed_command}");
-    let cmd: Command<i32, String> = Command::new(command);
+
+    let mut viz_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
+    viz_path.push("src");
+    viz_path.push("btree");
+    viz_path.push("tests");
+    viz_path.push("visualizer.md");
+    let visualizer = Visualizer::new(viz_path.to_str().unwrap());
+
+    let cmd = match Command::<i32, String>::new(trimmed_command) {
+        Some(cmd) => cmd,
+        None => return,
+    };
+
     if cmd.index_type.as_str() == "BTREE" || cmd.index_type.as_str() == "btree" {
         match cmd.index_function.as_str() {
             "INSERT" | "insert" => {
@@ -58,7 +89,10 @@ pub fn parse_command(index_session: &mut IndexSession, command: &str) {
                 };
 
                 index_session.btree.insert(Item { key, val });
-                println!("{:?}", index_session.btree);
+
+                if let Err(e) = visualizer.update(&index_session.btree) {
+                    eprintln!("Failed to update visualization: {e}");
+                }
             }
             "SEARCH" | "search" => {
                 let key = match cmd.key {
@@ -85,7 +119,10 @@ pub fn parse_command(index_session: &mut IndexSession, command: &str) {
                 match index_session.btree.delete(key) {
                     Ok(_) => {
                         println!("Successfully deleted key {key}");
-                        println!("Tree after deletion: {:?}", index_session.btree);
+
+                        if let Err(e) = visualizer.update(&index_session.btree) {
+                            eprintln!("Failed to update visualization: {e}");
+                        }
                     }
                     Err(e) => println!("Failed to delete key {key}: {e}"),
                 }