Did the feedback loop for the amplifier circuit

1 files changed, 116 insertions, 34 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,
+        }
     }
 }