Lightly refactor IntCode computer (day 5) puzzle

src/day_5/mod.rs

@@ -189,7 +189,12 @@ impl Create for Puzzle {
 
 impl Puzzle {
     fn solve_p1(&self) -> i32 {
-        let result = self.computer.run(stdin().lock(), stdout());
+        let result = if self.options.prompt {
+            self.computer.run(stdin().lock(), stdout())
+        } else {
+            self.computer.run(&b"1"[..], stdout())
+        };
+
         match result {
             Ok(res) => res[0],
             Err(e) => panic!(e),
@@ -197,7 +202,16 @@ impl Puzzle {
     }
 
     fn solve_p2(&self) -> i32 {
-        unimplemented!()
+        let result = if self.options.prompt {
+            self.computer.run(stdin().lock(), stdout())
+        } else {
+            self.computer.run(&b"5"[..], stdout())
+        };
+
+        match result {
+            Ok(res) => res[0],
+            Err(e) => panic!(e),
+        }
     }
 
     fn prompt_for_opcodes() -> Vec<i32> {
@@ -213,7 +227,7 @@ impl Puzzle {
             .split(',')
             .map(|code| match code.trim().parse::<i32>() {
                 Ok(i) => i,
-                Err(e) => panic!("failed to parse {}: {}", code, e),
+                Err(e) => panic!("Failed to parse {}: {}", code, e),
             })
             .collect()
     }
@@ -262,73 +276,36 @@ impl Puzzle {
     }
 }
 
-#[derive(Debug, Copy, Clone)]
-enum OpMode {
-    Position = 0,
-    Immediate,
+struct Computer {
+    opcodes: Vec<i32>,
 }
 
-impl TryFrom<i32> for OpMode {
-    type Error = i32;
+impl Computer {
+    fn run(&self, mut reader: impl BufRead, mut writer: impl Write) -> Result<Vec<i32>, String> {
+        let mut index: usize = 0;
+        let ref mut opcodes = self.opcodes.clone();
 
-    fn try_from(v: i32) -> Result<Self, Self::Error> {
-        trace!("try converting {} to enum OpMode", v);
+        while index < opcodes.len() {
+            let mut op = match Operation::new(opcodes, index) {
+                Ok(x) => x,
+                Err(e) => {
+                    return Err(format!(
+                        "failed to parse operation '{}': {}",
+                        opcodes[index], e
+                    ))
+                }
+            };
 
-        match v {
-            x if x == OpMode::Position as i32 => Ok(OpMode::Position),
-            x if x == OpMode::Immediate as i32 => Ok(OpMode::Immediate),
-            _ => Err(v),
+            match op.apply(index, &mut reader, &mut writer) {
+                Ok(Some(i)) => index = i,
+                Ok(None) => break,
+                Err(e) => return Err(e),
+            }
         }
-    }
-}
-
-#[derive(Debug)]
-enum OpCode {
-    Add = 1,
-    Mult,
-    Input,
-    Output,
-    JumpIfTrue,
-    JumpIfFalse,
-    LessThan,
-    Equals,
-    Halt = 99,
-}
-
-impl TryFrom<i32> for OpCode {
-    type Error = i32;
-
-    fn try_from(v: i32) -> Result<Self, Self::Error> {
-        trace!("try converting {} to enum OpCode", v);
 
-        match v {
-            x if x == OpCode::Add as i32 => Ok(OpCode::Add),
-            x if x == OpCode::Mult as i32 => Ok(OpCode::Mult),
-            x if x == OpCode::Input as i32 => Ok(OpCode::Input),
-            x if x == OpCode::Output as i32 => Ok(OpCode::Output),
-            x if x == OpCode::JumpIfTrue as i32 => Ok(OpCode::JumpIfTrue),
-            x if x == OpCode::JumpIfFalse as i32 => Ok(OpCode::JumpIfFalse),
-            x if x == OpCode::LessThan as i32 => Ok(OpCode::LessThan),
-            x if x == OpCode::Equals as i32 => Ok(OpCode::Equals),
-            x if x == OpCode::Halt as i32 => Ok(OpCode::Halt),
-            _ => Err(v),
-        }
-    }
-}
+        debug!("resulting opcodes: {:?}", opcodes);
 
-impl OpCode {
-    fn num_args(&self) -> i32 {
-        match self {
-            OpCode::Add => 3,
-            OpCode::Mult => 3,
-            OpCode::Input => 1,
-            OpCode::Output => 1,
-            OpCode::JumpIfTrue => 2,
-            OpCode::JumpIfFalse => 2,
-            OpCode::LessThan => 3,
-            OpCode::Equals => 3,
-            OpCode::Halt => 0,
-        }
+        Ok(opcodes.to_owned())
     }
 }
 
@@ -350,7 +327,7 @@ impl Operation<'_> {
 
         let operation = match OpCode::try_from(opcode % 100) {
             Ok(x) => x,
-            Err(e) => return Err(format!("Unknown opcode {}", e)),
+            Err(e) => return Err(format!("unknown opcode {}", e)),
         };
         opcode /= 100;
         trace!("operation: {:?}, opcode: {}", operation, opcode);
@@ -392,74 +369,85 @@ impl Operation<'_> {
     ) -> Result<Option<usize>, String> {
         debug!("applying {:?} at index {}", self, index);
 
+        let num_args = self.opcode.num_args();
+        let end_index = index + num_args as usize;
+
+        if end_index >= self.opcodes.len() {
+            return Err(String::from("can't operate outside opcode instruction set"));
+        }
+
+        let (arg1, arg2) = match num_args {
+            x if x >= 2 => (
+                self.value_from(self.param_1_mode, index + 1),
+                self.value_from(self.param_2_mode, index + 2),
+            ),
+            x if x == 1 && self.opcode == OpCode::Input => (0, 0),
+            x if x == 1 => (self.value_from(self.param_1_mode, index + 1), 0),
+            x if x == 0 => (0, 0),
+            x => return Err(format!("failed to get {} required args", x)),
+        };
+
         match &self.opcode {
             OpCode::Add => {
-                let arg1 = self.value_from(self.param_1_mode, index + 1);
-                let arg2 = self.value_from(self.param_2_mode, index + 2);
                 self.value_to(self.param_3_mode, index + 3, arg1 + arg2);
             }
             OpCode::Mult => {
-                let arg1 = self.value_from(self.param_1_mode, index + 1);
-                let arg2 = self.value_from(self.param_2_mode, index + 2);
                 self.value_to(self.param_3_mode, index + 3, arg1 * arg2);
             }
-            OpCode::Input => {
-                print!("Input: ");
-                writer.flush().unwrap();
-
-                let input = &mut String::new();
-                let _ = reader.read_line(input);
-                let num: i32 = match input.trim().parse() {
-                    Ok(x) => x,
-                    Err(e) => return Err(format!("failed to decode input {}: {}", input, e)),
-                };
-                self.value_to(self.param_1_mode, index + 1, num);
-            }
-            OpCode::Output => {
-                let arg1 = self.value_from(self.param_1_mode, index + 1);
-                print!("Output: ");
-                match write!(writer, "{}", arg1) {
-                    Ok(_) => {}
-                    Err(e) => return Err(format!("Error writing output {}: {}", arg1, e)),
-                }
-                print!("\n");
-            }
+            OpCode::Input => match self.get_input(reader, writer) {
+                Ok(x) => self.value_to(self.param_1_mode, index + 1, x),
+                Err(e) => return Err(e),
+            },
+            OpCode::Output => match self.put_output(arg1, writer) {
+                Ok(_) => {}
+                Err(e) => return Err(e),
+            },
             OpCode::JumpIfTrue => {
-                let arg1 = self.value_from(self.param_1_mode, index + 1);
-                let arg2 = self.value_from(self.param_2_mode, index + 2);
-
                 if arg1 != 0 {
                     return Ok(Some(arg2 as usize));
                 }
             }
             OpCode::JumpIfFalse => {
-                let arg1 = self.value_from(self.param_1_mode, index + 1);
-                let arg2 = self.value_from(self.param_2_mode, index + 2);
-
                 if arg1 == 0 {
                     return Ok(Some(arg2 as usize));
                 }
             }
             OpCode::LessThan => {
-                let arg1 = self.value_from(self.param_1_mode, index + 1);
-                let arg2 = self.value_from(self.param_2_mode, index + 2);
-
                 self.value_to(self.param_3_mode, index + 3, (arg1 < arg2) as i32);
             }
             OpCode::Equals => {
-                let arg1 = self.value_from(self.param_1_mode, index + 1);
-                let arg2 = self.value_from(self.param_2_mode, index + 2);
-
                 self.value_to(self.param_3_mode, index + 3, (arg1 == arg2) as i32);
             }
             OpCode::Halt => return Ok(None),
         }
 
-        Ok(Some(index + (self.opcode.num_args() as usize) + 1))
+        Ok(Some(end_index + 1))
+    }
+
+    fn put_output(&self, value: i32, writer: &mut impl Write) -> Result<(), String> {
+        print!("Output: ");
+        match write!(writer, "{}", value) {
+            Ok(_) => {}
+            Err(e) => return Err(format!("error writing output {}", e)),
+        }
+        print!("\n");
+        Ok(())
+    }
+
+    fn get_input(&self, reader: &mut impl BufRead, writer: &mut impl Write) -> Result<i32, String> {
+        print!("Input: ");
+        writer.flush().unwrap();
+
+        let input = &mut String::new();
+        let _ = reader.read_line(input);
+        match input.trim().parse() {
+            Ok(x) => return Ok(x),
+            Err(e) => return Err(format!("failed to decode input {}: {}", input, e)),
+        };
     }
 
     fn value_from(&mut self, opmode: OpMode, index: usize) -> i32 {
-        debug!("GET via OpMode {:?} from index {}", opmode, index);
+        debug!("get via OpMode {:?} from index {}", opmode, index);
 
         match opmode {
             OpMode::Position => {
@@ -476,7 +464,7 @@ impl Operation<'_> {
     }
 
     fn value_to(&mut self, opmode: OpMode, index: usize, value: i32) {
-        debug!("SET {} via OpMode {:?} to index {}", value, opmode, index);
+        debug!("set {} via OpMode {:?} to index {}", value, opmode, index);
 
         match opmode {
             OpMode::Position => {
@@ -502,36 +490,73 @@ impl Operation<'_> {
     }
 }
 
-struct Computer {
-    opcodes: Vec<i32>,
+#[derive(Debug, PartialEq)]
+enum OpCode {
+    Add = 1,
+    Mult,
+    Input,
+    Output,
+    JumpIfTrue,
+    JumpIfFalse,
+    LessThan,
+    Equals,
+    Halt = 99,
 }
 
-impl Computer {
-    fn run(&self, mut reader: impl BufRead, mut writer: impl Write) -> Result<Vec<i32>, String> {
-        let mut index: usize = 0;
-        let ref mut opcodes = self.opcodes.clone();
+impl TryFrom<i32> for OpCode {
+    type Error = i32;
 
-        while index < opcodes.len() {
-            let mut op = match Operation::new(opcodes, index) {
-                Ok(x) => x,
-                Err(e) => {
-                    return Err(format!(
-                        "failed to parse operation '{}': {}",
-                        opcodes[index], e
-                    ))
-                }
-            };
+    fn try_from(v: i32) -> Result<Self, Self::Error> {
+        trace!("try converting {} to enum OpCode", v);
 
-            match op.apply(index, &mut reader, &mut writer) {
-                Ok(Some(i)) => index = i,
-                Ok(None) => break,
-                Err(e) => return Err(e),
-            }
+        match v {
+            x if x == OpCode::Add as i32 => Ok(OpCode::Add),
+            x if x == OpCode::Mult as i32 => Ok(OpCode::Mult),
+            x if x == OpCode::Input as i32 => Ok(OpCode::Input),
+            x if x == OpCode::Output as i32 => Ok(OpCode::Output),
+            x if x == OpCode::JumpIfTrue as i32 => Ok(OpCode::JumpIfTrue),
+            x if x == OpCode::JumpIfFalse as i32 => Ok(OpCode::JumpIfFalse),
+            x if x == OpCode::LessThan as i32 => Ok(OpCode::LessThan),
+            x if x == OpCode::Equals as i32 => Ok(OpCode::Equals),
+            x if x == OpCode::Halt as i32 => Ok(OpCode::Halt),
+            _ => Err(v),
         }
+    }
+}
 
-        debug!("resulting opcodes: {:?}", opcodes);
+impl OpCode {
+    fn num_args(&self) -> i32 {
+        match self {
+            OpCode::Add => 3,
+            OpCode::Mult => 3,
+            OpCode::Input => 1,
+            OpCode::Output => 1,
+            OpCode::JumpIfTrue => 2,
+            OpCode::JumpIfFalse => 2,
+            OpCode::LessThan => 3,
+            OpCode::Equals => 3,
+            OpCode::Halt => 0,
+        }
+    }
+}
 
-        Ok(opcodes.to_owned())
+#[derive(Debug, Copy, Clone)]
+enum OpMode {
+    Position = 0,
+    Immediate,
+}
+
+impl TryFrom<i32> for OpMode {
+    type Error = i32;
+
+    fn try_from(v: i32) -> Result<Self, Self::Error> {
+        trace!("try converting {} to enum OpMode", v);
+
+        match v {
+            x if x == OpMode::Position as i32 => Ok(OpMode::Position),
+            x if x == OpMode::Immediate as i32 => Ok(OpMode::Immediate),
+            _ => Err(v),
+        }
     }
 }