Peter Norvigs spellcheck in Rust

Question

For a Rust project I'm working on I needed a (simple) spellcheck algorithm and I set out to re-implement Peter Norvigs spellcheck algorithm in Rust.

extern crate regex;

use regex::Regex;

use std::collections::HashMap;
use std::collections::HashSet;
use std::iter::FromIterator;
use std::path::Path;
use std::io::prelude::*;
use std::fs::File;

static ALPHABET : &'static str = "abcdefghijklmnopqrstuvwxyz";

fn deletes(s: &str) -> Vec<String> {
    (0..s.len())
        .map(|i| {
            let mut delete = s.to_owned();
            delete.remove(i);

            delete
        })
        .collect()
}

fn inserts(s: &str) -> Vec<String> {
    (0..s.len() + 1)
        .flat_map(|i| {
            ALPHABET.chars().map(|chr| {
                let mut insert = s.to_owned();
                insert.insert(i, chr);

                insert
            }).collect::<Vec<_>>()
        })
        .collect()
}

fn replaces(s: &str) -> Vec<String> {
    (0..s.len())
        .flat_map(|i| {
            ALPHABET.chars().map(|chr| {
                let mut replace = String::with_capacity(s.len());
                replace.push_str(&s[0..i]);
                replace.push(chr);
                replace.push_str(&s[i + 1..]);

                replace
            }).collect::<Vec<_>>()
        })
        .collect()
}

fn transposes(s: &str) -> Vec<String> {
    let bytes = s.as_bytes();

    (1..bytes.len())
        .map(|i| {
            let mut transpose = bytes.to_owned();
            transpose.swap(i - 1, i);

            String::from_utf8(transpose).expect("Invalid UTF-8")
        })
        .collect()
}

fn read_dictionary_words(dictionary_path: &str) -> String {
    let mut buffer = String::new();

    File::open(&Path::new(dictionary_path))
        .unwrap()
        .read_to_string(&mut buffer)
        .unwrap();

    buffer
}

fn word_count(all_words: &str) ->  HashMap<String, u32> {
    let re = Regex::new(r"\w+").unwrap();
    let lowercase = all_words.to_lowercase();

    let mut map = HashMap::new();

    for word in re.captures_iter(&lowercase) {
        // arg, the to_owned is pretty nasty here, as it creates a new owned string every time ...
        *map.entry(word[0].to_owned()).or_insert(0) += 1;
    }

    map
}

fn edits1(s: &str) -> HashSet<String> {
    deletes(s)
        .into_iter()
        .chain(replaces(s).into_iter())
        .chain(transposes(s).into_iter())
        .chain(inserts(s).into_iter())
        .collect()
}

fn edits2(s: &str) -> HashSet<String> {
    edits1(s)
        .iter()
        .flat_map(|e| edits1(e))
        .collect()
}

fn calc_p(word: &str, wc: &HashMap<String, u32>) -> f32 {
    if !wc.contains_key(word) {
        return 0.0;
    }

    let total:u32 = wc.values().sum();
    wc[word] as f32 / total as f32
}

fn known(candidates: HashSet<String>, wc: &HashMap<String, u32>) -> HashSet<String> {
    candidates
        .into_iter()
        .filter(|w| wc.contains_key(w))
        .collect::<HashSet<_>>()
}

fn create_candidates(s: &str, wc: &HashMap<String, u32>) -> HashSet<String> {
    if wc.contains_key(s) {
        return HashSet::from_iter(vec![s.to_owned()].into_iter());
    }

    let candidates = known(edits1(s), wc);

    if candidates.len() > 0 {
        return candidates;
    }

    let candidates = known(edits2(s), wc);

    if candidates.len() > 0 {
        return candidates;
    }

    HashSet::new()
}

fn correction(candidates: HashSet<String>, wc: &HashMap<String, u32>) -> Vec<(f32, String)> {
    let mut weighted_candidates = candidates
        .into_iter()
        .map(|w| (calc_p(&w, &wc), w))
        .collect::<Vec<_>>();

    weighted_candidates.sort_by(|a, b| b.partial_cmp(a).unwrap());

    weighted_candidates
}

fn main() {
    // This is the file from http://www.norvig.com/big.txt
    let all_words = read_dictionary_words("/tmp/big.txt");
    let wc = word_count(&all_words);

    let corrections = correction(create_candidates("speling", &wc), &wc);
    println!("{:?}", corrections.first().unwrap());

    let corrections = correction(create_candidates("korrectud", &wc), &wc);
    println!("{:?}", corrections.first().unwrap());
}

The output is:

(0.0000035855628, "spelling") // suggestion for "speling"
(0.00001254947, "corrected")  // suggestion for "korrectud"

Which seems to be roughly in line with what Peters Python version does.

The code is basically runnable, but you need to have the file big.txt which contains words from various text sources combined into one blob. The file can be found here (BEWARE: 6.5 MB text file ahead!)

• I refactored already the deletes, transposes, inserts, replaces functions thanks to another code review here.
• In general I would be interested if I got ownership right. I fought less with the borrow checker this time however.
• There is a line in the word_count function that's particular nasty, it does a string allocation every time for checking if the key is in the HashMap but puts it in only once.
• The performance is OK for my usecase but not really impressive.

All other comments/suggestions are welcome, too, of course.

My usecase is only about ASCII strings. I thought about using the ASCII crate types to make that explicit, but not in this version, it's a fair suggestion of course.

Answer 1

A few comments, since no one else has posted here:

Set up some proper error handling so you don't have so many unwraps and expects. I think the error_chain crate is the recommended/most popular way to do error handling, but is probably overkill for your case. However, error_chain also provides the quick_error macro, which allows you to easily define your own errors and wrap external errors. For very, very basic handling you can just do:

pub mod err {
    quick_error! {
        #[derive(Debug)]
        pub enum Error {
            Other(err: Box<::std::error::Error + Send + Sync>) {
                from(e: &'static str) -> (e.into())
                from(e: ::std::io::Error) -> (e.into())                
                description(err.description())
                display("{}", err)
            }    
        }
    }
    pub type Result<T> = ::std::result::Result<T, Error>;
}

And then wrap your functions in Result<_>, and you can use try/? everywhere.

---

Not really a big deal for your case since you only call it once, but it is good practice to wrap Regex::new in lazy_static to avoid compiling it multiple times, e.g.:

lazy_static! {
    static ref RE: Regex = Regex::new(r"\w+").unwrap();
}

---

You can work around the String allocation with the Entry API like (not sure if there is a more graceful way to do this...):

for word in re.captures_iter(&lowercase) {
    let word0 = word[0];
    if let Some(e) = map.get_mut(word0) {
        *e += 1;
        continue;
    }
    map.insert(word0.into(), 1);
}

---

I would suggest returning an Option or Result from known where the empty case returns a None like type, so that you don't have to check candidates.len() > 0 multiple times in create_candidates. You could chain the different candidates then return the default if none apply.

---

ALSO! If your performance seems bad, make sure that you are doing cargo run --release or cargo build --release to run your program.

Hope this helps—