add problem 06

problem06/src/main.rs

@@ -0,0 +1,457 @@
+use std::collections::{HashMap, HashSet};
+use std::sync::{Arc, Mutex, MutexGuard};
+
+use anyhow::{Result, ensure};
+use itertools::Itertools;
+use tokio::io::AsyncWriteExt;
+use tokio::sync::broadcast::{self, Receiver, Sender};
+use tokio::{
+    io::AsyncReadExt,
+    net::{TcpListener, TcpStream},
+};
+
+type PlateLog = HashMap<String, Vec<(Camera, u32)>>;
+type TicketLog = HashSet<(String, usize)>;
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    let listener = TcpListener::bind("0.0.0.0:10000").await?;
+
+    let plates = Arc::new(Mutex::new(HashMap::new()));
+    let days_with_issued_tickets = Arc::new(Mutex::new(TicketLog::new()));
+    let ticket_queue = Arc::new(Mutex::new(vec![]));
+    let (tx, _) = broadcast::channel(100_000);
+
+    loop {
+        let (socket, _) = listener.accept().await?;
+        let this_plates = plates.clone();
+        let this_tx = tx.clone();
+        let this_queue = ticket_queue.clone();
+        let this_days_tickets = days_with_issued_tickets.clone();
+
+        tokio::spawn(async move {
+            match handle_connection(socket, this_plates, this_queue, this_days_tickets, this_tx)
+                .await
+            {
+                Ok(()) => (),
+                Err(e) => {
+                    dbg!(e);
+                }
+            }
+        });
+    }
+}
+
+async fn handle_connection(
+    mut socket: TcpStream,
+    plates: Arc<Mutex<PlateLog>>,
+    ticket_queue: Arc<Mutex<Vec<Ticket>>>,
+    days_with_issued_tickets: Arc<Mutex<TicketLog>>,
+    tx: Sender<ChannelMessage>,
+) -> Result<()> {
+    let (mut reader, mut writer) = socket.split();
+
+    let mut heartbeat_interval = None;
+    let mut timer = tokio::time::interval(std::time::Duration::from_hours(100));
+    timer.tick().await;
+
+    loop {
+        tokio::select! {
+            _ = timer.tick() => {
+                writer.write_all(&[0x41]).await?;
+            }
+            byte = reader.read_u8() => {
+                match byte? {
+                    0x40 => {
+                        if heartbeat_interval.is_some() {
+                            send_error_message(&mut writer, "multiple heartbeat requests").await?;
+                            return Ok(());
+                        }
+                        let interval = reader.read_u32().await?;
+                        heartbeat_interval = Some(interval);
+                        if interval > 0 {
+                            timer = tokio::time::interval(std::time::Duration::from_millis(
+                                100.max((interval as u64 * 100) - 100),
+                            ));
+                        }
+                    }
+                    0x80 => {
+                        camera_main_loop(reader, writer, plates, ticket_queue, timer, heartbeat_interval, tx).await?;
+                        return Ok(());
+                    }
+                    0x81 => {
+                        let rx = tx.subscribe();
+                        dispatcher_main_loop(reader, writer, ticket_queue, days_with_issued_tickets, timer, heartbeat_interval, rx)
+                            .await?;
+                        return Ok(());
+                    }
+                    other => {
+                        return send_error_message(
+                            &mut writer,
+                            &format!("initialized with wrong message type 0x{other:x}"),
+                        )
+                        .await;
+                    }
+                };
+            }
+        }
+    }
+}
+
+async fn dispatcher_main_loop(
+    mut reader: tokio::net::tcp::ReadHalf<'_>,
+    mut writer: tokio::net::tcp::WriteHalf<'_>,
+    ticket_queue: Arc<Mutex<Vec<Ticket>>>,
+    days_with_issued_tickets: Arc<Mutex<TicketLog>>,
+    mut timer: tokio::time::Interval,
+    mut heartbeat_interval: Option<u32>,
+    mut rx: Receiver<ChannelMessage>,
+) -> Result<()> {
+    let num_roads = reader.read_u8().await?;
+    let mut roads = vec![];
+
+    for _ in 0..num_roads {
+        roads.push(reader.read_u16().await?);
+    }
+
+    work_ticket_queue(
+        &roads,
+        &mut writer,
+        &ticket_queue,
+        &days_with_issued_tickets,
+    )
+    .await?;
+
+    loop {
+        tokio::select! {
+            _ = timer.tick() => {
+                writer.write_all(&[0x41]).await?;
+            }
+            _ = rx.recv() => {
+                work_ticket_queue(&roads, &mut writer, &ticket_queue, &days_with_issued_tickets).await?;
+            }
+            byte = reader.read_u8() => {
+                match byte? {
+                    0x40 => {
+                        if heartbeat_interval.is_some() {
+                            send_error_message(&mut writer, "multiple heartbeat requests").await?;
+                            return Ok(());
+                        }
+                        heartbeat_interval = Some(reader.read_u32().await?);
+                        if let Some(interval) = heartbeat_interval
+                            && interval > 0
+                        {
+                            timer = tokio::time::interval(std::time::Duration::from_millis(
+                                100.max((interval * 100) as u64 - 100),
+                            ));
+                        }
+                    }
+                    _ => {
+                        return send_error_message(&mut writer, "received wrong message type").await;
+                    }
+                }
+            }
+
+        }
+    }
+}
+
+async fn work_ticket_queue(
+    roads: &[u16],
+    writer: &mut tokio::net::tcp::WriteHalf<'_>,
+    ticket_queue: &Arc<Mutex<Vec<Ticket>>>,
+    days_with_issued_tickets: &Arc<Mutex<TicketLog>>,
+) -> Result<()> {
+    let messages = generate_ticket_messages(ticket_queue, roads, days_with_issued_tickets).await?;
+
+    for message in messages {
+        writer.write_all(&message).await?;
+    }
+
+    Ok(())
+}
+
+async fn generate_ticket_messages(
+    queue: &Arc<Mutex<Vec<Ticket>>>,
+    roads: &[u16],
+    days_with_issued_tickets: &Arc<Mutex<TicketLog>>,
+) -> Result<Vec<Vec<u8>>> {
+    let mut queue = queue.lock().unwrap();
+
+    let mut days = days_with_issued_tickets.lock().unwrap();
+
+    let own_tickets = queue
+        .iter()
+        .positions(|ticket| roads.contains(&ticket.road));
+
+    let mut messages = vec![];
+    let mut to_remove = vec![];
+
+    for idx in own_tickets {
+        let ticket = queue.get(idx).unwrap();
+        to_remove.push(idx);
+
+        let day1 = timestamp_to_day(ticket.timestamp1);
+        let day2 = timestamp_to_day(ticket.timestamp2);
+
+        let key1 = (ticket.plate.clone(), day1);
+        let key2 = (ticket.plate.clone(), day2);
+
+        if days.contains(&key1) || days.contains(&key2) {
+            continue;
+        }
+
+        days.insert(key1);
+        days.insert(key2);
+
+        let to_wire = ticket_to_vec_u8(ticket);
+
+        messages.push(to_wire);
+    }
+
+    to_remove.sort();
+
+    for idx in to_remove.iter().rev() {
+        queue.remove(*idx);
+    }
+
+    Ok(messages)
+}
+
+fn timestamp_to_day(timestamp: u32) -> usize {
+    (timestamp / 86400) as _
+}
+
+fn ticket_to_vec_u8(ticket: &Ticket) -> Vec<u8> {
+    let mut to_wire = vec![0x21];
+
+    // plate
+    to_wire.push(ticket.plate.len() as u8);
+    for byte in ticket.plate.as_bytes() {
+        to_wire.push(*byte);
+    }
+
+    // road
+    for byte in ticket.road.to_be_bytes() {
+        to_wire.push(byte);
+    }
+
+    // mile 1, timestamp 1
+    for byte in ticket.mile1.to_be_bytes() {
+        to_wire.push(byte);
+    }
+    for byte in ticket.timestamp1.to_be_bytes() {
+        to_wire.push(byte);
+    }
+
+    // mile 2, timestamp 2
+    for byte in ticket.mile2.to_be_bytes() {
+        to_wire.push(byte);
+    }
+    for byte in ticket.timestamp2.to_be_bytes() {
+        to_wire.push(byte);
+    }
+
+    // speed
+    for byte in ticket.speed.to_be_bytes() {
+        to_wire.push(byte);
+    }
+
+    to_wire
+}
+
+async fn camera_main_loop(
+    mut reader: tokio::net::tcp::ReadHalf<'_>,
+    mut writer: tokio::net::tcp::WriteHalf<'_>,
+    plates: Arc<Mutex<PlateLog>>,
+    ticket_queue: Arc<Mutex<Vec<Ticket>>>,
+    mut timer: tokio::time::Interval,
+    mut heartbeat_interval: Option<u32>,
+    tx: Sender<ChannelMessage>,
+) -> Result<()> {
+    let myself = Camera {
+        road: reader.read_u16().await?,
+        mile: reader.read_u16().await?,
+        limit: reader.read_u16().await?,
+    };
+
+    loop {
+        tokio::select! {
+            _ = timer.tick() => {
+                writer.write_all(&[0x41]).await?;
+            }
+            byte = reader.read_u8() => {
+                match byte? {
+
+                    0x20 => {
+                        let plate = read_string(&mut reader).await?;
+                        let timestamp = reader.read_u32().await?;
+                        {
+                            let mut this_plates = plates.lock().unwrap();
+                            this_plates
+                                .entry(plate.clone())
+                                .and_modify(|v| v.push((myself.clone(), timestamp)))
+                                .or_insert(vec![(myself.clone(), timestamp)]);
+
+                            for ticket in find_possible_tickets(this_plates, &myself, plate, timestamp)
+                            {
+                                ticket_queue.lock().unwrap().push(ticket);
+                                let _ = tx.send(ChannelMessage::TicketIssued);
+                            }
+                        }
+                    }
+                    0x40 => {
+                        if heartbeat_interval.is_some() {
+                            send_error_message(&mut writer, "multiple heartbeat requests").await?;
+                            return Ok(());
+                        }
+                        heartbeat_interval = Some(reader.read_u32().await?);
+                        if let Some(interval) = heartbeat_interval
+                            && interval > 0
+                        {
+                            timer = tokio::time::interval(std::time::Duration::from_millis(
+                                100.max((interval * 100) as u64 - 100),
+                            ));
+                        }
+                    }
+
+                    _ => {
+                        return send_error_message(&mut writer, "received wrong message type").await;
+                    }
+                }
+            }
+        }
+    }
+}
+
+fn find_possible_tickets(
+    plates: MutexGuard<'_, PlateLog>,
+    this_camera: &Camera,
+    this_plate: String,
+    this_timestamp: u32,
+) -> Vec<Ticket> {
+    let relevant: Vec<_> = plates
+        .iter()
+        .filter(|(plate, _)| **plate == this_plate)
+        .flat_map(|(_, camera_timestamp_map)| camera_timestamp_map)
+        .filter(|(camera, _)| camera.road == this_camera.road)
+        .collect();
+
+    let violations = relevant
+        .iter()
+        .combinations(2)
+        .filter(|data| {
+            data.iter()
+                .any(|(_, timestamp)| *timestamp == this_timestamp)
+        })
+        .map(|data| (data.clone(), calculate_speed(data).unwrap()))
+        .filter(|(_, speed)| *speed > this_camera.limit * 100);
+
+    let mut result = vec![];
+
+    for violation in violations {
+        let (camera_timestamp_map, speed) = violation;
+        let ticket = Ticket::new(
+            this_plate.clone(),
+            this_camera.road,
+            camera_timestamp_map[0].0.mile,
+            camera_timestamp_map[0].1,
+            camera_timestamp_map[1].0.mile,
+            camera_timestamp_map[1].1,
+            speed,
+        );
+
+        result.push(ticket);
+    }
+
+    result
+}
+
+fn calculate_speed(timestamps: Vec<&&(Camera, u32)>) -> Option<u16> {
+    let mut data = timestamps.clone();
+
+    data.sort_by_key(|(_, time)| time);
+
+    let dt = (data[1].1 - data[0].1) as u64;
+    let ds = (data[1].0.mile as i64 - data[0].0.mile as i64).unsigned_abs();
+
+    if dt == 0 {
+        return None;
+    }
+
+    Some(((ds * 3600 * 100) / dt) as _)
+}
+
+async fn send_error_message(
+    writer: &mut tokio::net::tcp::WriteHalf<'_>,
+    message: &str,
+) -> Result<()> {
+    let mut to_wire = vec![0x10];
+
+    ensure!(message.len() <= 255, "message too long");
+    to_wire.push(message.len() as u8);
+
+    for byte in message.as_bytes() {
+        to_wire.push(*byte);
+    }
+
+    writer.write_all(&to_wire).await?;
+    Ok(())
+}
+
+async fn read_string(r: &mut (impl AsyncReadExt + Unpin)) -> Result<String> {
+    let len = r.read_u8().await?;
+    let mut buf = vec![0u8; len as usize];
+    r.read_exact(&mut buf).await?;
+    Ok(String::from_utf8(buf)?)
+}
+
+#[derive(Clone, Debug)]
+struct Camera {
+    road: u16,
+    mile: u16,
+    limit: u16,
+}
+
+#[derive(Clone, Debug)]
+struct Ticket {
+    plate: String,
+    road: u16,
+    mile1: u16,
+    timestamp1: u32,
+    mile2: u16,
+    timestamp2: u32,
+    speed: u16,
+}
+
+impl Ticket {
+    fn new(
+        plate: String,
+        road: u16,
+        mut mile1: u16,
+        mut timestamp1: u32,
+        mut mile2: u16,
+        mut timestamp2: u32,
+        speed: u16,
+    ) -> Self {
+        if timestamp2 < timestamp1 {
+            (mile1, mile2) = (mile2, mile1);
+            (timestamp1, timestamp2) = (timestamp2, timestamp1);
+        }
+        Self {
+            plate,
+            road,
+            mile1,
+            timestamp1,
+            mile2,
+            timestamp2,
+            speed,
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+#[non_exhaustive]
+enum ChannelMessage {
+    TicketIssued,
+}