refactoring

1 files changed, 56 insertions, 3 deletions

diff --git a/common/src/lib.rs b/common/src/lib.rs
index d0db5d1..616dfcc 100644
--- a/common/src/lib.rs
+++ b/common/src/lib.rs
@@ -1,7 +1,8 @@
-use std::io::{Stdout, Write};
+use std::io::{Read, Write};
 
-use msg::{MaelstromBodyOr, Message, MessageHeader};
+use msg::{MaelstromBody, MaelstromBodyOr, Message, MessageHeader};
 use serde::{Deserialize, Serialize};
+use serde_json::Deserializer;
 
 pub mod msg;
 
@@ -13,7 +14,7 @@ pub trait Handler {
         &mut self,
         header: MessageHeader,
         body: Self::Body,
-        writer: &mut MsgWriter<Self::Body>,
+        writer: &mut MsgWriter<impl Write>,
     ) -> ();
 }
 
@@ -39,3 +40,55 @@ impl<W: Write> MsgWriter<W> {
         self.writer.write(&[b'\n']).unwrap();
     }
 }
+
+pub fn run_with<T: Handler>(mut reader: impl Read, mut writer: impl Write) {
+    let (mut handler, mut msg_writer) = init_handler::<T, _, _>(&mut reader, &mut writer);
+
+    let deser = Deserializer::from_reader(reader);
+    for msg in deser.into_iter::<Message<T::Body>>() {
+        let msg = msg.unwrap();
+        match msg.body {
+            MaelstromBodyOr::Other { inner } => {
+                handler.handle(msg.header, inner, &mut msg_writer);
+            }
+            _ => todo!(),
+        };
+    }
+}
+
+pub fn init_handler<T: Handler, R: Read, W: Write>(reader: R, writer: W) -> (T, MsgWriter<W>) {
+    let deser = Deserializer::from_reader(reader);
+    let mut deser = deser.into_iter::<Message<()>>();
+    let Some(msg) = deser.next() else {
+        panic!("stream ended before init message");
+    };
+    let Ok(msg) = msg else {
+            panic!("{}", msg.unwrap_err());
+    };
+
+    let (node_id, node_ids, msg_id) = match msg.body {
+        MaelstromBodyOr::MaelstromBody {
+            inner:
+                MaelstromBody::Init {
+                    node_id,
+                    node_ids,
+                    msg_id,
+                },
+        } => (node_id, node_ids, msg_id),
+        _ => {
+            panic!("expected init message to be first message");
+        }
+    };
+
+    let mut writer = MsgWriter::new(node_id.clone(), writer);
+
+    writer.write(
+        msg.header.src,
+        &MaelstromBody::InitOk {
+            msg_id: 0,
+            in_reply_to: msg_id,
+        },
+    );
+
+    (T::init(node_id, node_ids, msg_id), writer)
+}