Focus on sfb

1 files changed, 52 insertions, 0 deletions

diff --git a/src/corpus.rs b/src/corpus.rs
index 20ddbb3..2fab5a4 100644
--- a/src/corpus.rs
+++ b/src/corpus.rs
@@ -51,6 +51,58 @@ impl Corpus {
         });
     }
 
+    pub fn layout_fitness(&self, layout: &Layout) -> u32 {
+        let mut sfb: u32 = 0;
+
+        // 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 += 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 += 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 += 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)];
+            }
+        }
+        sfb
+    }
+
     pub fn evaluate_layout(&self, layout: &Layout) -> Evaluation {
         let mut keypress: [u32; NUM_KEYS] = [0; NUM_KEYS];
         let mut sfb: [u32; 8] = [0; 8];