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 theHashMap
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.