[2016][rust][10.0] create bots that can take values

2016/rust/src/day_10.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use regex::Regex;
 
 use errors::*;
@@ -7,23 +9,41 @@ pub fn solve(input: &str) -> Result<String> {
 }
 
 struct Factory {
-  bots: Vec<usize>,
+  bots: HashMap<usize, Bot>,
 }
 
 impl Factory {
   fn parse(input: &str) -> Result<Self> {
-    let re = Regex::new(r"bot (\d+) gives low to (bot|output) (\d+) and high to (bot|output) (\d+)").unwrap();
+    let re =
+      Regex::new(r"bot (\d+) gives low to bot (\d+) and high to bot (\d+)")
+        .unwrap();
 
-    let mut bots = Vec::new();
+    let mut bots = HashMap::new();
     for line in input.lines() {
       if let Some(caps) = re.captures(line) {
-        bots.push(0);
+        let id = caps[1].parse().unwrap();
+        bots.insert(id,
+                    Bot {
+                      id: id,
+                      values: Vec::new(),
+                    });
       }
     }
     Ok(Factory { bots: bots })
   }
 }
 
+struct Bot {
+  id: usize,
+  values: Vec<usize>,
+}
+
+impl Bot {
+  fn push(&mut self, value: usize) {
+    self.values.push(value);
+  }
+}
+
 #[test]
 fn test_factory() {
   let input = "value 5 goes to bot 2
@@ -32,9 +52,13 @@ value 3 goes to bot 1
 bot 1 gives low to output 1 and high to bot 0
 bot 0 gives low to output 2 and high to output 0
 value 2 goes to bot 2";
-  let factory = Factory::parse(input).unwrap();
+  let mut factory = Factory::parse(input).unwrap();
+
+  assert_eq!(factory.bots.len(), 1);
 
-  assert_eq!(factory.bots.len(), 3);
+  let mut bot = factory.bots.get_mut(&2).unwrap();
+  bot.push(2);
+  assert_eq!(bot.values, vec![2]);
 
   // After running, outputs are 5, 2, 3.
 }