trie.rs

use std::fmt::Debug;
use std::fs::File;
use std::io::{BufWriter, Result};
use std::io::prelude::*;
use std::path::Path;

/// A generic case-insensitive Trie map permitting only alphabetical string keys.
pub struct Trie<T> {
    values: Vec<T>,
    /// An array of twenty-six optional tries, corresponding to each letter of the alphabet.
    /// A `None` for any value indicates that there is no such child node in the current trie.
    children: [Option<Box<Trie<T>>>; 26],
}

impl<T> Trie<T> {
    /// Creates a new trie map.
    pub fn new() -> Trie<T> {
        Trie {
            values: Vec::new(),
            children: Default::default(),
        }
    }

    /// Consumes this trie, creating a new trie with the specified key-value pair inserted.
    pub fn insert(mut self, key: &str, value: T) -> Trie<T> {
        self.insert_mut(key, value);
        self
    }

    /// Inserts the specified key-value pair into this trie via mutation.
    pub fn insert_mut(&mut self, key: &str, value: T) {
        match key.chars().next() {
            None => {
                // Adds the value to this node in the trie.
                self.values.push(value);
            },
            Some(c) => {
                let k = keyify(c);
                // Gets the k-th child or creates a new heap-allocated trie if it doesn't exist.
                let mut child = self.children[k].take().unwrap_or_else(|| Box::new(Trie::new()));
                // Inserts the remainder of the key into the k-th child.
                child.insert_mut(&key[1..], value);
                // Updates the k-th index in this trie to hold the new child.
                self.children[k] = Some(child);
            }
        }
    }
}

impl<T> Trie<T> where T: Debug {
    /// Recursively prints the values stored in the trie map.
    pub fn print(&self) {
        if self.values.len() > 0 {
            println!("{:?}", self.values);
        }
        for child in self.children.iter() {
            if let &Some(ref trie) = child {
                trie.print();
            }
        }
    }

    /// Writes the values stored in this trie map to a file at the specified path.
    pub fn write<P: AsRef<Path>>(&self, path: P) -> Result<()> {
       let file = try!(File::create(path));
       let mut writer = BufWriter::new(file);
       self.write_impl(&mut writer)
    }

    /// Recursively implements writing to a file. Mirrors print implementation.
    fn write_impl<W: Write>(&self, writer: &mut W) -> Result<()> {
        if self.values.len() > 0 {
            // format_args! is used to convert self.values into a properly formatted string.
            // write_fmt can write these formatted strings.
            // try! will return early if the operation within it fails.
            try!(writer.write_fmt(format_args!("{:?}\n", self.values)));
        }
        for child in self.children.iter() {
            if let &Some(ref trie) = child {
                try!(trie.write_impl(writer));
            }
        }
        Ok(())
    }
}

/// Converts an alphabetical character into an array index.
/// Panics on non-alphabetical characters.
fn keyify(ch: char) -> usize {
    match ch {
        'A' | 'a' => 0,
        'B' | 'b' => 1,
        'C' | 'c' => 2,
        'D' | 'd' => 3,
        'E' | 'e' => 4,
        'F' | 'f' => 5,
        'G' | 'g' => 6,
        'H' | 'h' => 7,
        'I' | 'i' => 8,
        'J' | 'j' => 9,
        'K' | 'k' => 10,
        'L' | 'l' => 11,
        'M' | 'm' => 12,
        'N' | 'n' => 13,
        'O' | 'o' => 14,
        'P' | 'p' => 15,
        'Q' | 'q' => 16,
        'R' | 'r' => 17,
        'S' | 's' => 18,
        'T' | 't' => 19,
        'U' | 'u' => 20,
        'V' | 'v' => 21,
        'W' | 'w' => 22,
        'X' | 'x' => 23,
        'Y' | 'y' => 24,
        'Z' | 'z' => 25,
        _ => panic!("Found non-alphabetical character: {}", ch)
    }
}