Solve day 5 puzzle 2

src/day_5/mod.rs

@@ -87,6 +87,67 @@ zero except the diagnostic code, the diagnostic program ran successfully.
 
 After providing 1 to the only input instruction and passing all the tests, what diagnostic code
 does the program produce?
+
+--- Part Two ---
+
+The air conditioner comes online! Its cold air feels good for a while, but then the TEST alarms
+start to go off. Since the air conditioner can't vent its heat anywhere but back into the
+spacecraft, it's actually making the air inside the ship warmer.
+
+Instead, you'll need to use the TEST to extend the thermal radiators. Fortunately, the diagnostic
+program (your puzzle input) is already equipped for this. Unfortunately, your Intcode computer is
+not.
+
+Your computer is only missing a few opcodes:
+
+    Opcode 5 is jump-if-true: if the first parameter is non-zero, it sets the instruction pointer
+    to the value from the second parameter. Otherwise, it does nothing.
+
+    Opcode 6 is jump-if-false: if the first parameter is zero, it sets the instruction pointer to
+    the value from the second parameter. Otherwise, it does nothing.
+
+    Opcode 7 is less than: if the first parameter is less than the second parameter, it stores 1 in
+    the position given by the third parameter. Otherwise, it stores 0.
+
+    Opcode 8 is equals: if the first parameter is equal to the second parameter, it stores 1 in the
+    position given by the third parameter. Otherwise, it stores 0.
+
+Like all instructions, these instructions need to support parameter modes as described above.
+
+Normally, after an instruction is finished, the instruction pointer increases by the number of
+values in that instruction. However, if the instruction modifies the instruction pointer, that
+value is used and the instruction pointer is not automatically increased.
+
+For example, here are several programs that take one input, compare it to the value 8, and then
+produce one output:
+
+    3,9,8,9,10,9,4,9,99,-1,8 - Using position mode, consider whether the input is equal to 8; output 1 (if it is) or 0 (if it is not).
+
+    3,9,7,9,10,9,4,9,99,-1,8 - Using position mode, consider whether the input is less than 8; output 1 (if it is) or 0 (if it is not).
+
+    3,3,1108,-1,8,3,4,3,99 - Using immediate mode, consider whether the input is equal to 8; output 1 (if it is) or 0 (if it is not).
+
+    3,3,1107,-1,8,3,4,3,99 - Using immediate mode, consider whether the input is less than 8; output 1 (if it is) or 0 (if it is not).
+
+Here are some jump tests that take an input, then output 0 if the input was zero or 1 if the input
+was non-zero:
+
+    3,12,6,12,15,1,13,14,13,4,13,99,-1,0,1,9 (using position mode)
+    3,3,1105,-1,9,1101,0,0,12,4,12,99,1 (using immediate mode)
+
+Here's a larger example:
+
+3,21,1008,21,8,20,1005,20,22,107,8,21,20,1006,20,31,1106,0,36,98,0,0,1002,21,125,20,4,20,1105,1,46,104,999,1105,1,46,1101,1000,1,20,4,20,1105,1,46,98,99
+
+The above example program uses an input instruction to ask for a single number. The program will
+then output 999 if the input value is below 8, output 1000 if the input value is equal to 8, or
+output 1001 if the input value is greater than 8.
+
+This time, when the TEST diagnostic program runs its input instruction to get the ID of the system
+to test, provide it 5, the ID for the ship's thermal radiator controller. This diagnostic test
+suite only outputs one number, the diagnostic code.
+
+What is the diagnostic code for system ID 5?
 */
 
 use log::{debug, trace};
@@ -227,6 +288,10 @@ enum OpCode {
     Mult,
     Input,
     Output,
+    JumpIfTrue,
+    JumpIfFalse,
+    LessThan,
+    Equals,
     Halt = 99,
 }
 
@@ -241,6 +306,10 @@ impl TryFrom<i32> for OpCode {
             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),
         }
@@ -250,10 +319,14 @@ impl TryFrom<i32> for OpCode {
 impl OpCode {
     fn num_args(&self) -> i32 {
         match self {
-            OpCode::Add => 4,
-            OpCode::Mult => 4,
-            OpCode::Input => 2,
-            OpCode::Output => 2,
+            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,
         }
     }
@@ -344,15 +417,45 @@ impl Operation<'_> {
             }
             OpCode::Output => {
                 let arg1 = self.value_from(self.param_1_mode, index + 1);
-                match write!(writer, "Output: {}", arg1) {
+                print!("Output: ");
+                match write!(writer, "{}", arg1) {
                     Ok(_) => {}
                     Err(e) => return Err(format!("Error writing output {}: {}", arg1, e)),
                 }
+                print!("\n");
+            }
+            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(self.opcode.num_args() as usize))
+        Ok(Some(index + (self.opcode.num_args() as usize) + 1))
     }
 
     fn value_from(&mut self, opmode: OpMode, index: usize) -> i32 {
@@ -420,7 +523,7 @@ impl Computer {
             };
 
             match op.apply(index, &mut reader, &mut writer) {
-                Ok(Some(i)) => index += i,
+                Ok(Some(i)) => index = i,
                 Ok(None) => break,
                 Err(e) => return Err(e),
             }
@@ -438,14 +541,11 @@ mod puzzle_1 {
 
     #[test]
     fn opcodes_sample_1() {
-        let buf_reader = b"";
-        let mut writer = Vec::new();
-
         assert_eq!(
             Computer {
                 opcodes: vec![1002, 4, 3, 4, 33],
             }
-            .run(&buf_reader[..], &mut writer)
+            .run(&[][..], &mut Vec::new())
             .unwrap(),
             vec![1002, 4, 3, 4, 99]
         );
@@ -453,14 +553,11 @@ mod puzzle_1 {
 
     #[test]
     fn opcodes_sample_2() {
-        let buf_reader = b"";
-        let mut writer = Vec::new();
-
         assert_eq!(
             Computer {
                 opcodes: vec![1, 9, 10, 3, 2, 3, 11, 0, 99, 30, 40, 50]
             }
-            .run(&buf_reader[..], &mut writer)
+            .run(&[][..], &mut Vec::new())
             .unwrap(),
             vec![3500, 9, 10, 70, 2, 3, 11, 0, 99, 30, 40, 50]
         );
@@ -468,14 +565,11 @@ mod puzzle_1 {
 
     #[test]
     fn opcodes_sample_3() {
-        let buf_reader = b"";
-        let mut writer = Vec::new();
-
         assert_eq!(
             Computer {
                 opcodes: vec![10001, 1, 1, 0, 99]
             }
-            .run(&buf_reader[..], &mut writer)
+            .run(&[][..], &mut Vec::new())
             .unwrap(),
             vec![10001, 1, 1, 2, 99]
         );
@@ -483,14 +577,11 @@ mod puzzle_1 {
 
     #[test]
     fn opcodes_sample_4() {
-        let buf_reader = b"";
-        let mut writer = Vec::new();
-
         assert_eq!(
             Computer {
                 opcodes: vec![1, 0, 0, 0, 99]
             }
-            .run(&buf_reader[..], &mut writer)
+            .run(&[][..], &mut Vec::new())
             .unwrap(),
             vec![2, 0, 0, 0, 99]
         );
@@ -510,7 +601,270 @@ mod puzzle_1 {
             vec![1234, 0, 4, 0, 99]
         );
 
-        let expected = b"Output: 1234";
+        let expected = b"1234";
+        assert_eq!(expected.to_vec(), writer)
+    }
+}
+
+#[cfg(test)]
+mod puzzle_2 {
+    use super::Computer;
+
+    #[test]
+    fn positional_eq_8_false() {
+        let buf_reader = b"3";
+        let mut writer = Vec::new();
+
+        assert_eq!(
+            Computer {
+                opcodes: vec![3, 9, 8, 9, 10, 9, 4, 9, 99, -1, 8]
+            }
+            .run(&buf_reader[..], &mut writer)
+            .unwrap(),
+            vec![3, 9, 8, 9, 10, 9, 4, 9, 99, 0, 8]
+        );
+
+        let expected = b"0";
+        assert_eq!(expected.to_vec(), writer)
+    }
+
+    #[test]
+    fn positional_eq_8_true() {
+        let buf_reader = b"8";
+        let mut writer = Vec::new();
+
+        assert_eq!(
+            Computer {
+                opcodes: vec![3, 9, 8, 9, 10, 9, 4, 9, 99, -1, 8]
+            }
+            .run(&buf_reader[..], &mut writer)
+            .unwrap(),
+            vec![3, 9, 8, 9, 10, 9, 4, 9, 99, 1, 8]
+        );
+
+        let expected = b"1";
+        assert_eq!(expected.to_vec(), writer)
+    }
+
+    #[test]
+    fn positional_lt_8_true() {
+        let buf_reader = b"3";
+        let mut writer = Vec::new();
+
+        assert_eq!(
+            Computer {
+                opcodes: vec![3, 9, 7, 9, 10, 9, 4, 9, 99, -1, 8]
+            }
+            .run(&buf_reader[..], &mut writer)
+            .unwrap(),
+            vec![3, 9, 7, 9, 10, 9, 4, 9, 99, 1, 8]
+        );
+
+        let expected = b"1";
+        assert_eq!(expected.to_vec(), writer)
+    }
+
+    #[test]
+    fn positional_lt_8_false() {
+        let buf_reader = b"8";
+        let mut writer = Vec::new();
+
+        assert_eq!(
+            Computer {
+                opcodes: vec![3, 9, 7, 9, 10, 9, 4, 9, 99, -1, 8]
+            }
+            .run(&buf_reader[..], &mut writer)
+            .unwrap(),
+            vec![3, 9, 7, 9, 10, 9, 4, 9, 99, 0, 8]
+        );
+
+        let expected = b"0";
+        assert_eq!(expected.to_vec(), writer)
+    }
+
+    #[test]
+    fn immediate_eq_8_false() {
+        let buf_reader = b"3";
+        let mut writer = Vec::new();
+
+        assert_eq!(
+            Computer {
+                opcodes: vec![3, 3, 1108, -1, 8, 3, 4, 3, 99]
+            }
+            .run(&buf_reader[..], &mut writer)
+            .unwrap(),
+            vec![3, 3, 1108, 0, 8, 3, 4, 3, 99]
+        );
+
+        let expected = b"0";
+        assert_eq!(expected.to_vec(), writer)
+    }
+
+    #[test]
+    fn immediate_eq_8_true() {
+        let buf_reader = b"8";
+        let mut writer = Vec::new();
+
+        assert_eq!(
+            Computer {
+                opcodes: vec![3, 3, 1108, -1, 8, 3, 4, 3, 99]
+            }
+            .run(&buf_reader[..], &mut writer)
+            .unwrap(),
+            vec![3, 3, 1108, 1, 8, 3, 4, 3, 99]
+        );
+
+        let expected = b"1";
+        assert_eq!(expected.to_vec(), writer)
+    }
+
+    #[test]
+    fn immediate_lt_8_false() {
+        let buf_reader = b"8";
+        let mut writer = Vec::new();
+
+        assert_eq!(
+            Computer {
+                opcodes: vec![3, 3, 1107, -1, 8, 3, 4, 3, 99]
+            }
+            .run(&buf_reader[..], &mut writer)
+            .unwrap(),
+            vec![3, 3, 1107, 0, 8, 3, 4, 3, 99]
+        );
+
+        let expected = b"0";
+        assert_eq!(expected.to_vec(), writer)
+    }
+
+    #[test]
+    fn immediate_lt_8_true() {
+        let buf_reader = b"3";
+        let mut writer = Vec::new();
+
+        assert_eq!(
+            Computer {
+                opcodes: vec![3, 3, 1107, -1, 8, 3, 4, 3, 99]
+            }
+            .run(&buf_reader[..], &mut writer)
+            .unwrap(),
+            vec![3, 3, 1107, 1, 8, 3, 4, 3, 99]
+        );
+
+        let expected = b"1";
+        assert_eq!(expected.to_vec(), writer)
+    }
+
+    #[test]
+    fn jump_print_0_if_0_true() {
+        let buf_reader = b"0";
+        let mut writer = Vec::new();
+
+        assert_eq!(
+            Computer {
+                opcodes: vec![3, 12, 6, 12, 15, 1, 13, 14, 13, 4, 13, 99, -1, 0, 1, 9]
+            }
+            .run(&buf_reader[..], &mut writer)
+            .unwrap(),
+            vec![3, 12, 6, 12, 15, 1, 13, 14, 13, 4, 13, 99, 0, 0, 1, 9]
+        );
+
+        let expected = b"0";
+        assert_eq!(expected.to_vec(), writer)
+    }
+
+    #[test]
+    fn jump_print_0_if_0_false() {
+        let buf_reader = b"21341234";
+        let mut writer = Vec::new();
+
+        assert_eq!(
+            Computer {
+                opcodes: vec![3, 12, 6, 12, 15, 1, 13, 14, 13, 4, 13, 99, -1, 0, 1, 9]
+            }
+            .run(&buf_reader[..], &mut writer)
+            .unwrap(),
+            vec![3, 12, 6, 12, 15, 1, 13, 14, 13, 4, 13, 99, 21341234, 1, 1, 9]
+        );
+
+        let expected = b"1";
+        assert_eq!(expected.to_vec(), writer)
+    }
+
+    #[test]
+    fn compared_to_8_lt() {
+        let buf_reader = b"1";
+        let mut writer = Vec::new();
+
+        assert_eq!(
+            Computer {
+                opcodes: vec![
+                    3, 21, 1008, 21, 8, 20, 1005, 20, 22, 107, 8, 21, 20, 1006, 20, 31, 1106, 0,
+                    36, 98, 0, 0, 1002, 21, 125, 20, 4, 20, 1105, 1, 46, 104, 999, 1105, 1, 46,
+                    1101, 1000, 1, 20, 4, 20, 1105, 1, 46, 98, 99
+                ]
+            }
+            .run(&buf_reader[..], &mut writer)
+            .unwrap(),
+            vec![
+                3, 21, 1008, 21, 8, 20, 1005, 20, 22, 107, 8, 21, 20, 1006, 20, 31, 1106, 0, 36,
+                98, 0, 1, 1002, 21, 125, 20, 4, 20, 1105, 1, 46, 104, 999, 1105, 1, 46, 1101, 1000,
+                1, 20, 4, 20, 1105, 1, 46, 98, 99
+            ]
+        );
+
+        let expected = b"999";
+        assert_eq!(expected.to_vec(), writer)
+    }
+
+    #[test]
+    fn compared_to_8_eq() {
+        let buf_reader = b"8";
+        let mut writer = Vec::new();
+
+        assert_eq!(
+            Computer {
+                opcodes: vec![
+                    3, 21, 1008, 21, 8, 20, 1005, 20, 22, 107, 8, 21, 20, 1006, 20, 31, 1106, 0,
+                    36, 98, 0, 0, 1002, 21, 125, 20, 4, 20, 1105, 1, 46, 104, 999, 1105, 1, 46,
+                    1101, 1000, 1, 20, 4, 20, 1105, 1, 46, 98, 99
+                ]
+            }
+            .run(&buf_reader[..], &mut writer)
+            .unwrap(),
+            vec![
+                3, 21, 1008, 21, 8, 20, 1005, 20, 22, 107, 8, 21, 20, 1006, 20, 31, 1106, 0, 36,
+                98, 1000, 8, 1002, 21, 125, 20, 4, 20, 1105, 1, 46, 104, 999, 1105, 1, 46, 1101,
+                1000, 1, 20, 4, 20, 1105, 1, 46, 98, 99
+            ]
+        );
+
+        let expected = b"1000";
+        assert_eq!(expected.to_vec(), writer)
+    }
+
+    #[test]
+    fn compared_to_8_gt() {
+        let buf_reader = b"9";
+        let mut writer = Vec::new();
+
+        assert_eq!(
+            Computer {
+                opcodes: vec![
+                    3, 21, 1008, 21, 8, 20, 1005, 20, 22, 107, 8, 21, 20, 1006, 20, 31, 1106, 0,
+                    36, 98, 0, 0, 1002, 21, 125, 20, 4, 20, 1105, 1, 46, 104, 999, 1105, 1, 46,
+                    1101, 1000, 1, 20, 4, 20, 1105, 1, 46, 98, 99
+                ]
+            }
+            .run(&buf_reader[..], &mut writer)
+            .unwrap(),
+            vec![
+                3, 21, 1008, 21, 8, 20, 1005, 20, 22, 107, 8, 21, 20, 1006, 20, 31, 1106, 0, 36,
+                98, 1001, 9, 1002, 21, 125, 20, 4, 20, 1105, 1, 46, 104, 999, 1105, 1, 46, 1101,
+                1000, 1, 20, 4, 20, 1105, 1, 46, 98, 99
+            ]
+        );
+
+        let expected = b"1001";
         assert_eq!(expected.to_vec(), writer)
     }
 }