From 83e99215fb50264885e3df526ae5b5660b95eaca Mon Sep 17 00:00:00 2001
From: Antoine Martin <antoine97.martin@gmail.com>
Date: Sat, 7 Dec 2019 23:04:50 +0100
Subject: [PATCH] 2019: day07: part 2

---
 aoc2019/src/day07.rs       | 89 ++++++++++++++++++++++++++++++++++++++
 aoc2019/src/intcode/mod.rs | 26 ++++++++---
 2 files changed, 110 insertions(+), 5 deletions(-)

diff --git a/aoc2019/src/day07.rs b/aoc2019/src/day07.rs
index 834bec8..3aded78 100644
--- a/aoc2019/src/day07.rs
+++ b/aoc2019/src/day07.rs
@@ -12,6 +12,7 @@ pub fn run() -> Result<String> {
     let mut res = String::with_capacity(128);
 
     writeln!(res, "part 1: {}", part1(INPUT)?)?;
+    writeln!(res, "part 2: {}", part2(INPUT)?)?;
 
     Ok(res)
 }
@@ -67,6 +68,79 @@ fn part1(input: &str) -> Result<i64> {
     Ok(res)
 }
 
+fn part2(input: &str) -> Result<i64> {
+    let memory = parse_memory(input)?;
+
+    let combinations = permutations(&[5, 6, 7, 8, 9]);
+
+    let mut res = 0;
+    for combination in combinations {
+        let intcodes = &mut [
+            Intcode::with_memory(memory.clone()),
+            Intcode::with_memory(memory.clone()),
+            Intcode::with_memory(memory.clone()),
+            Intcode::with_memory(memory.clone()),
+            Intcode::with_memory(memory.clone()),
+        ];
+
+        for (phase, intcode) in combination.iter().zip(intcodes.iter_mut()) {
+            intcode.add_input(*phase);
+        }
+        intcodes[0].add_input(0);
+
+        let mut signal = None;
+        let mut num_halted = 0;
+        loop {
+            for i in 0..(intcodes.len() - 1) {
+                let (first, second) = intcodes.split_at_mut(i + 1);
+                let intcode = &mut first[i];
+                let next = &mut second[0];
+
+                let halted = intcode.run_and_wait()?;
+                if halted {
+                    num_halted += 1;
+                }
+
+                for out in intcode.output.iter() {
+                    next.add_input(*out);
+                }
+                intcode.output.clear();
+            }
+
+            let last_index = intcodes.len() - 1;
+            let (first, second) = intcodes.split_at_mut(last_index);
+            let first = &mut first[0];
+            let last = &mut second[0];
+            let halted = last.run_and_wait()?;
+
+            if halted {
+                let out = last
+                    .output
+                    .last()
+                    .copied()
+                    .ok_or_else(|| err!("last amplifier halted without output"))?;
+                signal = Some(out);
+            } else {
+                for out in last.output.iter() {
+                    first.add_input(*out);
+                }
+                last.output.clear();
+            }
+
+            if let Some(signal) = signal {
+                res = std::cmp::max(res, signal);
+                break;
+            }
+
+            if num_halted >= 4 {
+                return Err(err!("all non final amplifiers halted, feedback loop stuck"));
+            }
+        }
+    }
+
+    Ok(res)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -87,4 +161,19 @@ mod tests {
     fn part1_real() {
         assert_eq!(part1(INPUT).unwrap(), 844468);
     }
+
+    const PROVIDED4: &str =
+        "3,26,1001,26,-4,26,3,27,1002,27,2,27,1,27,26,27,4,27,1001,28,-1,28,1005,28,6,99,0,0,5";
+    const PROVIDED5: &str = "3,52,1001,52,-5,52,3,53,1,52,56,54,1007,54,5,55,1005,55,26,1001,54,-5,54,1105,1,12,1,53,54,53,1008,54,0,55,1001,55,1,55,2,53,55,53,4,53,1001,56,-1,56,1005,56,6,99,0,0,0,0,10";
+
+    #[test]
+    fn part2_provided() {
+        assert_eq!(part2(PROVIDED4).unwrap(), 139629729);
+        assert_eq!(part2(PROVIDED5).unwrap(), 18216);
+    }
+
+    #[test]
+    fn part2_real() {
+        assert_eq!(part2(INPUT).unwrap(), 4215746);
+    }
 }
diff --git a/aoc2019/src/intcode/mod.rs b/aoc2019/src/intcode/mod.rs
index 8a5ee13..8cf6e9b 100644
--- a/aoc2019/src/intcode/mod.rs
+++ b/aoc2019/src/intcode/mod.rs
@@ -24,12 +24,14 @@ enum Opcode {
     Halt,
 }
 
+#[derive(Debug)]
 pub struct Intcode {
     pub memory: Vec<i64>,
     input: Vec<i64>,
     pub output: Vec<i64>,
     ip: usize,
     next_input: usize,
+    wait_input: bool,
 }
 
 impl Intcode {
@@ -46,6 +48,7 @@ impl Intcode {
             output: Vec::new(),
             ip: 0,
             next_input: 0,
+            wait_input: false,
         }
     }
 
@@ -136,7 +139,7 @@ impl Intcode {
         }
     }
 
-    pub fn run(&mut self) -> Result<()> {
+    fn exec(&mut self) -> Result<bool> {
         loop {
             if self.ip >= self.memory.len() {
                 return Err(err!("reached end of program without halting"));
@@ -165,10 +168,12 @@ impl Intcode {
                     let input = if self.next_input < self.input.len() {
                         let res = self.input[self.next_input];
                         self.next_input += 1;
-                        Ok(res)
+                        res
+                    } else if self.wait_input {
+                        break Ok(false);
                     } else {
-                        Err(err!("tried to read input but it was empty"))
-                    }?;
+                        break Err(err!("tried to read input but it was empty"));
+                    };
                     dst.set(&mut self.memory, input)?;
 
                     self.ip += 2;
@@ -223,11 +228,22 @@ impl Intcode {
 
                     self.ip += 4;
                 }
-                Opcode::Halt => break Ok(()),
+                Opcode::Halt => break Ok(true),
             }
         }
     }
 
+    pub fn run_and_wait(&mut self) -> Result<bool> {
+        self.wait_input = true;
+        self.exec()
+    }
+
+    pub fn run(&mut self) -> Result<()> {
+        self.wait_input = false;
+        self.exec()?;
+        Ok(())
+    }
+
     pub fn get_day02_output(&self) -> Option<i64> {
         self.memory.get(0).copied()
     }