feat: support for word lengths

Mobydack · Jan 27, 2024 · 3fff446 · 3fff446
1 parent 7731a5e
commit 3fff446
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 15 deletions.
diff --git a/src/generator/strategies/weight_index.rs b/src/generator/strategies/weight_index.rs
@@ -11,11 +11,29 @@ impl WeightIndexStrategy {
     }
 }
 
+impl WeightIndexStrategy {
+
+}
+
 impl Strategy for WeightIndexStrategy {
     fn execute(&self, model: &Model) -> String {
         let mut rng = thread_rng();
         let mut word = String::new();
-        let mut curr_token = Token::StartPosition;
+
+        let mut available_lengths = model.transitions
+            .keys()
+            .filter(|&token| match token {
+                Token::StartPosition(_len) => true,
+                _ => false
+            })
+            .map(|token| match token {
+                Token::StartPosition(length) => length,
+                _ => &0,
+            });
+        let length_distributions = WeightedIndex::new(available_lengths.clone()).unwrap();
+
+        let word_length = available_lengths.nth(length_distributions.sample(&mut rng)).unwrap();
+        let mut curr_token = Token::StartPosition(*word_length);
 
         loop {
             match model.transitions.get(&curr_token) {
@@ -27,7 +45,7 @@ impl Strategy for WeightIndexStrategy {
                     if let Some((token, _)) = weights.nth(distributions.sample(&mut rng)) {
                         match token {
                             Token::EndPosition => {
-                                if word.len() >= 3 {
+                                if word.len() >= *word_length {
                                     return word;
                                 } else {
                                     continue;

diff --git a/src/main.rs b/src/main.rs
@@ -3,16 +3,14 @@ mod model_def;
 
 use generator::prelude::*;
 use model_def::prelude::*;
+use std::fs;
 
 fn main() {
-    let data_preset = vec![
-        "Aaren", "Aarika", "Abagael", "Abagail", "Abbe", "Abbey", "Abbi", "Abbie", "Abby", "Abra",
-    ]
-    .into_iter();
+    let file_content = fs::read_to_string("./dataset/words.txt").unwrap();
     let mut model = Model::new();
     let strategy = WeightIndexStrategy::new();
 
-    model.apply_preset(data_preset);
+    model.apply_preset(file_content.split('\n'));
 
     println!("Generated word: {:?}", strategy.execute(&model));
 }
diff --git a/src/model_def/model.rs b/src/model_def/model.rs
@@ -38,15 +38,15 @@ impl Model {
         let mut parent_token_option: Option<Token> = None;
 
         normalized_atomic_word.0.chars().for_each(|character| {
-            let token = Token::get_token(&character);
+            let token = Token::get_token(&normalized_atomic_word.0, &character);
 
             self.transitions
                 .entry(token)
                 .or_insert(Box::new(Transition::new()));
 
             if let Some(parent_token) = parent_token_option {
-                self.transitions.entry(parent_token).and_modify(|children| {
-                    children
+                self.transitions.entry(parent_token).and_modify(|transition| {
+                    transition
                         .0
                         .entry(token)
                         .and_modify(|token_count| *token_count += 1)

diff --git a/src/model_def/token.rs b/src/model_def/token.rs
@@ -3,15 +3,15 @@ pub const END_POSITION_TOKEN: char = '$';
 
 #[derive(Debug, PartialEq, Eq, Copy, Clone, Hash)]
 pub enum Token {
-    StartPosition,
+    StartPosition(usize),
     EndPosition,
     Symbol(char),
 }
 
 impl Token {
-    pub fn get_token(character: &char) -> Token {
+    pub fn get_token(word: &str, character: &char) -> Token {
         match *character {
-            START_POSITION_TOKEN => Token::StartPosition,
+            START_POSITION_TOKEN => Token::StartPosition(word.len()),
             END_POSITION_TOKEN => Token::EndPosition,
             _ => Token::Symbol(*character),
         }
@@ -25,7 +25,8 @@ mod tests {
 
     #[test]
     fn should_get_correct_token_for_different_characters() {
-        let mut samples = vec![('$', Token::EndPosition), ('^', Token::StartPosition)];
+        let word = "word";
+        let mut samples = vec![('$', Token::EndPosition), ('^', Token::StartPosition(word.len()))];
 
         Alphanumeric
             .sample_string(&mut rand::thread_rng(), 16)
@@ -35,7 +36,7 @@ mod tests {
             });
 
         samples.iter().for_each(|(character, expected_token)| {
-            assert_eq!(Token::get_token(character), *expected_token)
+            assert_eq!(Token::get_token(word, character), *expected_token)
         });
     }
 }