1 files changed, 33 insertions, 40 deletions

diff --git a/src/play.rs b/src/play.rs
index c635595..2f364c0 100644
--- a/src/play.rs
+++ b/src/play.rs
@@ -1,9 +1,8 @@
 use crate::AppState;
-use crate::session::{HandObject, Session};
+use crate::session::{HandObject, PlayerColor, Session};
 use axum::extract::ws::{Message, Utf8Bytes, WebSocket};
 use futures_util::{SinkExt, StreamExt};
 use serde::{Deserialize, Serialize};
-use std::collections::HashMap;
 use std::sync::{Arc, Mutex, RwLock, Weak};
 use tokio::sync::broadcast::Receiver;
 use tokio::sync::broadcast::error::RecvError;
@@ -27,7 +26,7 @@ enum OutgoingPlayMessage {
 
 #[derive(Clone)]
 pub enum PlayUpdate {
-    HandUpdate(HashMap<String, Vec<HandObject>>),
+    HandUpdate([Vec<HandObject>; PlayerColor::COUNT]),
 }
 
 pub async fn handle_play(socket: WebSocket, app_state: Arc<AppState>) {
@@ -57,7 +56,7 @@ pub async fn handle_play(socket: WebSocket, app_state: Arc<AppState>) {
         let sender_tx = sender_tx.clone();
         tokio::spawn(async move {
             let mut player_session = None;
-            let player_color = Arc::new(RwLock::new(String::new()));
+            let player_color = Arc::new(RwLock::new(PlayerColor::Grey));
 
             while let Some(msg) = receiver.next().await {
                 let Ok(Message::Text(text)) = msg else {
@@ -70,7 +69,7 @@ pub async fn handle_play(socket: WebSocket, app_state: Arc<AppState>) {
                             let sessions = app_state.sessions.read().unwrap();
                             sessions
                                 .get(&id)
-                                .map(Arc::to_owned)
+                                .map(Arc::clone)
                                 .ok_or("Session did not exist")
                         };
 
@@ -79,8 +78,11 @@ pub async fn handle_play(socket: WebSocket, app_state: Arc<AppState>) {
                                 let (colors, update_rx) = {
                                     let session = session.lock().unwrap();
 
-                                    let colors: Vec<String> =
-                                        session.seats.keys().cloned().collect();
+                                    let colors: Vec<String> = session.seats.iter().enumerate()
+                                        .filter(|(_, hand)| hand.len() > 0)
+                                        .flat_map(|(index, _)| PlayerColor::try_from(index).ok())
+                                        .map(|color| String::from(color.as_ref()))
+                                        .collect();
                                     let update_rx = session.update_tx.subscribe();
 
                                     (colors, update_rx)
@@ -121,35 +123,23 @@ pub async fn handle_play(socket: WebSocket, app_state: Arc<AppState>) {
                         let Some(session) =
                             player_session.clone().and_then(|session| session.upgrade())
                         else {
-                            let response = OutgoingPlayMessage::Error;
-                            if sender_tx.send(response).await.is_err() {
-                                break;
-                            }
-                            break;
+                            let _ = sender_tx.send(OutgoingPlayMessage::Error).await;
+                            break
                         };
-                        let hand = session
-                            .lock()
-                            .unwrap()
-                            .seats
-                            .get(&color)
-                            .map(|seat| seat.to_owned());
-                        match hand {
-                            Some(hand) => {
-                                *player_color.write().unwrap() = color;
-                                if sender_tx
-                                    .send(OutgoingPlayMessage::Hand(hand))
-                                    .await
-                                    .is_err()
-                                {
-                                    break;
-                                }
-                            }
-                            None => {
-                                if sender_tx.send(OutgoingPlayMessage::Error).await.is_err() {
-                                    break;
-                                }
-                            }
+                        let Ok(color) = PlayerColor::try_from(color.as_str()) else {
+                            let _ = sender_tx.send(OutgoingPlayMessage::Error).await;
+                            break
                         };
+
+                        let hand = session.lock().unwrap().seats[&color].clone();
+                        *player_color.write().unwrap() = color;
+                        if sender_tx
+                            .send(OutgoingPlayMessage::Hand(hand))
+                            .await
+                            .is_err()
+                        {
+                            break;
+                        }
                     }
                     Err(err) => {
                         eprintln!(
@@ -173,23 +163,26 @@ async fn handle_update(
     mut update_rx: Receiver<PlayUpdate>,
     sender_tx: Sender<OutgoingPlayMessage>,
     _player_session: Weak<Mutex<Session>>,
-    player_color: Arc<RwLock<String>>,
+    player_color: Arc<RwLock<PlayerColor>>,
 ) {
     loop {
         match update_rx.recv().await {
             Ok(PlayUpdate::HandUpdate(hands)) => {
+                let colors: Vec<String> = hands.iter().enumerate()
+                    .filter(|(_, hand)| hand.len() > 0)
+                    .flat_map(|(index, _)| PlayerColor::try_from(index).ok())
+                    .map(|color| String::from(color.as_ref()))
+                    .collect();
                 let _ = sender_tx
                     .send(OutgoingPlayMessage::Initialize {
-                        colors: hands.keys().cloned().collect(),
+                        colors,
                     })
                     .await;
                 let hand = {
                     let color = player_color.read().unwrap();
-                    hands.get(&*color).map(ToOwned::to_owned)
-                };
-                if let Some(hand) = hand {
-                    let _ = sender_tx.send(OutgoingPlayMessage::Hand(hand)).await;
+                    hands[usize::from(&*color)].to_owned()
                 };
+                let _ = sender_tx.send(OutgoingPlayMessage::Hand(hand)).await;
             }
             Err(RecvError::Closed) => break,
             Err(RecvError::Lagged(_)) => continue,