1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
|
use std::io::{Read, Write};
use msg::{MaelstromBody, MaelstromBodyOr, Message, MessageHeader};
use serde::{Deserialize, Serialize};
use serde_json::Deserializer;
pub mod msg;
pub trait Handler {
type Body: Serialize + for<'a> Deserialize<'a>;
fn init(node_id: String, node_ids: Vec<String>, msg_id: usize) -> Self;
fn handle(
&mut self,
header: MessageHeader,
body: Self::Body,
writer: &mut MsgWriter<impl Write>,
) -> ();
}
pub struct MsgWriter<W> {
node_id: String,
writer: W,
}
impl<W: Write> MsgWriter<W> {
pub fn new(node_id: String, writer: W) -> Self {
Self { node_id, writer }
}
pub fn write<T: Serialize>(&mut self, dst: String, msg: &T) {
let msg = Message {
header: MessageHeader {
src: self.node_id.clone(),
dst,
},
body: MaelstromBodyOr::Other { inner: msg },
};
serde_json::to_writer(&mut self.writer, &msg).unwrap();
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)
}
|
