Did the feedback loop for the amplifier circuit

2 files changed, 167 insertions, 45 deletions

diff --git a/src/bin/day_7.rs b/src/bin/day_7.rs
index d412734..0abf215 100644
--- a/src/bin/day_7.rs
+++ b/src/bin/day_7.rs
@@ -1,8 +1,10 @@
 use aoc2019::*;
 use rpds::list;
 use rpds::list::List;
+use rpds::vector::Vector;
 use std::io;
 use std::io::prelude::*;
+use std::iter;
 use std::process;
 use structopt::StructOpt;
 
@@ -28,7 +30,10 @@ fn main() {
         .map(|x| exit_on_failed_assertion(x.trim().parse::<Intcode>(), "Invalid number"))
         .collect::<IntcodeProgram>();
 
-    let result = exit_on_failed_assertion(find_max_power(&program), "Program errored");
+    let result = exit_on_failed_assertion(
+        find_max_power(&program, opt.feedback_loop_mode),
+        "Program errored",
+    );
     println!("{}", result);
 }
 
@@ -42,50 +47,127 @@ fn exit_on_failed_assertion<A, E: std::error::Error>(data: Result<A, E>, message
     }
 }
 
-fn find_max_power(program: &IntcodeProgram) -> Result<Intcode, IntcodeProgramError> {
-    PhaseSetting::all()
-        .map(|phase| run_amplifiers(program, phase))
-        .collect::<Result<Vec<Intcode>, _>>()
-        .map(|powers| powers.into_iter().max().unwrap_or(0))
-}
-
-fn run_amplifiers(
-    program: &IntcodeProgram,
-    phase: PhaseSetting,
-) -> Result<Intcode, IntcodeProgramError> {
-    run_amp(program, phase, 0, 0)
-        .and_then(|pipe| run_amp(program, phase, 1, pipe))
-        .and_then(|pipe| run_amp(program, phase, 2, pipe))
-        .and_then(|pipe| run_amp(program, phase, 3, pipe))
-        .and_then(|pipe| run_amp(program, phase, 4, pipe))
-}
-
-fn run_amp(
+fn find_max_power(
     program: &IntcodeProgram,
-    phase: PhaseSetting,
-    n: usize,
-    input: Intcode,
+    feedback_loop_mode: bool,
 ) -> Result<Intcode, IntcodeProgramError> {
-    program
-        .with_input(list![phase.0[n], input])
-        .execute()
-        .and_then(|output_vec| output_vec.get(0).cloned().ok_or(IntcodeProgramError))
+    PhaseSetting::all(feedback_loop_mode)
+        .map(|phase| AmplifierArray::new(program, phase).execute())
+        .collect::<Result<Vec<Intcode>, _>>()
+        .map(|powers| powers.into_iter().max().unwrap_or(0))
 }
 
 #[derive(Debug, Clone, Copy)]
 struct PhaseSetting([Intcode; 5]);
 
 impl PhaseSetting {
-    fn all() -> impl Iterator<Item = PhaseSetting> {
-        (0..=4)
+    fn all(feedback_loop_mode: bool) -> impl Iterator<Item = PhaseSetting> {
+        if feedback_loop_mode {
+            PhaseSetting::permute(5, 10)
+        } else {
+            PhaseSetting::permute(0, 5)
+        }
+    }
+
+    fn permute(min: Intcode, max: Intcode) -> impl Iterator<Item = PhaseSetting> {
+        // This is an absolutely atrocious way to do the permutation,
+        // but luckily it's only 5 elements long.
+        (min..max)
             .flat_map(move |a| {
-                (0..=4).flat_map(move |b| {
-                    (0..=4).flat_map(move |c| {
-                        (0..=4)
-                            .flat_map(move |d| (0..=4).map(move |e| PhaseSetting([a, b, c, d, e])))
+                (min..max).flat_map(move |b| {
+                    (min..max).flat_map(move |c| {
+                        (min..max).flat_map(move |d| {
+                            (min..max).map(move |e| PhaseSetting([a, b, c, d, e]))
+                        })
                     })
                 })
             })
-            .filter(|phase| (0..=4).all(|x| phase.0.contains(&x)))
+            .filter(move |phase| (min..max).all(|x| phase.0.contains(&x)))
+    }
+}
+
+#[derive(Debug, Clone)]
+struct AmplifierArray {
+    amplifiers: Vector<IntcodeProgram>,
+}
+
+impl AmplifierArray {
+    fn new(program: &IntcodeProgram, phase: PhaseSetting) -> AmplifierArray {
+        AmplifierArray {
+            amplifiers: (0..5)
+                .map(|n| AmplifierArray::new_amp(program, phase, n))
+                .collect(),
+        }
+    }
+
+    fn new_amp(program: &IntcodeProgram, phase: PhaseSetting, n: usize) -> IntcodeProgram {
+        if n == 0 {
+            program.with_input(list![phase.0[n], 0])
+        } else {
+            program.with_input(list![phase.0[n]])
+        }
+    }
+
+    fn execute(&self) -> Result<Intcode, IntcodeProgramError> {
+        self.run_to_termination().output_into_result()
+    }
+
+    fn run_to_termination(&self) -> AmplifierArray {
+        iter::successors(Some(self.clone()), |p| Some(p.next()))
+            .find(|p| p.is_terminated())
+            .unwrap() // successors doesn't terminate, so this will never be none.
+    }
+
+    fn output_into_result(&self) -> Result<Intcode, IntcodeProgramError> {
+        self.amplifiers
+            .first()
+            .and_then(|amp| amp.input.first().cloned())
+            .ok_or(IntcodeProgramError)
+    }
+
+    fn is_terminated(&self) -> bool {
+        self.amplifiers.last().map(|amp| amp.halted).unwrap_or(true)
+    }
+
+    fn next(&self) -> AmplifierArray {
+        self.run_amplifiers().update_inputs()
+    }
+
+    fn run_amplifiers(&self) -> AmplifierArray {
+        AmplifierArray {
+            amplifiers: self
+                .amplifiers
+                .iter()
+                .map(|amp| amp.run_to_termination_or_input())
+                .collect(),
+        }
+    }
+
+    fn update_inputs(&self) -> AmplifierArray {
+        AmplifierArray {
+            amplifiers: self
+                .amplifiers
+                .iter()
+                .fold(
+                    (
+                        Vector::new(),
+                        self.amplifiers
+                            .last()
+                            .map(|a| a.output.iter().cloned().collect::<List<Intcode>>())
+                            .unwrap_or(List::new()),
+                    ),
+                    |(amps, piped_input), next_amp| {
+                        (
+                            amps.push_back(
+                                next_amp
+                                    .with_additional_input(piped_input)
+                                    .with_cleared_output(),
+                            ),
+                            next_amp.output.iter().cloned().collect::<List<Intcode>>(),
+                        )
+                    },
+                )
+                .0,
+        }
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
index 2591a54..717505e 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -9,11 +9,12 @@ pub type Intcode = i32;
 #[derive(Debug, Clone)]
 pub struct IntcodeProgram {
     instruction_pointer: usize,
-    error: bool,
-    halted: bool,
+    pub error: bool,
+    pub halted: bool,
+    pub awaiting_input: bool,
     memory: Vector<Intcode>,
-    input: List<Intcode>,
-    output: Vector<Intcode>,
+    pub input: List<Intcode>,
+    pub output: Vector<Intcode>,
 }
 
 impl FromIterator<Intcode> for IntcodeProgram {
@@ -22,6 +23,7 @@ impl FromIterator<Intcode> for IntcodeProgram {
             instruction_pointer: 0,
             error: false,
             halted: false,
+            awaiting_input: false,
             memory: iter.into_iter().collect(),
             input: List::new(),
             output: Vector::new(),
@@ -41,6 +43,20 @@ impl IntcodeProgram {
         }
     }
 
+    pub fn with_additional_input(&self, input: List<Intcode>) -> IntcodeProgram {
+        IntcodeProgram {
+            input: self.input.iter().chain(input.iter()).cloned().collect(),
+            ..self.clone()
+        }
+    }
+
+    pub fn with_cleared_output(&self) -> IntcodeProgram {
+        IntcodeProgram {
+            output: Vector::new(),
+            ..self.clone()
+        }
+    }
+
     pub fn execute(&self) -> Result<Vector<Intcode>, IntcodeProgramError> {
         self.run_to_termination().output_into_result()
     }
@@ -49,6 +65,22 @@ impl IntcodeProgram {
         self.run_to_termination().memory_0_into_result()
     }
 
+    pub fn run_to_termination(&self) -> IntcodeProgram {
+        iter::successors(Some(self.clear_await_input()), |p| {
+            Some(IntcodeProgram::next(&p))
+        })
+        .find(|p| p.halted)
+        .unwrap() // successors doesn't terminate, so this will never be none.
+    }
+
+    pub fn run_to_termination_or_input(&self) -> IntcodeProgram {
+        iter::successors(Some(self.clear_await_input()), |p| {
+            Some(IntcodeProgram::next(&p))
+        })
+        .find(|p| p.halted || p.awaiting_input)
+        .unwrap()
+    }
+
     fn with_instruction_pointer(&self, instruction_pointer: usize) -> IntcodeProgram {
         IntcodeProgram {
             instruction_pointer,
@@ -105,13 +137,6 @@ impl IntcodeProgram {
             Ok(self.memory.get(0).cloned())
         }
     }
-
-    fn run_to_termination(&self) -> IntcodeProgram {
-        iter::successors(Some(self.clone()), |p| Some(IntcodeProgram::next(&p)))
-            .find(|p| p.halted)
-            .unwrap() // successors doesn't terminate, so this will never be none.
-    }
-
     fn next(&self) -> IntcodeProgram {
         //eprintln!("{:?}", self);
         self.memory
@@ -155,6 +180,7 @@ impl IntcodeProgram {
                 .with_instruction_pointer_offset(2)
                 .with_memory_set(out as usize, input)
                 .with_input_consumed(),
+            (None, Some(_out)) => self.await_input(),
             _ => self.error(),
         }
     }
@@ -206,6 +232,20 @@ impl IntcodeProgram {
         }
     }
 
+    fn await_input(&self) -> IntcodeProgram {
+        IntcodeProgram {
+            awaiting_input: true,
+            ..self.clone()
+        }
+    }
+
+    fn clear_await_input(&self) -> IntcodeProgram {
+        IntcodeProgram {
+            awaiting_input: false,
+            ..self.clone()
+        }
+    }
+
     fn error(&self) -> IntcodeProgram {
         IntcodeProgram {
             halted: true,