Change: implement += for WordCounts

1 files changed, 26 insertions, 10 deletions

diff --git a/src/main.rs b/src/main.rs
index f376742..4045cd7 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,6 +2,7 @@ use std::env;
 use std::fs::File;
 use std::io;
 use std::io::BufRead;
+use std::ops;
 use std::collections::HashMap;
 
 
@@ -10,38 +11,40 @@ fn main() -> io::Result<()> {
     for filename in env::args().skip(1) {
         let mut f = File::open(&filename)?;
         let mut reader = io::BufReader::new(f);
-        print_words(&mut reader)?;
+        let counts = count_words(&mut reader)?;
+        counts.print()
     }
     Ok(())
 }
 
 
-fn print_words(reader: &mut io::BufReader<File>) -> io::Result<()> {
+fn count_words(reader: &mut io::BufReader<File>) -> io::Result<WordCounts> {
+    let mut counts = WordCounts::new();
     for line in reader.lines() {
         let line = line?;
-        count_words_in_line(&line);
+        counts += count_words_in_line(&line);
     }
-    Ok(())
+    Ok(counts)
 }
 
 
-fn count_words_in_line(line: &str) {
-    let mut words = WordCounts::new();
+fn count_words_in_line(line: &str) -> WordCounts {
+    let mut counts = WordCounts::new();
     let mut word = String::new();
     for c in line.chars() {
         if c.is_alphabetic() {
             word.push(c);
         } else {
             if !word.is_empty() {
-                words.count(word.clone());
+                counts.count(word.clone());
                 word.clear();
             }
         }
     }
     if !word.is_empty() {
-        words.count(word.clone());
+        counts.count(word.clone());
     }
-    words.print();
+    counts
 }
 
 
@@ -58,8 +61,12 @@ impl WordCounts {
     }
 
     fn count(&mut self, word: String) {
+        self.add(word, 1);
+    }
+
+    fn add(&mut self, word: String, count: u32) {
         let counter = self.counts.entry(word).or_insert(0);
-        *counter += 1;
+        *counter += count;
     }
 
     fn print(&self) {
@@ -68,3 +75,12 @@ impl WordCounts {
         }
     }
 }
+
+
+impl ops::AddAssign for WordCounts {
+    fn add_assign(&mut self, other: WordCounts) {
+        for (word, count) in other.counts.iter() {
+            self.add(word.to_string(), *count);
+        }
+    }
+}