Day 13: Pong AI done

2 files changed, 105 insertions, 22 deletions

diff --git a/src/bin/day_13.rs b/src/bin/day_13.rs
index a7d6128..ac1c478 100644
--- a/src/bin/day_13.rs
+++ b/src/bin/day_13.rs
@@ -1,8 +1,12 @@
 use aoc2019::*;
+use rpds::list;
+use rpds::list::List;
 use rpds::vector::Vector;
 use rpds::RedBlackTreeMap;
+use std::cmp::Ordering;
 use std::io;
 use std::io::prelude::*;
+use std::iter;
 use std::process;
 use structopt::StructOpt;
 
@@ -15,8 +19,8 @@ use structopt::StructOpt;
 ///
 /// See https://adventofcode.com/2019/day/13 for details.
 struct Opt {
-    #[structopt(short = "i", long = "input")]
-    input: Vec<Intcode>,
+    #[structopt(short = "q", long = "quarters")]
+    quarters: Option<Intcode>,
 }
 
 fn main() {
@@ -29,12 +33,18 @@ fn main() {
         .map(|x| exit_on_failed_assertion(x, "Error reading input"))
         .map(|x| exit_on_failed_assertion(String::from_utf8(x), "Input was not valid UTF-8"))
         .map(|x| exit_on_failed_assertion(x.trim().parse::<Intcode>(), "Invalid number"))
-        .collect::<IntcodeProgram>()
-        .with_input(opt.input.into_iter().collect());
+        .collect::<IntcodeProgram>();
 
-    let result = exit_on_failed_assertion(program.execute(), "Program errored");
-
-    println!("{}", count_blocks(result));
+    match opt.quarters {
+        Some(quarters) => {
+            let result = score_on_won_game(program.with_mem_0(quarters));
+            println!("{}", result);
+        }
+        None => {
+            let result = exit_on_failed_assertion(program.execute(), "Program errored");
+            println!("{}", count_blocks(&result));
+        }
+    }
 }
 
 fn exit_on_failed_assertion<A, E: std::error::Error>(data: Result<A, E>, message: &str) -> A {
@@ -47,24 +57,93 @@ fn exit_on_failed_assertion<A, E: std::error::Error>(data: Result<A, E>, message
     }
 }
 
-fn render_screen(output: Vector<Intcode>) -> RedBlackTreeMap<(Intcode, Intcode), Intcode> {
-    (0..output.len() / 3)
-        .map(|i| i * 3)
-        .map(|i| {
-            (
-                output[i].clone(),
-                output[i + 1].clone(),
-                output[i + 2].clone(),
-            )
-        })
-        .fold(RedBlackTreeMap::new(), |acc, (x, y, tile)| {
-            acc.insert((x, y), tile)
-        })
+#[derive(Default, Clone)]
+struct Screen {
+    screen: RedBlackTreeMap<(Intcode, Intcode), Intcode>,
+    previous_ball: (Intcode, Intcode),
+    ball: (Intcode, Intcode),
+    paddle: (Intcode, Intcode),
+    score: Intcode,
 }
 
-fn count_blocks(output: Vector<Intcode>) -> usize {
-    render_screen(output)
+impl Screen {
+    fn render(output: &Vector<Intcode>) -> Screen {
+        (0..output.len() / 3)
+            .map(|i| i * 3)
+            .map(|i| {
+                (
+                    output[i].clone(),
+                    output[i + 1].clone(),
+                    output[i + 2].clone(),
+                )
+            })
+            .fold(Screen::default(), |acc, (x, y, tile)| {
+                if x == Intcode::from(-1) && y == Intcode::from(0) {
+                    Screen { score: tile, ..acc }
+                } else if tile == Intcode::from(4) {
+                    Screen {
+                        ball: (x, y),
+                        previous_ball: acc.ball,
+                        ..acc
+                    }
+                } else if tile == Intcode::from(3) {
+                    Screen {
+                        paddle: (x.clone(), y.clone()),
+                        screen: acc.screen.insert((x, y), tile),
+                        ..acc
+                    }
+                } else if tile == Intcode::from(0) {
+                    Screen {
+                        screen: acc.screen.remove(&(x, y)),
+                        ..acc
+                    }
+                } else {
+                    Screen {
+                        screen: acc.screen.insert((x, y), tile),
+                        ..acc
+                    }
+                }
+            })
+    }
+
+    fn paddle_required_direction(&self) -> Intcode {
+        match self.paddle.0.cmp(&self.ball.0) {
+            Ordering::Less => Intcode::from(1),
+            Ordering::Equal => Intcode::from(0),
+            Ordering::Greater => Intcode::from(-1),
+        }
+    }
+}
+
+fn count_blocks(output: &Vector<Intcode>) -> usize {
+    Screen::render(output)
+        .screen
         .values()
         .filter(|val| **val == Intcode::from(2))
         .count()
 }
+
+fn score_on_won_game(program: IntcodeProgram) -> Intcode {
+    Screen::render(
+        &iter::successors(Some(program.run_to_termination_or_input()), |program| {
+            Some(next_game_state(program.clone()))
+        })
+        .take_while(|program| !program.halted)
+        .find(|program| count_blocks(&program.output) == 0)
+        .unwrap()
+        .output,
+    )
+    .score
+}
+
+fn next_game_state(program: IntcodeProgram) -> IntcodeProgram {
+    program_with_next_input(program.run_to_termination_or_input())
+}
+
+fn program_with_next_input(program: IntcodeProgram) -> IntcodeProgram {
+    program.with_input(list![next_input(&program.output)])
+}
+
+fn next_input(output: &Vector<Intcode>) -> Intcode {
+    Screen::render(output).paddle_required_direction()
+}
diff --git a/src/lib.rs b/src/lib.rs
index 6be1aba..5df59f4 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -71,6 +71,10 @@ pub fn bool_to_intcode(i: bool) -> Intcode {
 }
 
 impl IntcodeProgram {
+    pub fn with_mem_0(&self, val: Intcode) -> IntcodeProgram {
+        self.with_memory_set(0.into(), val)
+    }
+
     pub fn with_noun_verb_input(&self, noun: Intcode, verb: Intcode) -> IntcodeProgram {
         self.with_memory_set(1.into(), noun)
             .with_memory_set(2.into(), verb)