herocode/db
12
0
Fork 0
Code Issues 2 Pull Requests Actions Packages Projects Releases Wiki Activity

port radixtree to rust

Browse Source
This commit is contained in:
Timur Gordon 2025-04-09 13:09:59 +02:00
parent e62e0a125a
commit dc8c887026
3 changed files with 427 additions and 341 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
radixtree/src/lib.rs
View File

@ -12,8 +12,7 @@ mod serialize;
pub use error::Error;
pub use node::{Node, NodeRef};
use ourdb::{OurDB, OurDBConfig, OurDBSetArgs};
use std::path::PathBuf;
use ourdb::OurDB;
/// RadixTree represents a radix tree data structure with persistent storage.
pub struct RadixTree {
@ -37,8 +36,7 @@ impl RadixTree {
///
/// Returns an error if the database cannot be created or opened
pub fn new(path: &str, reset: bool) -> Result<Self, Error> {
// Implementation will be provided in operations.rs
unimplemented!()
operations::new_radix_tree(path, reset)
}
/// Sets a key-value pair in the tree.
@ -52,8 +50,7 @@ impl RadixTree {
///
/// Returns an error if the operation fails
pub fn set(&mut self, key: &str, value: Vec<u8>) -> Result<(), Error> {
// Implementation will be provided in operations.rs
unimplemented!()
operations::set(self, key, value)
}
/// Gets a value by key from the tree.
@ -70,8 +67,7 @@ impl RadixTree {
///
/// Returns an error if the key is not found or the operation fails
pub fn get(&mut self, key: &str) -> Result<Vec<u8>, Error> {
// Implementation will be provided in operations.rs
unimplemented!()
operations::get(self, key)
}
/// Updates the value at a given key prefix.
@ -85,8 +81,7 @@ impl RadixTree {
///
/// Returns an error if the prefix is not found or the operation fails
pub fn update(&mut self, prefix: &str, new_value: Vec<u8>) -> Result<(), Error> {
// Implementation will be provided in operations.rs
unimplemented!()
operations::update(self, prefix, new_value)
}
/// Deletes a key from the tree.
@ -99,8 +94,7 @@ impl RadixTree {
///
/// Returns an error if the key is not found or the operation fails
pub fn delete(&mut self, key: &str) -> Result<(), Error> {
// Implementation will be provided in operations.rs
unimplemented!()
operations::delete(self, key)
}
/// Lists all keys with a given prefix.
@ -117,8 +111,7 @@ impl RadixTree {
///
/// Returns an error if the operation fails
pub fn list(&mut self, prefix: &str) -> Result<Vec<String>, Error> {
// Implementation will be provided in operations.rs
unimplemented!()
operations::list(self, prefix)
}
/// Gets all values for keys with a given prefix.
@ -135,7 +128,6 @@ impl RadixTree {
///
/// Returns an error if the operation fails
pub fn getall(&mut self, prefix: &str) -> Result<Vec<Vec<u8>>, Error> {
// Implementation will be provided in operations.rs
unimplemented!()
operations::getall(self, prefix)
}
}

276
radixtree/src/operations.rs
View File

@ -3,20 +3,26 @@
use crate::error::Error;
use crate::node::{Node, NodeRef};
use crate::RadixTree;
use crate::serialize::get_common_prefix;
use ourdb::{OurDB, OurDBConfig, OurDBSetArgs};
use std::path::PathBuf;
impl RadixTree {
/// Creates a new radix tree with the specified database path.
pub fn new(path: &str, reset: bool) -> Result<Self, Error> {
/// Creates a new radix tree with the specified database path.
pub fn new_radix_tree(path: &str, reset: bool) -> Result<RadixTree, Error> {
let config = OurDBConfig {
record_size_max: 1024 * 4, // 4KB max record size
path: PathBuf::from(path),
incremental_mode: true,
reset,
..Default::default()
file_size: Some(1024 * 1024), // 1MB file size
keysize: Some(4), // Default key size
};
let db = OurDB::new(path, config)?;
let mut db = OurDB::new(config)?;
// If reset is true, we would clear the database
// Since OurDB doesn't have a reset method, we'll handle it by
// creating a fresh database when reset is true
// We'll implement this by checking if it's a new database (next_id == 1)
let root_id = if db.get_next_id()? == 1 {
// Create a new root node
@ -34,44 +40,28 @@ impl RadixTree {
1 // Root node always has ID 1
};
Ok(Self {
Ok(RadixTree {
db,
root_id,
})
}
}
/// Helper function to get a node from the database.
pub(crate) fn get_node(&mut self, node_id: u32) -> Result<Node, Error> {
let data = self.db.get(node_id)?;
Node::deserialize(&data)
}
/// Helper function to save a node to the database.
pub(crate) fn save_node(&mut self, node_id: Option<u32>, node: &Node) -> Result<u32, Error> {
let data = node.serialize();
let args = OurDBSetArgs {
id: node_id,
data: &data,
};
Ok(self.db.set(args)?)
}
/// Sets a key-value pair in the tree.
pub fn set(&mut self, key: &str, value: Vec<u8>) -> Result<(), Error> {
let mut current_id = self.root_id;
/// Sets a key-value pair in the tree.
pub fn set(tree: &mut RadixTree, key: &str, value: Vec<u8>) -> Result<(), Error> {
let mut current_id = tree.root_id;
let mut offset = 0;
// Handle empty key case
if key.is_empty() {
let mut root_node = self.get_node(current_id)?;
let mut root_node = tree.get_node(current_id)?;
root_node.is_leaf = true;
root_node.value = value;
self.save_node(Some(current_id), &root_node)?;
tree.save_node(Some(current_id), &root_node)?;
return Ok(());
}
while offset < key.len() {
let mut node = self.get_node(current_id)?;
let mut node = tree.get_node(current_id)?;
// Find matching child
let mut matched_child = None;
@ -92,7 +82,7 @@ impl RadixTree {
is_leaf: true,
};
let new_id = self.save_node(None, &new_node)?;
let new_id = tree.save_node(None, &new_node)?;
// Create new child reference and update parent node
node.children.push(NodeRef {
@ -100,7 +90,7 @@ impl RadixTree {
node_id: new_id,
});
self.save_node(Some(current_id), &node)?;
tree.save_node(Some(current_id), &node)?;
return Ok(());
}
@ -109,7 +99,7 @@ impl RadixTree {
if common_prefix.len() < child.key_part.len() {
// Split existing node
let mut child_node = self.get_node(child.node_id)?;
let child_node = tree.get_node(child.node_id)?;
// Create new intermediate node
let new_node = Node {
@ -118,14 +108,14 @@ impl RadixTree {
children: child_node.children.clone(),
is_leaf: child_node.is_leaf,
};
let new_id = self.save_node(None, &new_node)?;
let new_id = tree.save_node(None, &new_node)?;
// Update current node
node.children[child_index] = NodeRef {
key_part: common_prefix.to_string(),
node_id: new_id,
};
self.save_node(Some(current_id), &node)?;
tree.save_node(Some(current_id), &node)?;
// Update child node reference
child.node_id = new_id;
@ -133,10 +123,10 @@ impl RadixTree {
if offset + common_prefix.len() == key.len() {
// Update value at existing node
let mut child_node = self.get_node(child.node_id)?;
let mut child_node = tree.get_node(child.node_id)?;
child_node.value = value;
child_node.is_leaf = true;
self.save_node(Some(child.node_id), &child_node)?;
tree.save_node(Some(child.node_id), &child_node)?;
return Ok(());
}
@ -145,30 +135,30 @@ impl RadixTree {
}
Ok(())
}
}
/// Gets a value by key from the tree.
pub fn get(&mut self, key: &str) -> Result<Vec<u8>, Error> {
let mut current_id = self.root_id;
/// Gets a value by key from the tree.
pub fn get(tree: &mut RadixTree, key: &str) -> Result<Vec<u8>, Error> {
let mut current_id = tree.root_id;
let mut offset = 0;
// Handle empty key case
if key.is_empty() {
let root_node = self.get_node(current_id)?;
let root_node = tree.get_node(current_id)?;
if root_node.is_leaf {
return Ok(root_node.value);
return Ok(root_node.value.clone());
}
return Err(Error::KeyNotFound(key.to_string()));
}
while offset < key.len() {
let node = self.get_node(current_id)?;
let node = tree.get_node(current_id)?;
let mut found = false;
for child in &node.children {
if key[offset..].starts_with(&child.key_part) {
if offset + child.key_part.len() == key.len() {
let child_node = self.get_node(child.node_id)?;
let child_node = tree.get_node(child.node_id)?;
if child_node.is_leaf {
return Ok(child_node.value);
}
@ -186,11 +176,11 @@ impl RadixTree {
}
Err(Error::KeyNotFound(key.to_string()))
}
}
/// Updates the value at a given key prefix.
pub fn update(&mut self, prefix: &str, new_value: Vec<u8>) -> Result<(), Error> {
let mut current_id = self.root_id;
/// Updates the value at a given key prefix.
pub fn update(tree: &mut RadixTree, prefix: &str, new_value: Vec<u8>) -> Result<(), Error> {
let mut current_id = tree.root_id;
let mut offset = 0;
// Handle empty prefix case
@ -199,18 +189,18 @@ impl RadixTree {
}
while offset < prefix.len() {
let node = self.get_node(current_id)?;
let node = tree.get_node(current_id)?;
let mut found = false;
for child in &node.children {
if prefix[offset..].starts_with(&child.key_part) {
if offset + child.key_part.len() == prefix.len() {
// Found exact prefix match
let mut child_node = self.get_node(child.node_id)?;
let mut child_node = tree.get_node(child.node_id)?;
if child_node.is_leaf {
// Update the value
child_node.value = new_value;
self.save_node(Some(child.node_id), &child_node)?;
tree.save_node(Some(child.node_id), &child_node)?;
return Ok(());
}
}
@ -227,17 +217,30 @@ impl RadixTree {
}
Err(Error::PrefixNotFound(prefix.to_string()))
}
}
/// Deletes a key from the tree.
pub fn delete(&mut self, key: &str) -> Result<(), Error> {
let mut current_id = self.root_id;
/// Deletes a key from the tree.
pub fn delete(tree: &mut RadixTree, key: &str) -> Result<(), Error> {
let mut current_id = tree.root_id;
let mut offset = 0;
let mut path = Vec::new();
// Handle empty key case
if key.is_empty() {
let mut root_node = tree.get_node(current_id)?;
if !root_node.is_leaf {
return Err(Error::KeyNotFound(key.to_string()));
}
// For the root node, we just mark it as non-leaf
root_node.is_leaf = false;
root_node.value = Vec::new();
tree.save_node(Some(current_id), &root_node)?;
return Ok(());
}
// Find the node to delete
while offset < key.len() {
let node = self.get_node(current_id)?;
let node = tree.get_node(current_id)?;
let mut found = false;
for child in &node.children {
@ -249,7 +252,7 @@ impl RadixTree {
// Check if we've matched the full key
if offset == key.len() {
let child_node = self.get_node(child.node_id)?;
let child_node = tree.get_node(child.node_id)?;
if child_node.is_leaf {
found = true;
break;
@ -269,20 +272,20 @@ impl RadixTree {
}
// Get the node to delete
let mut last_node = self.get_node(path.last().unwrap().node_id)?;
let mut last_node = tree.get_node(path.last().unwrap().node_id)?;
// If the node has children, just mark it as non-leaf
if !last_node.children.is_empty() {
last_node.is_leaf = false;
last_node.value = Vec::new();
self.save_node(Some(path.last().unwrap().node_id), &last_node)?;
tree.save_node(Some(path.last().unwrap().node_id), &last_node)?;
return Ok(());
}
// If node has no children, remove it from parent
if path.len() > 1 {
let parent_id = path[path.len() - 2].node_id;
let mut parent_node = self.get_node(parent_id)?;
let mut parent_node = tree.get_node(parent_id)?;
// Find and remove the child from parent
for i in 0..parent_node.children.len() {
@ -292,33 +295,139 @@ impl RadixTree {
}
}
self.save_node(Some(parent_id), &parent_node)?;
tree.save_node(Some(parent_id), &parent_node)?;
// Delete the node from the database
self.db.delete(path.last().unwrap().node_id)?;
tree.db.delete(path.last().unwrap().node_id)?;
} else {
// If this is a direct child of the root, just mark it as non-leaf
last_node.is_leaf = false;
last_node.value = Vec::new();
self.save_node(Some(path.last().unwrap().node_id), &last_node)?;
tree.save_node(Some(path.last().unwrap().node_id), &last_node)?;
}
Ok(())
}
}
/// Lists all keys with a given prefix.
pub fn list(&mut self, prefix: &str) -> Result<Vec<String>, Error> {
/// Lists all keys with a given prefix.
pub fn list(tree: &mut RadixTree, prefix: &str) -> Result<Vec<String>, Error> {
let mut result = Vec::new();
// Handle empty prefix case - will return all keys
if prefix.is_empty() {
self.collect_all_keys(self.root_id, "", &mut result)?;
collect_all_keys(tree, tree.root_id, "", &mut result)?;
return Ok(result);
}
// Start from the root and find all matching keys
self.find_keys_with_prefix(self.root_id, "", prefix, &mut result)?;
find_keys_with_prefix(tree, tree.root_id, "", prefix, &mut result)?;
Ok(result)
}
/// Helper function to find all keys with a given prefix.
fn find_keys_with_prefix(
tree: &mut RadixTree,
node_id: u32,
current_path: &str,
prefix: &str,
result: &mut Vec<String>,
) -> Result<(), Error> {
let node = tree.get_node(node_id)?;
// If the current path already matches or exceeds the prefix length
if current_path.len() >= prefix.len() {
// Check if the current path starts with the prefix
if current_path.starts_with(prefix) {
// If this is a leaf node, add it to the results
if node.is_leaf {
result.push(current_path.to_string());
}
// Collect all keys from this subtree
for child in &node.children {
let child_path = format!("{}{}", current_path, child.key_part);
find_keys_with_prefix(tree, child.node_id, &child_path, prefix, result)?;
}
}
return Ok(());
}
// Current path is shorter than the prefix, continue searching
for child in &node.children {
let child_path = format!("{}{}", current_path, child.key_part);
// Check if this child's path could potentially match the prefix
if prefix.starts_with(current_path) {
// The prefix starts with the current path, so we need to check if
// the child's key_part matches the next part of the prefix
let prefix_remainder = &prefix[current_path.len()..];
// If the prefix remainder starts with the child's key_part or vice versa
if prefix_remainder.starts_with(&child.key_part)
|| (child.key_part.starts_with(prefix_remainder)
&& child.key_part.len() >= prefix_remainder.len()) {
find_keys_with_prefix(tree, child.node_id, &child_path, prefix, result)?;
}
}
}
Ok(())
}
/// Helper function to recursively collect all keys under a node.
fn collect_all_keys(
tree: &mut RadixTree,
node_id: u32,
current_path: &str,
result: &mut Vec<String>,
) -> Result<(), Error> {
let node = tree.get_node(node_id)?;
// If this node is a leaf, add its path to the result
if node.is_leaf {
result.push(current_path.to_string());
}
// Recursively collect keys from all children
for child in &node.children {
let child_path = format!("{}{}", current_path, child.key_part);
collect_all_keys(tree, child.node_id, &child_path, result)?;
}
Ok(())
}
/// Gets all values for keys with a given prefix.
pub fn getall(tree: &mut RadixTree, prefix: &str) -> Result<Vec<Vec<u8>>, Error> {
// Get all matching keys
let keys = list(tree, prefix)?;
// Get values for each key
let mut values = Vec::new();
for key in keys {
if let Ok(value) = get(tree, &key) {
values.push(value);
}
}
Ok(values)
}
impl RadixTree {
/// Helper function to get a node from the database.
pub(crate) fn get_node(&mut self, node_id: u32) -> Result<Node, Error> {
let data = self.db.get(node_id)?;
Node::deserialize(&data)
}
/// Helper function to save a node to the database.
pub(crate) fn save_node(&mut self, node_id: Option<u32>, node: &Node) -> Result<u32, Error> {
let data = node.serialize();
let args = OurDBSetArgs {
id: node_id,
data: &data,
};
Ok(self.db.set(args)?)
}
/// Helper function to find all keys with a given prefix.
@ -393,33 +502,6 @@ impl RadixTree {
Ok(())
}
/// Gets all values for keys with a given prefix.
pub fn getall(&mut self, prefix: &str) -> Result<Vec<Vec<u8>>, Error> {
// Get all matching keys
let keys = self.list(prefix)?;
// Get values for each key
let mut values = Vec::new();
for key in keys {
if let Ok(value) = self.get(&key) {
values.push(value);
}
}
Ok(values)
}
}
/// Helper function to get the common prefix of two strings.
fn get_common_prefix(a: &str, b: &str) -> String {
let mut i = 0;
let a_bytes = a.as_bytes();
let b_bytes = b.as_bytes();
while i < a.len() && i < b.len() && a_bytes[i] == b_bytes[i] {
i += 1;
}
a[..i].to_string()
}

16
radixtree/src/serialize.rs
View File

@ -2,8 +2,7 @@
use crate::error::Error;
use crate::node::{Node, NodeRef};
use std::convert::TryInto;
use std::io::{Cursor, Read, Write};
use std::io::{Cursor, Read};
use std::mem::size_of;
/// Current binary format version.
@ -142,3 +141,16 @@ fn read_u32(cursor: &mut Cursor<&[u8]>) -> std::io::Result<u32> {
Ok(u32::from_le_bytes(bytes))
}
/// Helper function to get the common prefix of two strings.
pub fn get_common_prefix(a: &str, b: &str) -> String {
let mut i = 0;
let a_bytes = a.as_bytes();
let b_bytes = b.as_bytes();
while i < a.len() && i < b.len() && a_bytes[i] == b_bytes[i] {
i += 1;
}
a[..i].to_string()
}