Rust trie that serializes to filesystem

fs_trie.rs

//! # A trie that can be saved to and loaded from a file
//!
//! This crate implements a Trie with char keys.
//! The trie can be saved to and loaded from a file on the local filesystem.
//! This allows the user to persist the trie between executions.
//!
//! Basic example:
//!
//! ```ignore
//! let trie_file = "/path/to/trie-file";
//! let mut trie = fs_trie::Trie::default();
//! trie.insert("abc", String::from("contents1"));
//! trie.insert("abd", String::from("contents2"));
//! trie.insert("hello", String::from("world"));
//! trie.save_to_file(trie_file).expect(
//!     "Couldn't save trie to file",
//! );
//! let trie2 = fs_trie::Trie::load_from_file(trie_file).expect("Couldn't load trie from file");
//! assert_eq!(trie, trie2);
//! ```

#![deny(missing_docs)]

#![feature(box_patterns)]
#![feature(test)]

#[cfg(test)]
#[macro_use]
extern crate quickcheck;

#[cfg(test)]
extern crate test;

#[cfg(test)]
extern crate tempfile;

extern crate serde;
#[macro_use]
extern crate serde_derive;
extern crate bincode;

use std::fs::OpenOptions;
use std::io::{self, BufReader, BufWriter};
use std::collections::HashMap;

use bincode::{serialize_into, deserialize_from, Infinite};

pub use bincode::Error as BincodeError;

type Result<T> = std::result::Result<T, BincodeError>;

/// The Trie struct. The children are a `std::collections::HashMap` of other Tries.
#[derive(Default, Debug, PartialEq, Serialize, Deserialize)]
pub struct Trie<V> {
    key: Option<char>,
    children: HashMap<Option<char>, Trie<V>>,
    contents: Option<V>,
}

impl<V> Trie<V> {
    /// Returns a Trie<V> from a file - can be empty or a previously saved trie.
    /// The type of V (the Trie values) must be known
    pub fn load_from_file(path: &str) -> Result<Self>
    where
        for<'de> V: serde::Serialize + serde::Deserialize<'de>,
    {
        let f = OpenOptions::new()
            .read(true)
            .write(true)
            .create(true)
            .open(path)
            .expect("Couldn't open trie file");
        let mut br = BufReader::new(f);
        match deserialize_from(&mut br, Infinite) {
            Ok(x) => Ok(x),
            Err(box bincode::ErrorKind::IoError(e)) => {
                if e.kind() == io::ErrorKind::UnexpectedEof {
                    return Ok(Trie {
                        key: None,
                        children: HashMap::new(),
                        contents: None,
                    });
                }
                Err(Box::new(bincode::ErrorKind::IoError(e)))
            }
            Err(e) => Err(e),
        }
    }

    /// Inserts a Trie entry. The &str key is split into its chars to generate
    /// the children.
    pub fn insert(&mut self, key: &str, contents: V) -> Option<V> {
        let mut chars = key.chars();
        let mut key_i_need = chars.next();
        if self.key == key_i_need {
            if chars.size_hint().0 == 0 {
                let ret = self.contents.take();
                self.contents = Some(contents);
                return ret;
            }
            key_i_need = chars.next();
        }
        if let Some(c) = self.children.get_mut(&key_i_need) {
            return c.insert(chars.as_str(), contents);
        }
        let mut trie = Trie {
            key: key_i_need,
            children: HashMap::new(),
            contents: None,
        };
        trie.insert(chars.as_str(), contents);
        self.children.insert(key_i_need, trie);
        None
    }

    /// Get an entry from the Trie. Traverses the trie with the chars in the key.
    pub fn get(&self, key: &str) -> Option<&V> {
        let mut chars = key.chars();
        let mut key_i_need = chars.next();
        if self.key == key_i_need {
            if chars.size_hint().0 == 0 {
                return self.contents.as_ref();
            }
            key_i_need = chars.next();
        }
        if let Some(c) = self.children.get(&key_i_need) {
            return c.get(chars.as_str());
        }
        None
    }

    /// Get a mutable reference to the Trie entry. Use this to modify entries.
    pub fn get_mut(&mut self, key: &str) -> Option<&mut V> {
        let mut chars = key.chars();
        let mut key_i_need = chars.next();
        if self.key == key_i_need {
            if chars.size_hint().0 == 0 {
                return self.contents.as_mut();
            }
            key_i_need = chars.next();
        }
        if let Some(c) = self.children.get_mut(&key_i_need) {
            return c.get_mut(chars.as_str());
        }
        None
    }

    /// Save the Trie to a file.
    pub fn save_to_file(&mut self, path: &str) -> Result<()>
    where
        for<'de> V: serde::Serialize + serde::Deserialize<'de>,
    {
        let f = OpenOptions::new()
            .read(true)
            .write(true)
            .create(true)
            .open(path)
            .expect("Couldn't open trie file");
        let mut bw = BufWriter::new(f);
        serialize_into(&mut bw, self, Infinite)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::NamedTempFile;
    use std::collections::BTreeMap;
    use quickcheck::TestResult;
    use test::Bencher;

    fn insertion_test_helper(mut v: Vec<(String, String)>, replace: bool) -> TestResult {
        let v = v.iter_mut()
            .map(|&mut (ref i, ref j)| {
                let v = if !replace { j } else { "this_will_be_replaced" };
                (i, v)
            })
            .collect::<Vec<_>>();
        let mut t = Trie::default();
        let mut bt = BTreeMap::new();
        for &(i, j) in v.iter() {
            assert_eq!(t.insert(i, j), bt.insert(i, j));
        }
        for &(i, _) in v.iter() {
            assert_eq!(t.get(i), bt.get(i));
        }
        TestResult::from_bool(true)
    }

    #[test]
    fn basic_insertion() -> () {
        let testcases = vec![
            (String::from("def"), String::from("contents1")),
            (String::from("abc"), String::from("contents2")),
            (String::from("abf"), String::from("contents3")),
        ];
        insertion_test_helper(testcases, false);
    }

    #[test]
    fn test_get_mut() -> () {
        let mut trie = Trie::default();
        trie.insert("abc", vec!["test1"]);

        {
            let thing_to_modify = trie.get_mut("abc").unwrap();
            thing_to_modify.push("test2");
        }

        assert_eq!(*trie.get("abc").unwrap(), vec!["test1", "test2"]);
    }

    quickcheck! {
        fn random_insertion(v: Vec<(String, String)>) -> TestResult {
            insertion_test_helper(v, false)
        }
        fn replace_insertion(v: Vec<(String, String)>) -> TestResult {
            insertion_test_helper(v, true)
        }
    }

    #[test]
    fn save_to_file_roundtrip() -> () {
        let trie_file = NamedTempFile::new().expect("failed to create temporary file");
        let trie_file_name = trie_file.path().to_str().unwrap();

        let mut trie = Trie::default();
        trie.insert("abc", String::from("contents1"));
        trie.insert("abd", String::from("contents2"));
        trie.insert("hello", String::from("world"));

        trie.save_to_file(trie_file_name).expect(
            "Couldn't save trie to file",
        );
        let trie2 = Trie::load_from_file(trie_file_name).expect("Couldn't load trie from file");
        assert_eq!(trie, trie2);
    }

    #[test]
    fn load_from_empty_file() -> () {
        let trie_file = NamedTempFile::new().expect("failed to create temporary file");
        let trie_file_name = trie_file.path().to_str().unwrap();

        let trie = Trie::<Trie<String>>::load_from_file(trie_file_name)
            .expect("Couldn't load trie from file");
        assert_eq!(trie.key, None);
        assert_eq!(trie.children.len(), 0);
        assert_eq!(trie.contents, None);
    }

    #[bench]
    fn bench_many_children(b: &mut Bencher) {
        let mut trie = Trie::default();
        let utf8_max_char = 128;
        let last_child = &String::from_utf8(vec![utf8_max_char - 1]).unwrap();
        for i in 0..utf8_max_char {
            trie.insert(&String::from_utf8(vec![i]).unwrap(), "");
        }
        assert_eq!((utf8_max_char) as usize, trie.children.len());
        b.iter(|| trie.get(&last_child));
    }
}