1 files changed, 47 insertions, 18 deletions

diff --git a/srchr/src/corpus.rs b/srchr/src/corpus.rs
index aa6e163..4428bc6 100644
--- a/srchr/src/corpus.rs
+++ b/srchr/src/corpus.rs
@@ -17,10 +17,11 @@ const CHAR_TO_INDEX: [usize; 128] = build_lookup_table();
 
 const NUM_BIGRAMS: usize = NUM_KEYS * NUM_KEYS;
 pub struct Corpus {
-    bigram_count: [u32; NUM_BIGRAMS],
+    bigram_percs: [f32; NUM_BIGRAMS],
+    skipgram_percs: [f32; NUM_BIGRAMS],
     character_count: [u32; 128],
     total_count: u32,
-    index_threshold: u32,
+    total_bigrams: u32,
 }
 
 fn pair_to_index(x: u8, y: u8) -> usize {
@@ -34,66 +35,94 @@ impl Corpus {
         self.character_count[c as usize]
     }
 
-    pub fn get_bigram_count(&self, x: u8, y: u8) -> u32 {
-        self.bigram_count[pair_to_index(x, y)]
+    pub fn get_bigram_perc(&self, x: u8, y: u8) -> f32 {
+        self.bigram_percs[pair_to_index(x, y)]
     }
 
-    pub fn get_index_threshold(&self) -> u32 {
-        self.index_threshold
+    pub fn get_skipgram_perc(&self, x: u8, y: u8) -> f32 {
+        self.skipgram_percs[pair_to_index(x, y)]
     }
 
     pub fn get_total_count(&self) -> u32 {
         self.total_count
     }
 
+    pub fn get_total_bigrams(&self) -> u32 {
+        self.total_bigrams
+    }
+
     pub fn load(path: &str) -> Result<Corpus, std::io::Error> {
         let contents = fs::read_to_string(path)?;
 
         let mut bigram_count = [0; NUM_BIGRAMS];
+        let mut skipgram_count = [0; NUM_BIGRAMS];
         let mut character_count = [0; 128];
         let mut total_count = 0;
+        let mut total_bigrams = 0;
+        let mut total_trigrams = 0;
 
         let mut last_char: Option<u8> = None;
+        let mut last_last_char: Option<u8> = None;
         for c in contents.chars() {
-            if let Some(c) = canonicalise(c) {
+            let current_char = canonicalise(c);
+            if let Some(c) = current_char {
                 if let Some(lc) = last_char {
-                    // We don't count these anyway
+                    if let Some(llc) = last_last_char {
+                        if llc != c {
+                            skipgram_count[pair_to_index(llc, c)] += 1;
+                            skipgram_count[pair_to_index(c, llc)] += 1;
+                        }
+                        total_trigrams += 1;
+                    }
                     if lc != c {
                         bigram_count[pair_to_index(c, lc)] += 1;
                         bigram_count[pair_to_index(lc, c)] += 1;
                     }
+                    total_bigrams += 1;
                 }
-                last_char = Some(c);
                 character_count[c as usize] += 1;
                 total_count += 1;
-            } else {
-                last_char = None;
             }
+            last_last_char = last_char;
+            last_char = current_char;
         }
 
-        let index_threshold = (total_count as f32 * DESIRED_INDEX_USAGE_PERCENT) as u32;
+        let mut bigram_percs = [0.0; NUM_BIGRAMS];
+        let mut skipgram_percs = [0.0; NUM_BIGRAMS];
+
+        for i in 0..NUM_BIGRAMS {
+            bigram_percs[i] = bigram_count[i] as f32 / total_bigrams as f32;
+            skipgram_percs[i] = skipgram_count[i] as f32 / total_trigrams as f32;
+        }
 
         return Ok(Corpus {
-            bigram_count,
+            skipgram_percs,
+            bigram_percs,
             character_count,
             total_count,
-            index_threshold,
+            total_bigrams,
         });
     }
 }
 
-fn dump_bigrams(corpus: &Corpus) -> Vec<(String, u32)> {
-    let mut result: Vec<(String, u32)> = Vec::new();
+fn dump_bigrams(corpus: &Corpus) -> Vec<(String, f32)> {
+    let mut result: Vec<(String, f32)> = Vec::new();
 
     for (i, &x) in KEY_CHARS.iter().enumerate() {
         for &y in &KEY_CHARS[i..] {
             let mut pair = String::from(x);
             pair.push(y);
-            result.push((pair, corpus.bigram_count[pair_to_index(x as u8, y as u8)]));
+            result.push((pair, corpus.bigram_percs[pair_to_index(x as u8, y as u8)]));
         }
     }
 
-    result.sort_by(|(_, c1), (_, c2)| c1.cmp(c2).reverse());
+    result.sort_by(|(_, c1), (_, c2)| {
+        if c2 < c1 {
+            std::cmp::Ordering::Less
+        } else {
+            std::cmp::Ordering::Greater
+        }
+    });
 
     return result;
 }