Last active
July 26, 2024 14:06
-
-
Save dharanad/2b77fc83c32e2ddd1118fcd2c4d11a3f to your computer and use it in GitHub Desktop.
A copy on write Trie
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
use std::collections::HashMap; | |
use std::hash::Hash; | |
use criterion::{black_box, Criterion, criterion_group, criterion_main}; | |
#[derive(Clone)] | |
struct TrieNode<T: Clone> { | |
value: Option<T>, | |
children: HashMap<char, Box<TrieNode<T>>>, | |
} | |
impl<T> TrieNode<T> | |
where | |
T: Clone, | |
{ | |
pub fn new(value: Option<T>, children: HashMap<char, Box<TrieNode<T>>>) -> Self { | |
Self { value, children } | |
} | |
fn child(&self, key: &char) -> Option<&Box<TrieNode<T>>> { | |
self.children.get(key) | |
} | |
fn value(&self) -> Option<T> { | |
self.value.clone() | |
} | |
} | |
pub struct Trie<T> | |
where | |
T: Clone, | |
{ | |
root: Option<Box<TrieNode<T>>>, | |
versions: Vec<Option<Box<TrieNode<T>>>>, | |
current_version: u32, | |
} | |
impl<T> Trie<T> | |
where | |
T: Clone, | |
{ | |
pub fn new() -> Self { | |
Self { | |
root: None, | |
versions: vec![], | |
current_version: 0, | |
} | |
} | |
} | |
impl<T> Trie<T> | |
where | |
T: Clone, | |
{ | |
fn update_version(&mut self) -> u32 { | |
self.current_version += 1; | |
self.current_version | |
} | |
pub fn insert(&mut self, key: &str, value: T) -> u32 { | |
let key_arr = key.chars().collect::<Vec<char>>(); | |
let mut new_root = Some(Self::insert_inner(self.root.as_ref(), &key_arr, 0, value)); | |
std::mem::swap(&mut self.root, &mut new_root); | |
self.versions.push(new_root); | |
return self.update_version(); | |
} | |
fn insert_inner( | |
root: Option<&Box<TrieNode<T>>>, | |
key: &[char], | |
idx: usize, | |
value: T, | |
) -> Box<TrieNode<T>> { | |
// get a children map since we have to clone this map | |
let mut children_map = root.map(|r| r.children.clone()).unwrap_or_else(HashMap::new); | |
if idx == key.len() { | |
Box::new(TrieNode::new(Some(value), children_map)) | |
} else { | |
children_map.insert( | |
key[idx], | |
Self::insert_inner(children_map.get(&key[idx]), key, idx + 1, value), | |
); | |
Box::new(TrieNode::new(root.and_then(|r| r.value.clone()), children_map)) | |
} | |
} | |
pub fn search(&self, version: u32, key: &str) -> Option<T> { | |
assert!(version < self.versions.len() as u32); | |
if let Some(ref ptr) = &self.versions[version as usize] { | |
let mut ptr = ptr; | |
let chars = key.chars().collect::<Vec<char>>(); | |
for c in chars.iter() { | |
match ptr.child(c) { | |
None => { | |
return None; | |
} | |
Some(c) => ptr = c, | |
} | |
} | |
return ptr.value.clone(); | |
} else { | |
None | |
} | |
} | |
} | |
#[cfg(test)] | |
mod tests { | |
use crate::trie::Trie; | |
#[test] | |
fn test_trie_insert() { | |
let mut trie = Trie::new(); | |
let keys = vec!["cat", "dog", "mouse", "tom"]; | |
let mut counter = 0; | |
let mut prev_version = 0; | |
for k in keys.into_iter() { | |
let new_version = trie.insert(k, counter); | |
counter += 1; | |
assert_eq!(new_version, prev_version + 1); | |
prev_version = new_version; | |
} | |
} | |
#[test] | |
fn test_trie_search() { | |
let mut trie = Trie::new(); | |
// value 100 101 102 103 | |
// version 1 2 3 4 | |
let keys = vec!["cat", "dog", "mouse", "tom"]; | |
let mut counter = 100; | |
for k in keys.into_iter() { | |
let new_version = trie.insert(k, counter); | |
counter += 1; | |
} | |
assert_eq!(trie.search(1, "cat"), Some(100)); | |
assert_eq!(trie.search(3, "mouse"), Some(102)); | |
assert_eq!(trie.search(2, "mouse"), None); | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment