1 files changed, 21 insertions, 72 deletions
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 { .. } => (),
+        };
     }
 }