A copy on write Trie

cow_trie.rs

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);
    }
}