1 files changed, 204 insertions, 0 deletions
diff --git a/srchr/src/corpus.rs b/srchr/src/corpus.rs
new file mode 100644
index 0000000..8e4b437
--- /dev/null
+++ b/srchr/src/corpus.rs
@@ -0,0 +1,204 @@
+use std::fmt;
+use std::fs;
+
+use crate::layout::*;
+
+const fn build_lookup_table() -> [usize; 128] {
+    let mut result = [0; 128];
+    let mut i = 0;
+
+    while i < NUM_KEYS {
+        result[KEY_CHARS[i] as usize] = i;
+        i += 1;
+    }
+    return result;
+}
+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],
+    character_count: [u32; 256],
+    total_count: u32,
+    index_threshold: u32,
+}
+
+fn pair_to_index(x: char, y: char) -> usize {
+    let xi = CHAR_TO_INDEX[x as usize];
+    let yi = CHAR_TO_INDEX[y as usize];
+    xi + NUM_KEYS * yi
+}
+
+impl Corpus {
+    pub fn get_character_count(&self, c: char) -> u32 {
+        self.character_count[c as usize]
+    }
+
+    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 character_count = [0; 256];
+        let mut total_count = 0;
+
+        let mut last_char = None;
+        for c in contents.chars() {
+            if let Some(c) = canonicalise(c) {
+                if let Some(lc) = last_char {
+                    // We don't count these anyway
+                    if lc != c {
+                        bigram_count[pair_to_index(c, lc)] += 1;
+                        bigram_count[pair_to_index(lc, c)] += 1;
+                    }
+                }
+                last_char = Some(c);
+                character_count[c as usize] += 1;
+                total_count += 1;
+            } else {
+                last_char = None;
+            }
+        }
+
+        let index_threshold = (total_count as f32 * 0.15) as u32;
+
+        return Ok(Corpus {
+            bigram_count,
+            character_count,
+            total_count,
+            index_threshold,
+        });
+    }
+
+    pub fn prelayout_fitness(&self, layout: &Prelayout) -> u32 {
+        let mut score: u32 = 0;
+
+        for i in 0..6 {
+            let keys = layout.get_standard_column(i);
+            let k1 = keys[0];
+            let k2 = keys[1];
+            let k3 = keys[2];
+
+            score += self.bigram_count[pair_to_index(k1, k2)]
+                + self.bigram_count[pair_to_index(k1, k3)]
+                + self.bigram_count[pair_to_index(k2, k3)];
+        }
+
+        for i in 0..2 {
+            let keys = layout.get_index_column(i);
+            let k1 = keys[0];
+            let k2 = keys[1];
+            let k3 = keys[2];
+            let k4 = keys[3];
+            let k5 = keys[4];
+            let k6 = keys[5];
+
+            // We want index usage!
+            let index_count: u32 = keys.iter().map(|&k| self.character_count[k as usize]).sum();
+            score += (self.index_threshold as i64 - index_count as i64).abs() as u32 / 64;
+
+            score += self.bigram_count[pair_to_index(k1, k2)]
+                + self.bigram_count[pair_to_index(k1, k3)]
+                + self.bigram_count[pair_to_index(k2, k3)]
+                + self.bigram_count[pair_to_index(k4, k5)]
+                + self.bigram_count[pair_to_index(k4, k6)]
+                + self.bigram_count[pair_to_index(k5, k6)]
+                + self.bigram_count[pair_to_index(k1, k4)]
+                + self.bigram_count[pair_to_index(k1, k5)]
+                + self.bigram_count[pair_to_index(k1, k6)]
+                + self.bigram_count[pair_to_index(k2, k4)]
+                + self.bigram_count[pair_to_index(k2, k5)]
+                + self.bigram_count[pair_to_index(k2, k6)]
+                + self.bigram_count[pair_to_index(k3, k4)]
+                + self.bigram_count[pair_to_index(k3, k5)]
+                + self.bigram_count[pair_to_index(k3, k6)];
+        }
+
+        score
+    }
+
+    pub fn evaluate_layout(&self, layout: &Layout) -> Evaluation {
+        let mut keypress: [u32; NUM_KEYS] = [0; NUM_KEYS];
+        let mut sfb: [u32; 8] = [0; 8];
+
+        for (i, &c) in self.character_count.iter().enumerate() {
+            if c > 0 {
+                keypress[layout.get_index((i as u8) as char)] = c;
+            }
+        }
+
+        // TODO: We make assumptions about NUM_KEYS here
+        for i in 0..8 {
+            let ind = if i < 4 { i } else { i + 2 };
+
+            let k1 = layout.get_key(ind + 10 * 0);
+            let k2 = layout.get_key(ind + 10 * 1);
+            let k3 = layout.get_key(ind + 10 * 2);
+
+            sfb[i] = self.bigram_count[pair_to_index(k1, k2)]
+                + self.bigram_count[pair_to_index(k1, k3)]
+                + self.bigram_count[pair_to_index(k2, k3)];
+
+            if i == 3 {
+                let k4 = layout.get_key(4 + 10 * 0);
+                let k5 = layout.get_key(4 + 10 * 1);
+                let k6 = layout.get_key(4 + 10 * 2);
+                sfb[i] += self.bigram_count[pair_to_index(k4, k5)]
+                    + self.bigram_count[pair_to_index(k4, k6)]
+                    + self.bigram_count[pair_to_index(k5, k6)]
+                    + self.bigram_count[pair_to_index(k1, k4)]
+                    + self.bigram_count[pair_to_index(k1, k5)]
+                    + self.bigram_count[pair_to_index(k1, k6)]
+                    + self.bigram_count[pair_to_index(k2, k4)]
+                    + self.bigram_count[pair_to_index(k2, k5)]
+                    + self.bigram_count[pair_to_index(k2, k6)]
+                    + self.bigram_count[pair_to_index(k3, k4)]
+                    + self.bigram_count[pair_to_index(k3, k5)]
+                    + self.bigram_count[pair_to_index(k3, k6)];
+            } else if i == 4 {
+                let k4 = layout.get_key(5 + 10 * 0);
+                let k5 = layout.get_key(5 + 10 * 1);
+                let k6 = layout.get_key(5 + 10 * 2);
+                sfb[i] += self.bigram_count[pair_to_index(k4, k5)]
+                    + self.bigram_count[pair_to_index(k5, k6)]
+                    + self.bigram_count[pair_to_index(k4, k6)]
+                    + self.bigram_count[pair_to_index(k1, k4)]
+                    + self.bigram_count[pair_to_index(k1, k5)]
+                    + self.bigram_count[pair_to_index(k1, k6)]
+                    + self.bigram_count[pair_to_index(k2, k4)]
+                    + self.bigram_count[pair_to_index(k2, k5)]
+                    + self.bigram_count[pair_to_index(k2, k6)]
+                    + self.bigram_count[pair_to_index(k3, k4)]
+                    + self.bigram_count[pair_to_index(k3, k5)]
+                    + self.bigram_count[pair_to_index(k3, k6)];
+            }
+        }
+
+        return Evaluation::new(keypress, self.total_count, sfb);
+    }
+}
+
+fn dump_bigrams(corpus: &Corpus) -> Vec<(String, u32)> {
+    let mut result: Vec<(String, u32)> = 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, y)]));
+        }
+    }
+
+    result.sort_by(|(_, c1), (_, c2)| c1.cmp(c2).reverse());
+
+    return result;
+}
+
+impl fmt::Display for Corpus {
+    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+        formatter.write_str(&format!(
+            "Counted {} characters, top 5 bigrams {:?}",
+            self.total_count,
+            &dump_bigrams(self)[0..5]
+        ))
+    }
+}