Added relative mode support. Currently isn't working.

1 files changed, 111 insertions, 21 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 5a5f9b3..57d0cb1 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -12,6 +12,7 @@ pub type Intcode = BigInt;
 #[derive(Debug, Clone)]
 pub struct IntcodeProgram {
     instruction_pointer: Intcode,
+    relative_base: Intcode,
     pub error: Option<IntcodeProgramError>,
     pub halted: bool,
     pub awaiting_input: bool,
@@ -24,6 +25,7 @@ impl FromIterator<Intcode> for IntcodeProgram {
     fn from_iter<I: IntoIterator<Item = Intcode>>(iter: I) -> Self {
         IntcodeProgram {
             instruction_pointer: Intcode::from(0),
+            relative_base: Intcode::from(0),
             error: None,
             halted: false,
             awaiting_input: false,
@@ -103,6 +105,13 @@ impl IntcodeProgram {
         }
     }
 
+    fn with_relative_base(&self, relative_base: Intcode) -> IntcodeProgram {
+        IntcodeProgram {
+            relative_base,
+            ..self.clone()
+        }
+    }
+
     fn with_memory_set(&self, address: Intcode, value: Intcode) -> IntcodeProgram {
         IntcodeProgram {
             memory: self.memory.insert(address, value),
@@ -156,6 +165,7 @@ impl IntcodeProgram {
                 6 => self.jump_if_false(opcode),
                 7 => self.less_than(opcode),
                 8 => self.equals(opcode),
+                9 => self.set_relative_base(opcode),
                 99 => self.halt(),
                 unknown => self.error(IntcodeProgramError::InvalidOpcode(unknown.clone())),
             })
@@ -163,7 +173,11 @@ impl IntcodeProgram {
     }
 
     fn add(&self, mode: &Intcode) -> IntcodeProgram {
-        match (self.get(1, mode), self.get(2, mode), self.get_immediate(3)) {
+        match (
+            self.get(1, mode),
+            self.get(2, mode),
+            self.get_literal(3, mode),
+        ) {
             (in1, in2, out) => self
                 .with_instruction_pointer_offset(4)
                 .with_memory_set(out, in1 + in2),
@@ -171,15 +185,19 @@ impl IntcodeProgram {
     }
 
     fn multiply(&self, mode: &Intcode) -> IntcodeProgram {
-        match (self.get(1, mode), self.get(2, mode), self.get_immediate(3)) {
+        match (
+            self.get(1, mode),
+            self.get(2, mode),
+            self.get_literal(3, mode),
+        ) {
             (in1, in2, out) => self
                 .with_instruction_pointer_offset(4)
                 .with_memory_set(out, in1 * in2),
         }
     }
 
-    fn input(&self, _mode: &Intcode) -> IntcodeProgram {
-        match (self.input.first().cloned(), self.get_immediate(1)) {
+    fn input(&self, mode: &Intcode) -> IntcodeProgram {
+        match (self.input.first().cloned(), self.get_literal(1, mode)) {
             (Some(input), out) => self
                 .with_instruction_pointer_offset(2)
                 .with_memory_set(out, input)
@@ -208,7 +226,11 @@ impl IntcodeProgram {
     }
 
     fn less_than(&self, mode: &Intcode) -> IntcodeProgram {
-        match (self.get(1, mode), self.get(2, mode), self.get_immediate(3)) {
+        match (
+            self.get(1, mode),
+            self.get(2, mode),
+            self.get_literal(3, mode),
+        ) {
             (in1, in2, out) => self.with_instruction_pointer_offset(4).with_memory_set(
                 out,
                 if in1 < in2 {
@@ -221,7 +243,11 @@ impl IntcodeProgram {
     }
 
     fn equals(&self, mode: &Intcode) -> IntcodeProgram {
-        match (self.get(1, mode), self.get(2, mode), self.get_immediate(3)) {
+        match (
+            self.get(1, mode),
+            self.get(2, mode),
+            self.get_literal(3, mode),
+        ) {
             (in1, in2, out) => self.with_instruction_pointer_offset(4).with_memory_set(
                 out,
                 if in1 == in2 {
@@ -233,6 +259,12 @@ impl IntcodeProgram {
         }
     }
 
+    fn set_relative_base(&self, mode: &Intcode) -> IntcodeProgram {
+        match self.get(1, mode) {
+            base_change => self.with_relative_base(self.relative_base.clone() + base_change),
+        }
+    }
+
     fn halt(&self) -> IntcodeProgram {
         IntcodeProgram {
             halted: true,
@@ -266,13 +298,29 @@ impl IntcodeProgram {
         match mode
             .to_radix_le(10)
             .1
-            .get(pointer_offset + 2)
+            .get(pointer_offset + 1)
             .cloned()
             .unwrap_or(0)
         {
             0 => self.get_position(pointer_offset),
             1 => self.get_immediate(pointer_offset),
-            _ => Intcode::from(0), // TODO: Relative mode
+            2 => self.get_relative(pointer_offset),
+            _ => Intcode::from(0),
+        }
+    }
+
+    fn get_literal(&self, pointer_offset: usize, mode: &Intcode) -> Intcode {
+        match mode
+            .to_radix_le(10)
+            .1
+            .get(pointer_offset + 1)
+            .cloned()
+            .unwrap_or(0)
+        {
+            0 => self.get_immediate(pointer_offset),
+            1 => self.get_immediate(pointer_offset),
+            2 => self.get_immediate_relative(pointer_offset),
+            _ => Intcode::from(0),
         }
     }
 
@@ -289,6 +337,17 @@ impl IntcodeProgram {
             .cloned()
             .unwrap_or(Intcode::from(0))
     }
+
+    fn get_relative(&self, pointer_offset: usize) -> Intcode {
+        self.memory
+            .get(&(self.get_immediate(pointer_offset) + self.relative_base.clone()))
+            .cloned()
+            .unwrap_or(Intcode::from(0))
+    }
+
+    fn get_immediate_relative(&self, pointer_offset: usize) -> Intcode {
+        self.get_immediate(pointer_offset) + self.relative_base.clone()
+    }
 }
 
 #[derive(Debug, PartialEq, Clone)]
@@ -313,22 +372,53 @@ impl std::error::Error for IntcodeProgramError {}
 mod tests {
     use super::*;
 
-    #[test]
-    fn day_2_example_1() {
-        let program: IntcodeProgram = vec![1, 0, 0, 0, 99]
+    fn i32_vec_to_intcode_program(input: Vec<i32>) -> IntcodeProgram {
+        input.into_iter().map(Intcode::from).collect()
+    }
+
+    fn i32_vec_to_intcode_memory(input: Vec<i32>) -> RedBlackTreeMap<Intcode, Intcode> {
+        input
             .into_iter()
-            .map(Intcode::from)
-            .collect::<IntcodeProgram>()
-            .run_to_termination();
+            .enumerate()
+            .map(|(i, val)| (Intcode::from(i), Intcode::from(val)))
+            .collect()
+    }
+
+    fn test_example_program(
+        before_execution: Vec<i32>,
+        after_execution: Vec<i32>,
+    ) -> IntcodeProgram {
+        let program = i32_vec_to_intcode_program(before_execution).run_to_termination();
 
         assert_eq!(program.error, None);
-        assert_eq!(
-            program.memory,
-            vec![2, 0, 0, 0, 99]
-                .into_iter()
-                .enumerate()
-                .map(|(i, val)| (Intcode::from(i), Intcode::from(val)))
-                .collect::<RedBlackTreeMap<Intcode, Intcode>>()
+        assert_eq!(program.memory, i32_vec_to_intcode_memory(after_execution));
+        program
+    }
+
+    #[test]
+    fn day_2_example_1() {
+        test_example_program(vec![1, 0, 0, 0, 99], vec![2, 0, 0, 0, 99]);
+    }
+
+    #[test]
+    fn day_2_example_2() {
+        test_example_program(vec![2, 3, 0, 3, 99], vec![2, 3, 0, 6, 99]);
+    }
+
+    #[test]
+    fn day_5_example_1() {
+        test_example_program(vec![1002, 4, 3, 4, 33], vec![1002, 4, 3, 4, 99]);
+    }
+
+    #[test]
+    fn day_9_example_1() {
+        let program = test_example_program(
+            vec![
+                109, 1, 204, -1, 1001, 100, 1, 100, 1008, 100, 16, 101, 1006, 101, 0, 99,
+            ],
+            vec![
+                109, 1, 204, -1, 1001, 100, 1, 100, 1008, 100, 16, 101, 1006, 101, 0, 99,
+            ],
         );
     }
 }