trie.rs

pub struct TrieNode {
    vec : Vec<(char, Box<TrieNode>)>,
    word : bool,
}

impl TrieNode {
    pub fn new() -> TrieNode {
        TrieNode {
            vec: vec![],
            word: false,
        }
    }
}

pub fn get_char_node_for_insert(trie : *mut TrieNode, char : char) -> *mut TrieNode {
    let trie_ref : &mut TrieNode = unsafe { &mut *trie };
    match trie_ref.vec.binary_search_by(|&(char_,_)| char_.cmp(&char)) {
        Ok(idx) => &mut *trie_ref.vec[idx].1,
        Err(idx) => {
            trie_ref.vec.insert(idx, (char, Box::new(TrieNode { vec: vec![], word: false })));
            &mut *trie_ref.vec[idx].1
        },
    }
}

pub fn get_char_node_for_lookup(trie : &TrieNode, char : char) -> Option<&TrieNode> {
    match trie.vec.binary_search_by(|&(char_,_)| char_.cmp(&char)) {
        Ok(idx) => Some(&trie.vec[idx].1),
        Err(_) => None,
    }
}

pub fn trie_insert(trie : &mut TrieNode, str : &str) {
    let mut trie_ptr : *mut TrieNode = &mut *trie;
    for char in str.chars() {
        trie_ptr = get_char_node_for_insert(trie_ptr, char);
    }

    unsafe { (*trie_ptr).word = true; }
}

pub fn to_strings(prefix : &str, trie : &TrieNode) -> Vec<String> {
    let mut ret = {
        if trie.word {
            vec![prefix.to_owned()]
        } else {
            vec![]
        }
    };

    for &(c, ref t) in &trie.vec {
        let mut prefix_ = prefix.to_owned();
        prefix_.push(c);
        ret.extend_from_slice(&to_strings(&prefix_, &*t));
    }

    ret
}

pub fn to_pfx_strings(prefix : &str, mut trie : &TrieNode) -> Vec<String> {
    for char in prefix.chars() {
        if let Some(trie_) = get_char_node_for_lookup(trie, char) {
            trie = trie_;
        } else {
            return vec![];
        }
    }
    to_strings("", trie)
}

#[cfg(test)]
mod tests {

extern crate test;

use super::*;

#[test]
fn trie_test_1() {
    let mut trie = TrieNode::new();
    trie_insert(&mut trie, "yada yada");
    assert_eq!(vec!["yada yada"], to_strings("", &trie));
}

#[test]
fn trie_test_2() {
    let mut trie = TrieNode::new();
    trie_insert(&mut trie, "foo");
    trie_insert(&mut trie, "bar");
    trie_insert(&mut trie, "baz");
    assert_eq!(vec!["bar", "baz", "foo"], to_strings("", &trie));
}

#[test]
fn trie_test_3() {
    let mut trie = TrieNode::new();
    trie_insert(&mut trie, "foo");
    trie_insert(&mut trie, "bar");
    trie_insert(&mut trie, "baz");
    assert_eq!(vec!["ar", "az"], to_pfx_strings("b", &trie));
}

} // tests

#[cfg(bench)]
mod benchs {

extern crate test;

use self::test::Bencher;
use std::fs::File;
use std::io::Read;
use super::*;

#[bench]
fn bench_trie_build(b : &mut Bencher) {

    // Total words: 305,089
    // 117,701,680 ns (0.1 seconds)
    // (before reversing the list: 116,795,268 ns (0.1 seconds))

    let mut contents = String::new();
    let mut words : Vec<&str> = vec![];
    {
        let mut file = File::open("/usr/share/dict/american").unwrap();
        file.read_to_string(&mut contents).unwrap();
        words.extend(contents.lines());
    }

    b.iter(|| {
        let mut trie = TrieNode::new();
        // Note that we insert the words in reverse order here. Since the
        // dictionary is already sorted, we end up benchmarking the best case.
        // Since that best case is never really happens in practice, the number
        // is practically useless. Worst case is at least giving an upper bound.
        for word in words.iter().rev() {
            trie_insert(&mut trie, word);
        }
        trie
    });
}

#[bench]
fn bench_hashset_build(b : &mut Bencher) {

    // Total words: 305,089
    // 40,292,006 ns (0.04 seconds)

    use std::collections::HashSet;

    let mut contents = String::new();
    let mut words : Vec<&str> = vec![];
    {
        let mut file = File::open("/usr/share/dict/american").unwrap();
        file.read_to_string(&mut contents).unwrap();
        words.extend(contents.lines());
    }

    b.iter(|| {
        let mut set = HashSet::new();
        for word in words.iter() {
            set.insert(word);
        }
        set
    });
}

#[bench]
fn bench_trie_lookup(b : &mut Bencher) {

    // Total:     305,089 words
    // Returning:     235 words
    // 140,717 ns (0.14 ms)

    let mut contents = String::new();
    let mut words : Vec<&str> = vec![];
    {
        let mut file = File::open("/usr/share/dict/american").unwrap();
        file.read_to_string(&mut contents).unwrap();
        words.extend(contents.lines());
    }

    let mut trie = TrieNode::new();
    for word in words {
        trie_insert(&mut trie, word);
    }

    b.iter(|| {
        to_pfx_strings("abs", &trie)
    });
}

} // benchs