make internal nodes also contains values

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

Author: Arshdeep54

src/btree/README.md

@@ -6,3 +6,4 @@
 - https://www.cloudcentric.dev/implementing-a-b-tree-in-go/
 - https://youtu.be/7MqaHGWRS3E?si=6INHVHPWdd6J6pgu
 - https://nimrodshn.medium.com/writing-a-storage-engine-in-rust-writing-a-persistent-btree-part-1-916b6f3e2934
+- https://planetscale.com/blog/btrees-and-database-indexes

src/btree/mod.rs

@@ -91,15 +91,11 @@ impl Btree {
             .take()
             .expect("Called split_root on an empty tree.");
 
+        let (mid_item, new_node) = old_root.split(&mut self.pager).unwrap();
         let new_root_id = self.pager.allocate_page().unwrap();
         let mut new_root = Node::new(new_root_id);
 
-        let (mid_item, new_node) = old_root.split(&mut self.pager).unwrap();
-        let nav_item = Item {
-            key: mid_item.key,
-            val: String::new(), // Internal nodes don't store values
-        };
-        new_root.insert_item_at(0, nav_item);
+        new_root.insert_item_at(0, mid_item);
         new_root.insert_child_at(0, *old_root);
         new_root.insert_child_at(1, new_node);
         self.root = Some(Box::new(new_root));
@@ -117,16 +113,8 @@ impl Btree {
             let (pos, found) = current_node.search(key);
 
             if found {
-                if current_node.is_leaf() {
-                    let val = &current_node.items[pos as usize].val;
-                    return Ok(val.clone());
-                } else {
-                    current_node_opt = current_node
-                        .children
-                        .get(pos as usize + 1)
-                        .map(|boxed_node| &**boxed_node);
-                    continue;
-                }
+                let val = &current_node.items[pos as usize].val;
+                return Ok(val.clone());
             }
 
             current_node_opt = current_node

src/btree/node.rs

@@ -80,16 +80,7 @@ impl Node {
             insert_pos += 1;
         }
 
-        let item_to_insert = if !self.is_leaf() {
-            Item {
-                key: item.key,
-                val: String::new(),
-            }
-        } else {
-            item
-        };
-
-        self.items.insert(insert_pos, item_to_insert);
+        self.items.insert(insert_pos, item);
         self.num_items += 1;
     }
 
@@ -141,17 +132,7 @@ impl Node {
 
         if self.children[pos as usize].num_items >= MAX_ITEMS {
             let (mid_item, new_node) = self.children[pos as usize].split(pager).unwrap();
-
-            if !self.is_leaf() {
-                let nav_item = Item {
-                    key: mid_item.key,
-                    val: String::new(),
-                };
-                self.insert_item_at(pos, nav_item);
-            } else {
-                self.insert_item_at(pos, mid_item);
-            }
-
+            self.insert_item_at(pos, mid_item);
             self.insert_child_at(pos + 1, new_node);
 
             let condition = item.key - self.items[pos as usize].key;
@@ -196,9 +177,8 @@ impl Node {
 
         let mut all_items = Vec::new();
         all_items.extend(original_items);
-        all_items.extend(next_child.items);
         all_items.push(separator);
-
+        all_items.extend(next_child.items);
         all_items.sort_by_key(|item| item.key);
         child.items = all_items;
 
@@ -214,27 +194,20 @@ impl Node {
         let (left, right) = self.children.split_at_mut(pos as usize);
         let sibling = &mut left[left.len() - 1];
         let child = &mut right[0];
-
         let sibling_item = sibling.items.pop().unwrap();
 
         let parent_item = self.items[pos as usize - 1].clone();
 
-        self.items[pos as usize - 1] = Item {
-            key: sibling_item.key,
-            val: String::new(),
-        };
+        self.items[pos as usize - 1] = sibling_item;
 
-        if child.is_leaf() {
-            child.items.insert(0, sibling_item);
-        } else {
-            child.items.insert(0, parent_item);
-        }
+        child.items.insert(0, parent_item);
 
         if !sibling.is_leaf() {
             if let Some(last_child) = sibling.children.pop() {
                 child.children.insert(0, last_child);
             }
         }
+
         sibling.num_items -= 1;
         child.num_items += 1;
 
@@ -253,16 +226,9 @@ impl Node {
 
         let parent_item = self.items[pos as usize].clone();
 
-        self.items[pos as usize] = Item {
-            key: sibling_item.key,
-            val: String::new(),
-        };
+        self.items[pos as usize] = sibling_item;
 
-        if child.is_leaf() {
-            child.items.push(sibling_item);
-        } else {
-            child.items.push(parent_item);
-        }
+        child.items.push(parent_item);
 
         if !sibling.is_leaf() {
             if let Some(first_child) = sibling.children.first().cloned() {
@@ -289,12 +255,10 @@ impl Node {
                 items: self.items.clone(),
             }
         } else {
-            let keys = self.items.iter().map(|it| it.key).collect();
-            let children = self.children.iter().map(|c| c.id).collect();
             Page::Internal {
                 id: self.id,
-                keys,
-                children,
+                items: self.items.clone(),
+                children: self.children.iter().map(|c| c.id).collect(),
             }
         }
     }
@@ -313,22 +277,15 @@ impl Node {
             }
             Page::Internal {
                 id,
-                keys,
-                children: child_ids,
+                items,
+                children,
             } => {
-                let items: Vec<Item> = keys
-                    .iter()
-                    .map(|k| Item {
-                        key: *k,
-                        val: String::new(),
-                    })
-                    .collect();
                 let num_items = items.len() as i32;
-                let num_children = child_ids.len() as i32;
+                let num_children = children.len() as i32;
 
                 Node {
                     id: *id,
-                    items,
+                    items: items.clone(),
                     children: Vec::new(),
                     num_items,
                     num_children,

src/btree/paging.rs

@@ -10,7 +10,7 @@ pub type PageID = u32;
 pub enum Page {
     Internal {
         id: PageID,
-        keys: Vec<i32>,
+        items: Vec<Item>,
         children: Vec<PageID>,
     },
     Leaf {
@@ -81,17 +81,28 @@ impl Pager {
         let mut buf = vec![0u8; self.page_size];
         let page_id: u32;
         match page {
-            Page::Internal { id, keys, children } => {
+            Page::Internal {
+                id,
+                items,
+                children,
+            } => {
                 page_id = *id;
 
                 buf[0] = 1; // 1-> internal type
-                let key_count = keys.len() as u32;
-                buf[1..5].copy_from_slice(&key_count.to_le_bytes()); // next 4 bytes-> keycount in little endian format
+                let items_count = items.len() as u32;
+                buf[1..5].copy_from_slice(&items_count.to_le_bytes());
 
                 let mut offset = 5;
-                for k in keys {
-                    buf[offset..offset + 4].copy_from_slice(&k.to_le_bytes());
+                for item in items {
+                    buf[offset..offset + 4].copy_from_slice(&item.key.to_le_bytes());
+                    offset += 4;
+
+                    let val_bytes = item.val.as_bytes();
+                    let len = val_bytes.len() as u32;
+                    buf[offset..offset + 4].copy_from_slice(&len.to_le_bytes());
                     offset += 4;
+                    buf[offset..offset + len as usize].copy_from_slice(val_bytes);
+                    offset += len as usize;
                 }
 
                 for c in children {
@@ -127,9 +138,7 @@ impl Pager {
         self.file.write_all(&buf)?;
         Ok(())
     }
-}
 
-impl Pager {
     pub fn read_page(&mut self, page_id: PageID) -> Result<Page> {
         let mut buf = vec![0u8; self.page_size];
         self.file.seek(std::io::SeekFrom::Start(
@@ -139,45 +148,26 @@ impl Pager {
 
         let page_type = buf[0];
         if page_type == 1 {
-            let keycount = u32::from_le_bytes(buf[1..5].try_into().unwrap());
-
-            if keycount > ((self.page_size - 5) / 8) as u32 {
-                return Err(std::io::Error::new(
-                    std::io::ErrorKind::InvalidData,
-                    format!(
-                        "Invalid key count: {} (max allowed: {})",
-                        keycount,
-                        (self.page_size - 5) / 8
-                    ),
-                ));
-            }
+            let items_count = u32::from_le_bytes(buf[1..5].try_into().unwrap());
 
             let mut offset = 5;
+            let mut items = Vec::with_capacity(items_count as usize);
 
-            let mut keys = Vec::with_capacity(keycount as usize);
-            for _ in 0..keycount {
-                if offset + 4 > self.page_size {
-                    return Err(std::io::Error::new(
-                        std::io::ErrorKind::InvalidData,
-                        "Buffer overflow while reading keys",
-                    ));
-                }
+            for _ in 0..items_count {
+                let key = i32::from_le_bytes(buf[offset..offset + 4].try_into().unwrap());
+                offset += 4;
 
-                keys.push(i32::from_le_bytes(
-                    buf[offset..offset + 4].try_into().unwrap(),
-                ));
+                let val_len =
+                    u32::from_le_bytes(buf[offset..offset + 4].try_into().unwrap()) as usize;
                 offset += 4;
-            }
+                let val = String::from_utf8_lossy(&buf[offset..offset + val_len]).to_string();
+                offset += val_len;
 
-            let mut children = Vec::with_capacity((keycount + 1) as usize);
-            for _ in 0..=keycount {
-                if offset + 4 > self.page_size {
-                    return Err(std::io::Error::new(
-                        std::io::ErrorKind::InvalidData,
-                        "Buffer overflow while reading children",
-                    ));
-                }
+                items.push(Item { key, val });
+            }
 
+            let mut children = Vec::with_capacity((items_count + 1) as usize);
+            for _ in 0..=items_count {
                 children.push(u32::from_le_bytes(
                     buf[offset..offset + 4].try_into().unwrap(),
                 ));
@@ -186,58 +176,21 @@ impl Pager {
 
             Ok(Page::Internal {
                 id: page_id,
-                keys,
+                items,
                 children,
             })
         } else if page_type == 0 {
             let items_count = u32::from_le_bytes(buf[1..5].try_into().unwrap());
-
-            if items_count > ((self.page_size - 5) / 20) as u32 {
-                return Err(std::io::Error::new(
-                    std::io::ErrorKind::InvalidData,
-                    format!(
-                        "Invalid items count: {} (max allowed: {})",
-                        items_count,
-                        (self.page_size - 5) / 20
-                    ),
-                ));
-            }
-
             let mut items = Vec::with_capacity(items_count as usize);
             let mut offset = 5;
 
             for _ in 0..items_count {
-                if offset + 4 > self.page_size {
-                    return Err(std::io::Error::new(
-                        std::io::ErrorKind::InvalidData,
-                        "Buffer overflow while reading item key",
-                    ));
-                }
-
                 let key = i32::from_le_bytes(buf[offset..offset + 4].try_into().unwrap());
                 offset += 4;
 
-                if offset + 4 > self.page_size {
-                    return Err(std::io::Error::new(
-                        std::io::ErrorKind::InvalidData,
-                        "Buffer overflow while reading value length",
-                    ));
-                }
-
                 let val_len =
                     u32::from_le_bytes(buf[offset..offset + 4].try_into().unwrap()) as usize;
                 offset += 4;
-
-                if val_len > self.page_size || offset + val_len > self.page_size {
-                    return Err(std::io::Error::new(
-                        std::io::ErrorKind::InvalidData,
-                        format!(
-                            "Invalid value length: {} (buffer size: {})",
-                            val_len, self.page_size
-                        ),
-                    ));
-                }
-
                 let val = String::from_utf8_lossy(&buf[offset..offset + val_len]).to_string();
                 offset += val_len;