1 files changed, 137 insertions, 0 deletions
diff --git a/src/corpus.rs b/src/corpus.rs
new file mode 100644
index 0000000..20ddbb3
--- /dev/null
+++ b/src/corpus.rs
@@ -0,0 +1,137 @@
+use std::fmt;
+use std::fs;
+
+use crate::config::*;
+use crate::layout::*;
+
+const NUM_BIGRAMS: usize = NUM_KEYS * NUM_KEYS;
+
+pub struct Corpus {
+    bigram_count: [u32; NUM_BIGRAMS],
+    character_count: [u32; NUM_KEYS],
+    total_count: 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 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; NUM_KEYS];
+        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[CHAR_TO_INDEX[c as usize]] += 1;
+                total_count += 1;
+            } else {
+                last_char = None;
+            }
+        }
+
+        return Ok(Corpus {
+            bigram_count,
+            character_count,
+            total_count,
+        });
+    }
+
+    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() {
+            keypress[layout.translate_index(i)] = 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]
+        ))
+    }
+}