1 files changed, 34 insertions, 8 deletions

diff --git a/src/dictionary.rs b/src/dictionary.rs
index e651e71..563621d 100644
--- a/src/dictionary.rs
+++ b/src/dictionary.rs
@@ -34,19 +34,29 @@ impl Dictionary {
         Self { last_index : i - 1, data }
     }
 
-    pub fn set(&mut self, name : String) {
-        if let std::collections::hash_map::Entry::Vacant(e) = self.data.entry(name) {
+    pub fn set(&mut self, name : &String) {
+        if !self.data.contains_key(name) {
             self.last_index += 1;
-            e.insert(self.last_index as u64);
+            self.data.insert(name.clone(), self.last_index as u64);
         }
     }
 
-    pub fn get(&self, name : String) -> Option<&u64> {
-        self.data.get(&name)
+    pub fn set_and_get(&mut self, name : &String) -> u64 {
+        if !self.data.contains_key(name) {
+            self.last_index += 1;
+            self.data.insert(name.clone(), self.last_index as u64);
+            self.last_index as u64
+        } else {
+            *self.data.get(name).unwrap()
+        }
     }
 
-    pub fn iter(&self) -> &HashMap<String, u64> {
-        &self.data
+    pub fn get(&self, name : &String) -> Option<&u64> {
+        self.data.get(name)
+    }
+
+    pub fn iter(&self) -> std::collections::hash_map::Iter<'_, String, u64> {
+        self.data.iter()
     }
 
     pub fn to_list(&self) -> Vec<String> {
@@ -61,7 +71,7 @@ impl Dictionary {
         v
     }
 
-    pub fn vectorize_word_list(&self, words : Vec<String>) -> FileVector {
+    pub fn vectorize_word_list(&self, words : &Vec<&String>) -> FileVector {
         let mut fv = FileVector::new();
 
         for word in words {
@@ -76,4 +86,20 @@ impl Dictionary {
 
         fv
     }
+
+    pub fn insert_words_and_vectorize_word_list(&mut self, words : &Vec<&String>) -> FileVector {
+        let mut fv = FileVector::new();
+
+        for word in words {
+            let i = self.set_and_get(word);
+            if !fv.contains_key(&i) {
+                fv.insert(i, 1);
+            } else {
+                let c : u64 = *fv.get(&i).unwrap();
+                fv.insert(i, c + 1);
+            }
+        }
+
+        fv
+    }
 }