refactoring

2 files changed, 77 insertions, 75 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)
+}
diff --git a/echo/src/main.rs b/echo/src/main.rs
index 381099c..5ae72f3 100644
--- a/echo/src/main.rs
+++ b/echo/src/main.rs
@@ -1,8 +1,11 @@
-use std::io;
+use std::io::{self, Write};
 
-use common::{msg::*, send_msg, Handler};
+use common::{msg::*, run_with, Handler, MsgWriter};
 use serde::{Deserialize, Serialize};
-use serde_json::Deserializer;
+
+fn main() {
+    run_with::<EchoHandler>(io::stdin(), io::stdout());
+}
 
 #[derive(Debug, Serialize, Deserialize)]
 #[serde(tag = "type")]
@@ -17,61 +20,6 @@ pub enum EchoBody {
         echo: String,
     },
 }
-fn main() {
-    let mut stdout = io::stdout();
-    let deser = Deserializer::from_reader(io::stdin());
-    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");
-        }
-    };
-
-    send_msg(
-        &mut stdout,
-        &Message {
-            header: msg.header.flip(),
-            body: MaelstromBodyOr::MaelstromBody::<()> {
-                inner: MaelstromBody::InitOk {
-                    msg_id: 0,
-                    in_reply_to: msg_id,
-                },
-            },
-        },
-    );
-
-    let mut handler = EchoHandler::init(node_id, node_ids, msg_id);
-
-    drop(deser);
-
-    let deser = Deserializer::from_reader(io::stdin());
-    for msg in deser.into_iter::<Message<EchoBody>>() {
-        let msg = msg.unwrap();
-        match msg.body {
-            MaelstromBodyOr::Other { inner } => {
-                if let Some(out) = handler.handle(msg.header, inner) {
-                    send_msg(&mut stdout, &out);
-                }
-            }
-            _ => todo!(),
-        };
-    }
-}
 
 pub struct EchoHandler {
     next_msg_id: usize,
@@ -84,25 +32,26 @@ impl Handler for EchoHandler {
         EchoHandler { next_msg_id: 1 }
     }
 
-    fn handle(&mut self, header: MessageHeader, body: Self::Body) -> Option<Message<EchoBody>> {
+    fn handle(
+        &mut self,
+        header: MessageHeader,
+        body: Self::Body,
+        writer: &mut MsgWriter<impl Write>,
+    ) {
         match body {
             EchoBody::Echo { msg_id, echo } => {
-                let msg = Message {
-                    header: header.flip(),
-                    body: MaelstromBodyOr::Other {
-                        inner: EchoBody::EchoOk {
-                            msg_id: self.next_msg_id,
-                            in_reply_to: msg_id,
-                            echo,
-                        },
+                writer.write(
+                    header.src,
+                    &EchoBody::EchoOk {
+                        msg_id: self.next_msg_id,
+                        in_reply_to: msg_id,
+                        echo,
                     },
-                };
+                );
 
                 self.next_msg_id += 1;
-
-                Some(msg)
             }
-            EchoBody::EchoOk { .. } => None,
-        }
+            EchoBody::EchoOk { .. } => (),
+        };
     }
 }