From 8e74c1813580fbcf39c5ecd93997baa127df44fa Mon Sep 17 00:00:00 2001
From: ryan <ryan@distrust.co>
Date: Fri, 18 Aug 2023 01:11:54 -0500
Subject: [PATCH] keyfork-mnemonic-generate: split wordlist parsing into its
 own thing

---
 keyfork-mnemonic-generate/src/main.rs | 71 +++++++++++++++++----------
 1 file changed, 44 insertions(+), 27 deletions(-)

diff --git a/keyfork-mnemonic-generate/src/main.rs b/keyfork-mnemonic-generate/src/main.rs
index 78f93a2..032bc8f 100644
--- a/keyfork-mnemonic-generate/src/main.rs
+++ b/keyfork-mnemonic-generate/src/main.rs
@@ -91,17 +91,40 @@ fn ensure_offline() {
     }
 }
 
-// TODO: Can a Mnemonic be formatted using a wordlist, without allocating or without storing the
-// entire word list?
-//
-// NOTE: Yes, use a Language and pass a reference to this.
-struct Mnemonic {
-    pub words: Vec<usize>,
-    wordlist: Vec<String>,
+struct Wordlist(Vec<String>);
+
+impl Default for Wordlist {
+    fn default() -> Self {
+        // TODO: English is the only supported language.
+        let wordlist_file = include_str!("wordlist.txt");
+        Wordlist(
+            wordlist_file
+                .lines()
+                .skip(1)
+                .map(|x| x.trim().to_string())
+                .collect(),
+        )
+    }
 }
 
-impl Mnemonic {
-    fn from_entropy(bytes: &[u8]) -> Result<Mnemonic> {
+impl Wordlist {
+    fn get_word(&self, word: usize) -> Option<String> {
+        self.0.get(word).cloned()
+    }
+
+    #[cfg(test)]
+    fn into_inner(self) -> Vec<String> {
+        self.0
+    }
+}
+
+struct Mnemonic<'a> {
+    pub words: Vec<usize>,
+    wordlist: &'a Wordlist,
+}
+
+impl<'a> Mnemonic<'a> {
+    fn from_entropy(bytes: &[u8], wordlist: &'a Wordlist) -> Result<Mnemonic<'a>> {
         let bit_count = bytes.len() * 8;
         let hash = generate_slice_hash(bytes);
 
@@ -132,16 +155,15 @@ impl Mnemonic {
             .map(|chunk| bitslice_to_usize(chunk.try_into().expect("11 bit chunks")))
             .collect::<Vec<_>>();
 
-        let wordlist = build_wordlist();
         Ok(Mnemonic { words, wordlist })
     }
 }
 
-impl Display for Mnemonic {
+impl<'a> Display for Mnemonic<'a> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let mut iter = self.words.iter().peekable();
         while let Some(word_index) = iter.next() {
-            let word = &self.wordlist[*word_index];
+            let word = self.wordlist.get_word(*word_index).expect("word");
             write!(f, "{word}")?;
             if iter.peek().is_some() {
                 write!(f, " ")?;
@@ -157,15 +179,6 @@ fn generate_slice_hash(data: &[u8]) -> Vec<u8> {
     hasher.finalize().to_vec()
 }
 
-fn build_wordlist() -> Vec<String> {
-    let wordlist_file = include_str!("wordlist.txt");
-    wordlist_file
-        .lines()
-        .skip(1)
-        .map(|x| x.trim().to_string())
-        .collect()
-}
-
 fn bitslice_to_usize(bitslice: [bool; 11]) -> usize {
     let mut index = 0usize;
     for i in 0..11 {
@@ -196,7 +209,8 @@ fn main() -> Result<()> {
     let entropy = &mut [0u8; 256 / 8];
     random_handle.read_exact(&mut entropy[..])?;
 
-    let mnemonic = Mnemonic::from_entropy(&entropy[..bit_size / 8])?;
+    let wordlist = Wordlist::default();
+    let mnemonic = Mnemonic::from_entropy(&entropy[..bit_size / 8], &wordlist)?;
 
     println!("{mnemonic}");
 
@@ -209,7 +223,7 @@ mod tests {
 
     #[test]
     fn has_good_wordlist() {
-        let wordlist = build_wordlist();
+        let wordlist = Wordlist::default().into_inner();
         assert_eq!(
             wordlist.len(),
             2usize.pow(11),
@@ -228,7 +242,8 @@ mod tests {
             };
             let hex = hex::decode(hex_.as_str().unwrap()).unwrap();
 
-            let mnemonic = Mnemonic::from_entropy(&hex)?;
+            let wordlist = Wordlist::default();
+            let mnemonic = Mnemonic::from_entropy(&hex, &wordlist)?;
 
             assert_eq!(mnemonic.to_string(), seed.as_str().unwrap());
         }
@@ -240,7 +255,8 @@ mod tests {
         let mut random_handle = File::open("/dev/random").unwrap();
         let entropy = &mut [0u8; 256 / 8];
         random_handle.read_exact(&mut entropy[..]).unwrap();
-        let my_mnemonic = super::Mnemonic::from_entropy(&entropy[..256 / 8]).unwrap();
+        let wordlist = Wordlist::default();
+        let my_mnemonic = super::Mnemonic::from_entropy(&entropy[..256 / 8], &wordlist).unwrap();
         let their_mnemonic = bip39::Mnemonic::from_entropy(&entropy[..256 / 8]).unwrap();
         /*
         let my_words = my_mnemonic.words.clone();
@@ -252,13 +268,14 @@ mod tests {
     #[test]
     fn count_to_get_duplicate_words() {
         let mut count = 0.;
-        let tests = 10_000;
+        let tests = 100_000;
         let mut random_handle = File::open("/dev/random").unwrap();
         let entropy = &mut [0u8; 256 / 8];
+        let wordlist = Wordlist::default();
         for _ in 0..tests {
             random_handle.read_exact(&mut entropy[..]).unwrap();
             // let bits = generate_slice_hash(entropy);
-            let mnemonic = Mnemonic::from_entropy(&entropy[..256 / 8]).unwrap();
+            let mnemonic = Mnemonic::from_entropy(&entropy[..256 / 8], &wordlist).unwrap();
             let mut words = mnemonic.words.clone();
             words.sort();
             assert_eq!(words.len(), 24);