Day 23: A network of intcodes
[advent-of-code-2019.git] / src / bin / day_23.rs
1 use aoc2019::*;
2 use rpds::list;
3 use rpds::list::List;
4 use rpds::vector::Vector;
5 use std::io;
6 use std::io::prelude::*;
7 use std::iter;
8 use std::process;
9 use structopt::StructOpt;
10
11 #[derive(Debug, StructOpt)]
12 #[structopt(name = "Day 23: Category Six")]
13 /// Executes an Intcode program on a network of computers
14 ///
15 /// See https://adventofcode.com/2019/day/23 for details.
16 struct Opt {
17     #[structopt(short = "n", long = "network-size")]
18     network_size: usize,
19     #[structopt(long = "nat")]
20     send_nat: bool,
21 }
22
23 fn main() {
24     let stdin = io::stdin();
25     let opt = Opt::from_args();
26
27     let program: IntcodeProgram = stdin
28         .lock()
29         .split(b',')
30         .map(|x| exit_on_failed_assertion(x, "Error reading input"))
31         .map(|x| exit_on_failed_assertion(String::from_utf8(x), "Input was not valid UTF-8"))
32         .map(|x| exit_on_failed_assertion(x.trim().parse::<Intcode>(), "Invalid number"))
33         .collect::<IntcodeProgram>();
34
35     let network = Network::new(program, opt.network_size);
36
37     if opt.send_nat {
38         println!("{}", network.first_repeated_nat_packet().y);
39     } else {
40         println!("{}", network.first_nat_packet().y);
41     }
42 }
43
44 fn exit_on_failed_assertion<A, E: std::error::Error>(data: Result<A, E>, message: &str) -> A {
45     match data {
46         Ok(data) => data,
47         Err(e) => {
48             eprintln!("{}: {}", message, e);
49             process::exit(1);
50         }
51     }
52 }
53
54 #[derive(Debug, Clone)]
55 struct Network {
56     computers: Vector<IntcodeProgram>,
57     nat: Option<Packet>,
58     is_potentially_idle: bool,
59     is_idle: bool,
60 }
61
62 impl Network {
63     fn new(program: IntcodeProgram, network_size: usize) -> Network {
64         Network {
65             computers: (0..network_size)
66                 .map(|ip| program.with_input(list![ip.into(), Intcode::from(-1)]))
67                 .collect(),
68             nat: None,
69             is_potentially_idle: false,
70             is_idle: false,
71         }
72     }
73
74     fn first_nat_packet(&self) -> Packet {
75         iter::successors(Some(self.clone()), |network| Some(network.next()))
76             .filter_map(|network| network.nat)
77             .next()
78             .unwrap()
79     }
80
81     fn first_repeated_nat_packet(&self) -> Packet {
82         self.sent_nat_packets()
83             .zip(self.sent_nat_packets().skip(1))
84             // .inspect(|(nat, _p2)| eprintln!("{}", nat.y))
85             .find(|(p1, p2)| p1.y == p2.y)
86             .unwrap()
87             .0
88     }
89
90     fn sent_nat_packets<'a>(&'a self) -> impl Iterator<Item = Packet> + 'a {
91         iter::successors(Some(self.clone()), |network| {
92             Some(network.with_nat_packet_sent().run_to_network_idle())
93         })
94         .filter_map(|network| network.nat)
95     }
96
97     fn run_to_network_idle(&self) -> Network {
98         iter::successors(Some(self.clone()), |network| Some(network.next()))
99             .find(|network| network.is_idle)
100             .unwrap()
101     }
102
103     fn with_nat_packet_sent(&self) -> Network {
104         if self.nat.is_some() {
105             Network {
106                 computers: self
107                     .computers
108                     .iter()
109                     .enumerate()
110                     .map(|(ip, computer)| {
111                         computer
112                             .with_additional_input(
113                                 self.nat
114                                     .iter()
115                                     .filter(|packet| packet.dest == Intcode::from(ip))
116                                     .flat_map(|packet| vec![packet.x.clone(), packet.y.clone()])
117                                     .collect(),
118                             )
119                             .run_to_termination_or_input()
120                     })
121                     .collect(),
122                 nat: self.nat.clone(),
123                 is_potentially_idle: false,
124                 is_idle: false,
125             }
126         } else {
127             self.clone()
128         }
129     }
130
131     fn next(&self) -> Network {
132         Network {
133             computers: self
134                 .computers
135                 .iter()
136                 .enumerate()
137                 .map(|(ip, computer)| {
138                     computer
139                         .with_cleared_output()
140                         .with_additional_input(if self.empty_queues() {
141                             list![Intcode::from(-1)]
142                         } else {
143                             list![]
144                         })
145                         .with_additional_input(
146                             self.pending_packets()
147                                 .filter(|packet| packet.dest == Intcode::from(ip))
148                                 .flat_map(|packet| vec![packet.x, packet.y])
149                                 .collect(),
150                         )
151                         .run_to_termination_or_input()
152                 })
153                 .collect(),
154             nat: self
155                 .pending_packets()
156                 .filter(|packet| packet.is_nat())
157                 .map(|packet| packet.with_dest(0.into()))
158                 .last()
159                 .or_else(|| self.nat.clone()),
160             is_potentially_idle: self.empty_queues(),
161             is_idle: self.is_potentially_idle && self.empty_queues(),
162         }
163     }
164
165     fn pending_packets<'a>(&'a self) -> impl Iterator<Item = Packet> + 'a {
166         self.computers.iter().flat_map(|computer| {
167             computer
168                 .output
169                 .iter()
170                 .cloned()
171                 .collect::<Vec<Intcode>>()
172                 .chunks_exact(3)
173                 .map(|packet| Packet {
174                     dest: packet[0].clone(),
175                     x: packet[1].clone(),
176                     y: packet[2].clone(),
177                 })
178                 .collect::<Vec<Packet>>()
179         })
180     }
181
182     fn pending_input<'a>(&'a self) -> impl Iterator<Item = Intcode> + 'a {
183         self.computers
184             .iter()
185             .flat_map(|computer| computer.input.iter().cloned().collect::<Vec<Intcode>>())
186     }
187
188     fn empty_queues(&self) -> bool {
189         self.pending_packets().count() == 0 && self.pending_input().count() == 0
190     }
191 }
192
193 #[derive(Debug, Clone)]
194 struct Packet {
195     dest: Intcode,
196     x: Intcode,
197     y: Intcode,
198 }
199
200 impl Packet {
201     fn is_nat(&self) -> bool {
202         self.dest == 255.into()
203     }
204
205     fn with_dest(&self, dest: Intcode) -> Packet {
206         Packet {
207             dest,
208             ..self.clone()
209         }
210     }
211 }