...

tst/src/error.rs

@@ -0,0 +1,36 @@
+//! Error types for the TST module.
+
+use thiserror::Error;
+use std::io;
+
+/// Error type for TST operations.
+#[derive(Debug, Error)]
+pub enum Error {
+    /// Error from OurDB operations.
+    #[error("OurDB error: {0}")]
+    OurDB(#[from] ourdb::Error),
+    
+    /// Error when a key is not found.
+    #[error("Key not found: {0}")]
+    KeyNotFound(String),
+    
+    /// Error when a prefix is not found.
+    #[error("Prefix not found: {0}")]
+    PrefixNotFound(String),
+    
+    /// Error during serialization.
+    #[error("Serialization error: {0}")]
+    Serialization(String),
+    
+    /// Error during deserialization.
+    #[error("Deserialization error: {0}")]
+    Deserialization(String),
+    
+    /// Error for invalid operations.
+    #[error("Invalid operation: {0}")]
+    InvalidOperation(String),
+    
+    /// IO error.
+    #[error("IO error: {0}")]
+    IO(#[from] io::Error),
+}

tst/src/lib.rs

@@ -0,0 +1,122 @@
+//! TST is a space-optimized tree data structure that enables efficient string key operations
+//! with persistent storage using OurDB as a backend.
+//!
+//! This implementation provides a persistent ternary search tree that can be used for efficient
+//! string key operations, such as auto-complete, routing tables, and more.
+
+mod error;
+mod node;
+mod operations;
+mod serialize;
+
+pub use error::Error;
+pub use node::TSTNode;
+
+use ourdb::OurDB;
+
+/// TST represents a ternary search tree data structure with persistent storage.
+pub struct TST {
+    /// Database for persistent storage
+    db: OurDB,
+    
+    /// Database ID of the root node
+    root_id: Option<u32>,
+}
+
+impl TST {
+    /// Creates a new TST with the specified database path.
+    ///
+    /// # Arguments
+    ///
+    /// * `path` - The path to the database directory
+    /// * `reset` - Whether to reset the database if it exists
+    ///
+    /// # Returns
+    ///
+    /// A new `TST` instance
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the database cannot be created or opened
+    pub fn new(path: &str, reset: bool) -> Result<Self, Error> {
+        operations::new_tst(path, reset)
+    }
+
+    /// Sets a key-value pair in the tree.
+    ///
+    /// # Arguments
+    ///
+    /// * `key` - The key to set
+    /// * `value` - The value to set
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the operation fails
+    pub fn set(&mut self, key: &str, value: Vec<u8>) -> Result<(), Error> {
+        operations::set(self, key, value)
+    }
+
+    /// Gets a value by key from the tree.
+    ///
+    /// # Arguments
+    ///
+    /// * `key` - The key to get
+    ///
+    /// # Returns
+    ///
+    /// The value associated with the key
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the key is not found or the operation fails
+    pub fn get(&mut self, key: &str) -> Result<Vec<u8>, Error> {
+        operations::get(self, key)
+    }
+
+    /// Deletes a key from the tree.
+    ///
+    /// # Arguments
+    ///
+    /// * `key` - The key to delete
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the key is not found or the operation fails
+    pub fn delete(&mut self, key: &str) -> Result<(), Error> {
+        operations::delete(self, key)
+    }
+
+    /// Lists all keys with a given prefix.
+    ///
+    /// # Arguments
+    ///
+    /// * `prefix` - The prefix to search for
+    ///
+    /// # Returns
+    ///
+    /// A list of keys that start with the given prefix
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the operation fails
+    pub fn list(&mut self, prefix: &str) -> Result<Vec<String>, Error> {
+        operations::list(self, prefix)
+    }
+
+    /// Gets all values for keys with a given prefix.
+    ///
+    /// # Arguments
+    ///
+    /// * `prefix` - The prefix to search for
+    ///
+    /// # Returns
+    ///
+    /// A list of values for keys that start with the given prefix
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the operation fails
+    pub fn getall(&mut self, prefix: &str) -> Result<Vec<Vec<u8>>, Error> {
+        operations::getall(self, prefix)
+    }
+}

tst/src/node.rs

@@ -0,0 +1,49 @@
+//! Node types for the TST module.
+
+/// Represents a node in the ternary search tree.
+#[derive(Debug, Clone, PartialEq)]
+pub struct TSTNode {
+    /// The character stored at this node.
+    pub character: char,
+    
+    /// Value stored at this node (empty if not end of key).
+    pub value: Vec<u8>,
+    
+    /// Whether this node represents the end of a key.
+    pub is_end_of_key: bool,
+    
+    /// Reference to the left child node (for characters < current character).
+    pub left_id: Option<u32>,
+    
+    /// Reference to the middle child node (for next character in key).
+    pub middle_id: Option<u32>,
+    
+    /// Reference to the right child node (for characters > current character).
+    pub right_id: Option<u32>,
+}
+
+impl TSTNode {
+    /// Creates a new node.
+    pub fn new(character: char, value: Vec<u8>, is_end_of_key: bool) -> Self {
+        Self {
+            character,
+            value,
+            is_end_of_key,
+            left_id: None,
+            middle_id: None,
+            right_id: None,
+        }
+    }
+    
+    /// Creates a new root node.
+    pub fn new_root() -> Self {
+        Self {
+            character: '\0', // Use null character for root
+            value: Vec::new(),
+            is_end_of_key: false,
+            left_id: None,
+            middle_id: None,
+            right_id: None,
+        }
+    }
+}

tst/src/operations.rs

@@ -0,0 +1,418 @@
+//! Implementation of TST operations.
+
+use crate::error::Error;
+use crate::node::TSTNode;
+use crate::TST;
+use ourdb::{OurDB, OurDBConfig, OurDBSetArgs};
+use std::path::PathBuf;
+
+/// Creates a new TST with the specified database path.
+pub fn new_tst(path: &str, reset: bool) -> Result<TST, Error> {
+    // If the path exists and reset is true, remove it first
+    let path_buf = PathBuf::from(path);
+    if path_buf.exists() && reset {
+        std::fs::remove_dir_all(&path_buf)?;
+    }
+    
+    // Create the directory if it doesn't exist
+    std::fs::create_dir_all(&path_buf)?;
+    
+    let config = OurDBConfig {
+        path: path_buf,
+        incremental_mode: true,
+        file_size: Some(1024 * 1024), // 10MB file size for better performance with large datasets
+        keysize: Some(4), // Use keysize=4 (default)
+    };
+    
+    let mut db = OurDB::new(config)?;
+    
+    let root_id = if db.get_next_id()? == 1 || reset {
+        // Create a new root node
+        let root = TSTNode::new_root();
+        let root_id = db.set(OurDBSetArgs {
+            id: None,
+            data: &root.serialize(),
+        })?;
+        
+        Some(root_id)
+    } else {
+        // Use existing root node
+        Some(1) // Root node always has ID 1
+    };
+    
+    Ok(TST {
+        db,
+        root_id,
+    })
+}
+
+/// Sets a key-value pair in the tree.
+pub fn set(tree: &mut TST, key: &str, value: Vec<u8>) -> Result<(), Error> {
+    if key.is_empty() {
+        return Err(Error::InvalidOperation("Empty key not allowed".to_string()));
+    }
+    
+    let root_id = match tree.root_id {
+        Some(id) => id,
+        None => return Err(Error::InvalidOperation("Tree not initialized".to_string())),
+    };
+    
+    let chars: Vec<char> = key.chars().collect();
+    set_recursive(tree, root_id, &chars, 0, value)?;
+    
+    Ok(())
+}
+
+/// Recursive helper function for setting a key-value pair.
+fn set_recursive(tree: &mut TST, node_id: u32, chars: &[char], pos: usize, value: Vec<u8>) -> Result<u32, Error> {
+    let mut node = tree.get_node(node_id)?;
+    
+    if pos >= chars.len() {
+        // We've reached the end of the key
+        node.is_end_of_key = true;
+        node.value = value;
+        return tree.save_node(Some(node_id), &node);
+    }
+    
+    let current_char = chars[pos];
+    
+    if node.character == '\0' {
+        // Root node or empty node, set the character
+        node.character = current_char;
+        let node_id = tree.save_node(Some(node_id), &node)?;
+        
+        // Continue with the next character
+        if pos + 1 < chars.len() {
+            let new_node = TSTNode::new(chars[pos + 1], Vec::new(), false);
+            let new_id = tree.save_node(None, &new_node)?;
+            
+            let mut updated_node = tree.get_node(node_id)?;
+            updated_node.middle_id = Some(new_id);
+            tree.save_node(Some(node_id), &updated_node)?;
+            
+            return set_recursive(tree, new_id, chars, pos + 1, value);
+        } else {
+            // This is the last character
+            let mut updated_node = tree.get_node(node_id)?;
+            updated_node.is_end_of_key = true;
+            updated_node.value = value;
+            return tree.save_node(Some(node_id), &updated_node);
+        }
+    }
+    
+    if current_char < node.character {
+        // Go left
+        if let Some(left_id) = node.left_id {
+            return set_recursive(tree, left_id, chars, pos, value);
+        } else {
+            // Create new left node
+            let new_node = TSTNode::new(current_char, Vec::new(), false);
+            let new_id = tree.save_node(None, &new_node)?;
+            
+            // Update current node
+            node.left_id = Some(new_id);
+            tree.save_node(Some(node_id), &node)?;
+            
+            return set_recursive(tree, new_id, chars, pos, value);
+        }
+    } else if current_char > node.character {
+        // Go right
+        if let Some(right_id) = node.right_id {
+            return set_recursive(tree, right_id, chars, pos, value);
+        } else {
+            // Create new right node
+            let new_node = TSTNode::new(current_char, Vec::new(), false);
+            let new_id = tree.save_node(None, &new_node)?;
+            
+            // Update current node
+            node.right_id = Some(new_id);
+            tree.save_node(Some(node_id), &node)?;
+            
+            return set_recursive(tree, new_id, chars, pos, value);
+        }
+    } else {
+        // Character matches, go middle (next character)
+        if pos + 1 >= chars.len() {
+            // This is the last character
+            node.is_end_of_key = true;
+            node.value = value;
+            return tree.save_node(Some(node_id), &node);
+        }
+        
+        if let Some(middle_id) = node.middle_id {
+            return set_recursive(tree, middle_id, chars, pos + 1, value);
+        } else {
+            // Create new middle node
+            let new_node = TSTNode::new(chars[pos + 1], Vec::new(), false);
+            let new_id = tree.save_node(None, &new_node)?;
+            
+            // Update current node
+            node.middle_id = Some(new_id);
+            tree.save_node(Some(node_id), &node)?;
+            
+            return set_recursive(tree, new_id, chars, pos + 1, value);
+        }
+    }
+}
+
+/// Gets a value by key from the tree.
+pub fn get(tree: &mut TST, key: &str) -> Result<Vec<u8>, Error> {
+    if key.is_empty() {
+        return Err(Error::InvalidOperation("Empty key not allowed".to_string()));
+    }
+    
+    let root_id = match tree.root_id {
+        Some(id) => id,
+        None => return Err(Error::InvalidOperation("Tree not initialized".to_string())),
+    };
+    
+    let chars: Vec<char> = key.chars().collect();
+    let node_id = find_node(tree, root_id, &chars, 0)?;
+    
+    let node = tree.get_node(node_id)?;
+    if node.is_end_of_key {
+        Ok(node.value.clone())
+    } else {
+        Err(Error::KeyNotFound(key.to_string()))
+    }
+}
+
+/// Finds a node by key.
+fn find_node(tree: &mut TST, node_id: u32, chars: &[char], pos: usize) -> Result<u32, Error> {
+    let node = tree.get_node(node_id)?;
+    
+    if pos >= chars.len() {
+        return Ok(node_id);
+    }
+    
+    let current_char = chars[pos];
+    
+    if current_char < node.character {
+        // Go left
+        if let Some(left_id) = node.left_id {
+            find_node(tree, left_id, chars, pos)
+        } else {
+            Err(Error::KeyNotFound(chars.iter().collect()))
+        }
+    } else if current_char > node.character {
+        // Go right
+        if let Some(right_id) = node.right_id {
+            find_node(tree, right_id, chars, pos)
+        } else {
+            Err(Error::KeyNotFound(chars.iter().collect()))
+        }
+    } else {
+        // Character matches
+        if pos + 1 >= chars.len() {
+            // This is the last character
+            Ok(node_id)
+        } else if let Some(middle_id) = node.middle_id {
+            // Go to next character
+            find_node(tree, middle_id, chars, pos + 1)
+        } else {
+            Err(Error::KeyNotFound(chars.iter().collect()))
+        }
+    }
+}
+
+/// Deletes a key from the tree.
+pub fn delete(tree: &mut TST, key: &str) -> Result<(), Error> {
+    if key.is_empty() {
+        return Err(Error::InvalidOperation("Empty key not allowed".to_string()));
+    }
+    
+    let root_id = match tree.root_id {
+        Some(id) => id,
+        None => return Err(Error::InvalidOperation("Tree not initialized".to_string())),
+    };
+    
+    let chars: Vec<char> = key.chars().collect();
+    let node_id = find_node(tree, root_id, &chars, 0)?;
+    
+    let mut node = tree.get_node(node_id)?;
+    
+    if !node.is_end_of_key {
+        return Err(Error::KeyNotFound(key.to_string()));
+    }
+    
+    // If the node has a middle child, just mark it as not end of key
+    if node.middle_id.is_some() || node.left_id.is_some() || node.right_id.is_some() {
+        node.is_end_of_key = false;
+        node.value = Vec::new();
+        tree.save_node(Some(node_id), &node)?;
+        return Ok(());
+    }
+    
+    // Otherwise, we need to remove the node and update its parent
+    // This is more complex and would require tracking the path to the node
+    // For simplicity, we'll just mark it as not end of key for now
+    node.is_end_of_key = false;
+    node.value = Vec::new();
+    tree.save_node(Some(node_id), &node)?;
+    
+    Ok(())
+}
+
+/// Lists all keys with a given prefix.
+pub fn list(tree: &mut TST, prefix: &str) -> Result<Vec<String>, Error> {
+    let root_id = match tree.root_id {
+        Some(id) => id,
+        None => return Err(Error::InvalidOperation("Tree not initialized".to_string())),
+    };
+    
+    let mut result = Vec::new();
+    
+    // Handle empty prefix case - will return all keys
+    if prefix.is_empty() {
+        collect_all_keys(tree, root_id, String::new(), &mut result)?;
+        return Ok(result);
+    }
+    
+    // Find the node corresponding to the prefix
+    let chars: Vec<char> = prefix.chars().collect();
+    let node_id = match find_prefix_node(tree, root_id, &chars, 0) {
+        Ok(id) => id,
+        Err(_) => return Ok(Vec::new()), // Prefix not found, return empty list
+    };
+    
+    // Collect all keys from the subtree
+    collect_keys_with_prefix(tree, node_id, prefix.to_string(), &mut result)?;
+    
+    Ok(result)
+}
+
+/// Finds the node corresponding to a prefix.
+fn find_prefix_node(tree: &mut TST, node_id: u32, chars: &[char], pos: usize) -> Result<u32, Error> {
+    if pos >= chars.len() {
+        return Ok(node_id);
+    }
+    
+    let node = tree.get_node(node_id)?;
+    let current_char = chars[pos];
+    
+    if current_char < node.character {
+        // Go left
+        if let Some(left_id) = node.left_id {
+            find_prefix_node(tree, left_id, chars, pos)
+        } else {
+            Err(Error::PrefixNotFound(chars.iter().collect()))
+        }
+    } else if current_char > node.character {
+        // Go right
+        if let Some(right_id) = node.right_id {
+            find_prefix_node(tree, right_id, chars, pos)
+        } else {
+            Err(Error::PrefixNotFound(chars.iter().collect()))
+        }
+    } else {
+        // Character matches
+        if pos + 1 >= chars.len() {
+            // This is the last character of the prefix
+            Ok(node_id)
+        } else if let Some(middle_id) = node.middle_id {
+            // Go to next character
+            find_prefix_node(tree, middle_id, chars, pos + 1)
+        } else {
+            Err(Error::PrefixNotFound(chars.iter().collect()))
+        }
+    }
+}
+
+/// Collects all keys with a given prefix.
+fn collect_keys_with_prefix(
+    tree: &mut TST,
+    node_id: u32,
+    current_path: String,
+    result: &mut Vec<String>,
+) -> Result<(), Error> {
+    let node = tree.get_node(node_id)?;
+    
+    // If this node is an end of key, add it to the result
+    if node.is_end_of_key {
+        result.push(current_path.clone());
+    }
+    
+    // Recursively collect keys from all children
+    if let Some(left_id) = node.left_id {
+        collect_all_keys(tree, left_id, current_path.clone(), result)?;
+    }
+    
+    if let Some(middle_id) = node.middle_id {
+        let mut new_path = current_path.clone();
+        new_path.push(node.character);
+        collect_all_keys(tree, middle_id, new_path, result)?;
+    }
+    
+    if let Some(right_id) = node.right_id {
+        collect_all_keys(tree, right_id, current_path.clone(), result)?;
+    }
+    
+    Ok(())
+}
+
+/// Recursively collects all keys under a node.
+fn collect_all_keys(
+    tree: &mut TST,
+    node_id: u32,
+    current_path: String,
+    result: &mut Vec<String>,
+) -> Result<(), Error> {
+    let node = tree.get_node(node_id)?;
+    
+    let mut new_path = current_path.clone();
+    new_path.push(node.character);
+    
+    // If this node is an end of key, add it to the result
+    if node.is_end_of_key {
+        result.push(new_path.clone());
+    }
+    
+    // Recursively collect keys from all children
+    if let Some(left_id) = node.left_id {
+        collect_all_keys(tree, left_id, current_path.clone(), result)?;
+    }
+    
+    if let Some(middle_id) = node.middle_id {
+        collect_all_keys(tree, middle_id, new_path.clone(), result)?;
+    }
+    
+    if let Some(right_id) = node.right_id {
+        collect_all_keys(tree, right_id, current_path.clone(), result)?;
+    }
+    
+    Ok(())
+}
+
+/// Gets all values for keys with a given prefix.
+pub fn getall(tree: &mut TST, 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 TST {
+    /// Helper function to get a node from the database.
+    pub(crate) fn get_node(&mut self, node_id: u32) -> Result<TSTNode, Error> {
+        let data = self.db.get(node_id)?;
+        TSTNode::deserialize(&data)
+    }
+
+    /// Helper function to save a node to the database.
+    pub(crate) fn save_node(&mut self, node_id: Option<u32>, node: &TSTNode) -> Result<u32, Error> {
+        let data = node.serialize();
+        let args = OurDBSetArgs {
+            id: node_id,
+            data: &data,
+        };
+        Ok(self.db.set(args)?)
+    }
+}

tst/src/serialize.rs

@@ -0,0 +1,134 @@
+//! Serialization and deserialization for TST nodes.
+
+use crate::error::Error;
+use crate::node::TSTNode;
+
+/// Current binary format version.
+const VERSION: u8 = 1;
+
+impl TSTNode {
+    /// Serializes a node to bytes for storage.
+    pub fn serialize(&self) -> Vec<u8> {
+        let mut buffer = Vec::new();
+        
+        // Version
+        buffer.push(VERSION);
+        
+        // Character (as UTF-32)
+        let char_bytes = (self.character as u32).to_le_bytes();
+        buffer.extend_from_slice(&char_bytes);
+        
+        // Is end of key
+        buffer.push(if self.is_end_of_key { 1 } else { 0 });
+        
+        // Value (only if is_end_of_key)
+        if self.is_end_of_key {
+            let value_len = (self.value.len() as u32).to_le_bytes();
+            buffer.extend_from_slice(&value_len);
+            buffer.extend_from_slice(&self.value);
+        } else {
+            // Zero length
+            buffer.extend_from_slice(&[0, 0, 0, 0]);
+        }
+        
+        // Child pointers
+        let left_id = self.left_id.unwrap_or(0).to_le_bytes();
+        buffer.extend_from_slice(&left_id);
+        
+        let middle_id = self.middle_id.unwrap_or(0).to_le_bytes();
+        buffer.extend_from_slice(&middle_id);
+        
+        let right_id = self.right_id.unwrap_or(0).to_le_bytes();
+        buffer.extend_from_slice(&right_id);
+        
+        buffer
+    }
+
+    /// Deserializes bytes to a node.
+    pub fn deserialize(data: &[u8]) -> Result<Self, Error> {
+        if data.len() < 14 { // Minimum size: version + char + is_end + value_len + 3 child IDs
+            return Err(Error::Deserialization("Data too short".to_string()));
+        }
+        
+        let mut pos = 0;
+        
+        // Version
+        let version = data[pos];
+        pos += 1;
+        
+        if version != VERSION {
+            return Err(Error::Deserialization(format!("Unsupported version: {}", version)));
+        }
+        
+        // Character
+        let char_bytes = [data[pos], data[pos+1], data[pos+2], data[pos+3]];
+        let char_code = u32::from_le_bytes(char_bytes);
+        let character = char::from_u32(char_code)
+            .ok_or_else(|| Error::Deserialization("Invalid character".to_string()))?;
+        pos += 4;
+        
+        // Is end of key
+        let is_end_of_key = data[pos] != 0;
+        pos += 1;
+        
+        // Value length
+        let value_len_bytes = [data[pos], data[pos+1], data[pos+2], data[pos+3]];
+        let value_len = u32::from_le_bytes(value_len_bytes) as usize;
+        pos += 4;
+        
+        // Value
+        let value = if value_len > 0 {
+            if pos + value_len > data.len() {
+                return Err(Error::Deserialization("Value length exceeds data".to_string()));
+            }
+            data[pos..pos+value_len].to_vec()
+        } else {
+            Vec::new()
+        };
+        pos += value_len;
+        
+        // Child pointers
+        if pos + 12 > data.len() {
+            return Err(Error::Deserialization("Data too short for child pointers".to_string()));
+        }
+        
+        let left_id_bytes = [data[pos], data[pos+1], data[pos+2], data[pos+3]];
+        let left_id = u32::from_le_bytes(left_id_bytes);
+        pos += 4;
+        
+        let middle_id_bytes = [data[pos], data[pos+1], data[pos+2], data[pos+3]];
+        let middle_id = u32::from_le_bytes(middle_id_bytes);
+        pos += 4;
+        
+        let right_id_bytes = [data[pos], data[pos+1], data[pos+2], data[pos+3]];
+        let right_id = u32::from_le_bytes(right_id_bytes);
+        
+        Ok(TSTNode {
+            character,
+            value,
+            is_end_of_key,
+            left_id: if left_id == 0 { None } else { Some(left_id) },
+            middle_id: if middle_id == 0 { None } else { Some(middle_id) },
+            right_id: if right_id == 0 { None } else { Some(right_id) },
+        })
+    }
+}
+
+/// Gets the common prefix of two strings.
+pub fn get_common_prefix(a: &str, b: &str) -> String {
+    let mut result = String::new();
+    let a_chars: Vec<char> = a.chars().collect();
+    let b_chars: Vec<char> = b.chars().collect();
+    
+    let min_len = a_chars.len().min(b_chars.len());
+    
+    for i in 0..min_len {
+        if a_chars[i] == b_chars[i] {
+            result.push(a_chars[i]);
+        } else {
+            break;
+        }
+    }
+    
+    result
+}