From a193837f0d9405f1b60e4a9165b9a96acffd0a2c Mon Sep 17 00:00:00 2001 From: timeshifter Date: Sun, 19 Jul 2026 20:21:01 +0200 Subject: [PATCH] add problem 11 --- problem11/Cargo.lock | 242 ++++++++++++++++++ problem11/Cargo.toml | 10 + problem11/src/main.rs | 565 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 817 insertions(+) create mode 100644 problem11/Cargo.lock create mode 100644 problem11/Cargo.toml create mode 100644 problem11/src/main.rs diff --git a/problem11/Cargo.lock b/problem11/Cargo.lock new file mode 100644 index 0000000..719e136 --- /dev/null +++ b/problem11/Cargo.lock @@ -0,0 +1,242 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "anyhow" +version = "1.0.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + +[[package]] +name = "bytes" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "itertools" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b4baf93f58d4425749ca49a51c50ebab072c5df6994d08fed93541c331481dc" +dependencies = [ + "either", +] + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "mio" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d65c71f1ce40ab09135ce117d742b9f8a19ff91a41a8b57ed50bc2de59c427" +dependencies = [ + "libc", + "wasi", + "windows-sys", +] + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "problem11" +version = "0.1.0" +dependencies = [ + "anyhow", + "bytes", + "itertools", + "tokio", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "socket2" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "tokio" +version = "1.53.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d988bcd52dbe076d3d46903332f58c912b87a2c49b1428419a5845154762ffee" +dependencies = [ + "bytes", + "libc", + "mio", + "parking_lot", + "pin-project-lite", + "signal-hook-registry", + "socket2", + "tokio-macros", + "windows-sys", +] + +[[package]] +name = "tokio-macros" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6328af13490e73a9b4694030fafd93f8c8c6a9dede33e821c3fc63eddf8042ba" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] diff --git a/problem11/Cargo.toml b/problem11/Cargo.toml new file mode 100644 index 0000000..8df3920 --- /dev/null +++ b/problem11/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "problem11" +version = "0.1.0" +edition = "2024" + +[dependencies] +anyhow = "1" +tokio = { version = "1", features = ["full"] } +bytes = "*" +itertools = "*" diff --git a/problem11/src/main.rs b/problem11/src/main.rs new file mode 100644 index 0000000..be0b095 --- /dev/null +++ b/problem11/src/main.rs @@ -0,0 +1,565 @@ +use std::{collections::HashMap, fmt::Display, sync::Arc}; + +use anyhow::{Result, anyhow, bail}; +use bytes::BytesMut; +use itertools::Itertools; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, + sync::{Mutex, MutexGuard}, +}; + +const MAX_MSG_LEN: usize = 1 << 20; +const UPSTREAM: &str = "pestcontrol.protohackers.com:20547"; + +type Policies = HashMap<(Action, u32, String), u32>; // action, site, species, policyid + +#[tokio::main] +async fn main() -> Result<()> { + let listener = TcpListener::bind("0.0.0.0:10000").await?; + let policies = Arc::new(Mutex::new(Policies::new())); + + loop { + let (mut socket, _) = listener.accept().await?; + let this_policies = policies.clone(); + let mut hello_has_been_sent = false; + + tokio::spawn(async move { + match handle_connection(&mut socket, this_policies, &mut hello_has_been_sent).await { + Ok(()) => (), + Err(e) => { + if !hello_has_been_sent { + socket + .write_all(&encode_message(Message::Hello).unwrap()) + .await + .unwrap(); + } + socket + .write_all( + &encode_message(Message::Error { + message: e.root_cause().to_string(), + }) + .unwrap(), + ) + .await + .unwrap(); + } + } + }); + } +} + +async fn handle_connection( + socket: &mut TcpStream, + policies: Arc>, + hello_has_been_sent: &mut bool, +) -> Result<()> { + let (mut reader, mut writer) = socket.split(); + let mut buf = BytesMut::with_capacity(4096); + + loop { + while let Some(frame) = try_extract_frame(&mut buf)? { + let msg = decode_message(&frame)?; + match msg { + Message::Hello => { + let reply = encode_message(Message::Hello)?; + writer.write_all(&reply).await?; + *hello_has_been_sent = true; + } + Message::Error { message } => { + bail!("client sent error: {message}"); + } + Message::SiteVisit { site, populations } => { + match handle_site_visit(site, populations, &policies).await { + Ok(()) => (), + Err(e) => { + if e.root_cause().to_string() + == DataError::ConflictingPopulationReport.to_string() + { + writer + .write_all(&encode_message(Message::Error { + message: e.root_cause().to_string(), + })?) + .await?; + } + } + } + } + msg => bail!("unexpected client message: {msg:?}"), + }; + } + + if reader.read_buf(&mut buf).await? == 0 { + if !buf.is_empty() { + return Err(anyhow!("connection closed mid-message")); + } + return Ok(()); + } + } +} + +async fn handle_site_visit( + reported_site: u32, + reported_populations: Vec<(String, u32)>, + policies: &Arc>, +) -> Result<()> { + let mut upstream = TcpStream::connect(UPSTREAM).await?; + + if conflicted(&reported_populations) { + return Err(anyhow!(DataError::ConflictingPopulationReport)); + } + + let buffer = encode_message(Message::Hello)?; + upstream.write_all(&buffer).await?; + + let mut read_buf = BytesMut::with_capacity(4096); + loop { + upstream.read_buf(&mut read_buf).await?; + if let Some(frame) = try_extract_frame(&mut read_buf)? { + if decode_message(&frame)? == Message::Hello { + break; + } else { + bail!("expected hello, received something else"); + } + } + } + let buf = encode_message(Message::DialAuthority { + site: reported_site, + })?; + upstream.write_all(&buf).await?; + + upstream.read_buf(&mut read_buf).await?; + + { + if let Some(frame) = try_extract_frame(&mut read_buf)? { + let msg = decode_message(&frame)?; + let goal_policies = match msg { + Message::TargetPopulations { site, populations } => { + work_out_policies(site, populations, reported_site, reported_populations) + } + m => { + return Err(anyhow!( + "unexpected upstream message after site visit: {m:?}" + )); + } + }; + let mut established_policies = policies.lock().await; + for (species, action) in goal_policies { + if established_policies.contains_key(&( + action.clone(), + reported_site, + species.clone(), + )) { + continue; + } else { + let policy_id = create_policy(&species, &action, &mut upstream).await?; + let entry = (action, reported_site, species); + if let Some(policy_id) = policy_id { + established_policies.insert(entry.clone(), policy_id); + } + remove_conflicting_policies(&mut established_policies, entry, &mut upstream) + .await?; + } + } + } + } + + Ok(()) +} + +fn conflicted(reported_populations: &[(String, u32)]) -> bool { + reported_populations + .iter() + .combinations(2) + .any(|vec| vec[0].0 == vec[1].0 && vec[0].1 != vec[1].1) +} + +async fn remove_conflicting_policies( + established_policies: &mut MutexGuard<'_, HashMap<(Action, u32, String), u32>>, + entry: (Action, u32, String), + upstream: &mut TcpStream, +) -> Result<()> { + let actions_to_remove = match entry.0 { + Action::Conserve => vec![Action::Cull], + Action::Cull => vec![Action::Conserve], + Action::Maintain => vec![Action::Cull, Action::Conserve], + }; + + for action_to_remove in actions_to_remove { + let mut this_entry = entry.clone(); + this_entry.0 = action_to_remove; + if let Some(policy_id) = established_policies.remove(&this_entry) { + let msg = Message::DeletePolicy { policy: policy_id }; + let buf = encode_message(msg)?; + upstream.write_all(&buf).await?; + + let mut buf = BytesMut::with_capacity(4096); + upstream.read_buf(&mut buf).await?; + if let Some(frame) = try_extract_frame(&mut buf)? { + let msg = decode_message(&frame)?; + match msg { + Message::Ok => (), + _ => bail!("expected ok after DeletePolicy, received something else"), + } + } + } + } + + Ok(()) +} + +async fn create_policy( + species: &str, + action: &Action, + upstream: &mut TcpStream, +) -> Result> { + if *action == Action::Maintain { + return Ok(None); + } + + let msg = Message::CreatePolicy { + species: species.to_string(), + action: action.to_owned(), + }; + + let bytes = encode_message(msg)?; + upstream.write_all(&bytes).await?; + + let mut buf = BytesMut::with_capacity(4096); + loop { + upstream.read_buf(&mut buf).await?; + if let Some(frame) = try_extract_frame(&mut buf)? { + let msg = decode_message(&frame)?; + match msg { + Message::PolicyResult { policy } => return Ok(Some(policy)), + _ => bail!("expected 0x57 PolicyResult, got something else"), + } + } + } +} + +fn work_out_policies( + site: u32, + target_populations: Vec<(String, u32, u32)>, + reported_site: u32, + reported_populations: Vec<(String, u32)>, +) -> Vec<(String, Action)> { + assert!(site == reported_site); + let mut goal_policies = vec![]; + + for (species, count) in &reported_populations { + match target_populations.iter().find(|(s, _, _)| *species == *s) { + Some((_, min, max)) => { + let action = if (*min..=*max).contains(count) { + Action::Maintain + } else if count < min { + Action::Conserve + } else { + Action::Cull + }; + goal_policies.push((species.clone(), action)) + } + None => { + continue; + } + } + } + + for (species, min, _) in target_populations { + if min > 0 && reported_populations.iter().all(|(s, _)| *s != species) { + goal_policies.push((species, Action::Conserve)); + } + } + + goal_policies +} + +fn encode_message(msg: Message) -> Result, ProtoError> { + let mut w = Writer::new(); + + let msg_bytes = match msg { + Message::Hello => { + w.str("pestcontrol"); + w.u32(1); + write_message(0x50, &w.buf) + } + Message::Error { message } => { + w.str(&message); + write_message(0x51, &w.buf) + } + Message::Ok => write_message(0x52, &w.buf), + Message::DialAuthority { site } => { + w.u32(site); + write_message(0x53, &w.buf) + } + Message::TargetPopulations { site, populations } => { + w.u32(site); + w.array(&populations, |w, v| { + w.str(&v.0); + w.u32(v.1); + w.u32(v.2) + }); + write_message(0x54, &w.buf) + } + Message::CreatePolicy { species, action } => { + w.str(&species); + match action { + Action::Conserve => w.u8(0xa0), + Action::Cull => w.u8(0x90), + Action::Maintain => panic!("maintain is not for sending"), + }; + write_message(0x55, &w.buf) + } + Message::DeletePolicy { policy } => { + w.u32(policy); + write_message(0x56, &w.buf) + } + Message::PolicyResult { policy } => { + w.u32(policy); + write_message(0x57, &w.buf) + } + Message::SiteVisit { site, populations } => { + w.u32(site); + w.array(&populations, |w, v| { + w.str(&v.0); + w.u32(v.1) + }); + write_message(0x58, &w.buf) + } + }; + + Ok(msg_bytes) +} + +fn decode_message(frame: &[u8]) -> Result { + let sum: u32 = frame.iter().map(|&b| b as u32).sum(); + if !sum.is_multiple_of(256) { + return Err(ProtoError::WrongChecksum); + } + + let type_byte = frame[0]; + let mut r = Reader::new(&frame[5..frame.len() - 1]); + let msg = match type_byte { + 0x50 => { + if r.str()? != "pestcontrol" || r.u32()? != 1 { + return Err(ProtoError::Malformed); + } + Message::Hello + } + 0x51 => Message::Error { message: r.str()? }, + 0x52 => Message::Ok, + 0x53 => Message::DialAuthority { site: r.u32()? }, + 0x54 => { + let site = r.u32()?; + let populations = r.array(|r| Ok((r.str()?, r.u32()?, r.u32()?)))?; + Message::TargetPopulations { site, populations } + } + 0x55 => { + let species = r.str()?; + let action = match r.u8()? { + 0x90 => Action::Cull, + 0xa0 => Action::Conserve, + _ => { + return Err(ProtoError::InvalidPolicy); + } + }; + Message::CreatePolicy { species, action } + } + 0x56 => Message::DeletePolicy { policy: r.u32()? }, + 0x57 => Message::PolicyResult { policy: r.u32()? }, + 0x58 => { + let site = r.u32()?; + let populations = r.array(|r| Ok((r.str()?, r.u32()?)))?; + Message::SiteVisit { site, populations } + } + _ => return Err(ProtoError::WrongTypeByte), + }; + + if !r.finished() { + return Err(ProtoError::Malformed); + } + + Ok(msg) +} + +fn try_extract_frame(buf: &mut BytesMut) -> Result, ProtoError> { + if buf.len() < 5 { + return Ok(None); // less than type byte + length + } + let len = u32::from_be_bytes(buf[1..5].try_into().unwrap()) as usize; + + if !(6..MAX_MSG_LEN).contains(&len) { + return Err(ProtoError::Malformed); + } + + if buf.len() < len { + return Ok(None); + } + + Ok(Some(buf.split_to(len))) +} + +#[derive(Debug)] +struct Writer { + buf: Vec, +} + +impl Writer { + fn u8(&mut self, val: u8) { + self.buf.push(val); + } + + fn u32(&mut self, val: u32) { + self.buf.extend_from_slice(&val.to_be_bytes()); + } + + fn str(&mut self, s: &str) { + self.u32(s.len() as u32); + self.buf.extend_from_slice(s.as_bytes()); + } + + fn array(&mut self, items: &[T], mut elem: impl FnMut(&mut Self, &T)) { + self.u32(items.len() as u32); + for item in items { + elem(self, item); + } + } + + fn new() -> Self { + Self { + buf: Vec::with_capacity(4096), + } + } +} + +#[derive(Debug)] +struct Reader<'a> { + buf: &'a [u8], + pos: usize, +} + +impl<'a> Reader<'a> { + fn u8(&mut self) -> Result { + Ok(self.take(1)?[0]) + } + + fn u32(&mut self) -> Result { + let b = self.take(4)?; + Ok(u32::from_be_bytes(b.try_into().unwrap())) + } + + fn str(&mut self) -> Result { + let len = self.u32()? as usize; + let bytes = self.take(len)?; + // spec says ASCII, so this can be a straight cast rather than utf8 validation + Ok(bytes.iter().map(|&b| b as char).collect()) + } + + fn array( + &mut self, + mut elem: impl FnMut(&mut Self) -> Result, + ) -> Result, ProtoError> { + let len = self.u32()? as usize; + (0..len).map(|_| elem(self)).collect() + } + + fn take(&mut self, n: usize) -> Result<&'a [u8], ProtoError> { + if self.pos + n > self.buf.len() { + return Err(ProtoError::Malformed); + } + let s = &self.buf[self.pos..self.pos + n]; + self.pos += n; + Ok(s) + } + + fn finished(&self) -> bool { + self.pos == self.buf.len() + } + + fn new(buf: &'a [u8]) -> Self { + Self { buf, pos: 0 } + } +} + +#[derive(Debug)] +enum ProtoError { + Malformed, + WrongTypeByte, + WrongChecksum, + InvalidPolicy, +} + +#[derive(Debug)] +enum DataError { + ConflictingPopulationReport, +} + +impl Display for DataError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + DataError::ConflictingPopulationReport => f.write_str("conflicting population report"), + } + } +} + +impl std::error::Error for DataError {} + +impl Display for ProtoError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ProtoError::Malformed => f.write_str("malformed packet"), + ProtoError::WrongTypeByte => f.write_str("wrong type byte"), + ProtoError::WrongChecksum => f.write_str("checksum error"), + ProtoError::InvalidPolicy => f.write_str("invalid policy"), + } + } +} + +impl std::error::Error for ProtoError {} + +fn write_message(type_byte: u8, payload: &[u8]) -> Vec { + let mut out = Vec::with_capacity(payload.len() + 6); + out.push(type_byte); + out.extend_from_slice(&((payload.len() + 6) as u32).to_be_bytes()); + out.extend_from_slice(payload); + let sum: u32 = out.iter().map(|&b| b as u32).sum(); + out.push((256 - sum % 256) as u8); + out +} + +#[derive(Debug, PartialEq)] +enum Message { + Hello, + Error { + message: String, + }, + Ok, + DialAuthority { + site: u32, + }, + TargetPopulations { + site: u32, + populations: Vec<(String, u32, u32)>, + }, + CreatePolicy { + species: String, + action: Action, + }, + DeletePolicy { + policy: u32, + }, + PolicyResult { + policy: u32, + }, + SiteVisit { + site: u32, + populations: Vec<(String, u32)>, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +enum Action { + Conserve, + Cull, + Maintain, +}