1use std::io;
42use std::net::{ToSocketAddrs, UdpSocket};
43use std::time::{Duration, Instant};
44
45use crate::{EventStream, EventStreamBuilder};
46
47const MAX_PAYLOAD_BYTES: usize = 1400;
51
52const MAX_WORDS: usize = MAX_PAYLOAD_BYTES / 4;
54
55const NO_TIMESTAMP_FLAG: u32 = 0x8000_0000;
57
58const POLARITY_FLAG: u32 = 0x0000_8000;
60
61const COORD_MASK: u32 = 0x7FFF;
63
64#[derive(Clone, Copy, Debug, PartialEq, Eq)]
66pub struct WireFormat {
67 pub timestamps: bool,
70 pub host_endian: bool,
72}
73
74impl Default for WireFormat {
75 fn default() -> Self {
76 Self {
77 timestamps: false,
78 host_endian: true,
79 }
80 }
81}
82
83impl WireFormat {
84 pub fn host_endian() -> Self {
86 Self::default()
87 }
88
89 pub fn network_endian() -> Self {
91 Self {
92 timestamps: false,
93 host_endian: false,
94 }
95 }
96
97 pub fn with_timestamps(mut self) -> Self {
99 self.timestamps = true;
100 self
101 }
102
103 fn encode_word(self, word: u32) -> [u8; 4] {
104 if self.host_endian {
105 word.to_ne_bytes()
106 } else {
107 word.to_be_bytes()
108 }
109 }
110
111 fn decode_word(self, bytes: [u8; 4]) -> u32 {
112 if self.host_endian {
113 u32::from_ne_bytes(bytes)
114 } else {
115 u32::from_be_bytes(bytes)
116 }
117 }
118}
119
120fn encode_event(x: u16, y: u16, polarity: bool, timestamped: bool) -> u32 {
122 let x = u32::from(x) & COORD_MASK;
123 let y = u32::from(y) & COORD_MASK;
124 let high = if timestamped { x } else { x | 0x8000 };
126 let low = if polarity { y | POLARITY_FLAG } else { y };
127 (high << 16) | low
128}
129
130fn decode_event(word: u32) -> (u16, u16, bool) {
132 let x = ((word >> 16) & COORD_MASK) as u16;
133 let y = (word & COORD_MASK) as u16;
134 let polarity = word & POLARITY_FLAG != 0;
135 (x, y, polarity)
136}
137
138pub struct UdpSender {
140 socket: UdpSocket,
141 format: WireFormat,
142}
143
144impl UdpSender {
145 pub fn connect(target: impl ToSocketAddrs, format: WireFormat) -> io::Result<Self> {
147 let socket = UdpSocket::bind("0.0.0.0:0")?;
150 socket.connect(target)?;
151 Ok(Self { socket, format })
152 }
153
154 pub fn local_addr(&self) -> io::Result<std::net::SocketAddr> {
156 self.socket.local_addr()
157 }
158
159 pub fn send(&self, stream: &EventStream) -> io::Result<usize> {
170 let words_per_event = if self.format.timestamps { 2 } else { 1 };
171 let events_per_packet = MAX_WORDS / words_per_event;
172 if events_per_packet == 0 {
173 return Err(io::Error::new(
174 io::ErrorKind::InvalidInput,
175 "wire format does not fit in a datagram",
176 ));
177 }
178
179 let (xs, ys, ts, ps) = (stream.xs(), stream.ys(), stream.ts(), stream.ps());
180 let mut buffer = Vec::with_capacity(MAX_PAYLOAD_BYTES);
181 let mut sent = 0;
182
183 for chunk in (0..stream.len())
184 .collect::<Vec<_>>()
185 .chunks(events_per_packet)
186 {
187 buffer.clear();
188 for &index in chunk {
189 let word = encode_event(xs[index], ys[index], ps[index], self.format.timestamps);
190 buffer.extend_from_slice(&self.format.encode_word(word));
191 if self.format.timestamps {
192 buffer.extend_from_slice(&self.format.encode_word(ts[index] as u32));
196 }
197 }
198 if !buffer.is_empty() {
199 self.socket.send(&buffer)?;
200 sent += chunk.len();
201 }
202 }
203 Ok(sent)
204 }
205}
206
207pub struct UdpReceiver {
209 socket: UdpSocket,
210 width: usize,
211 height: usize,
212 format: WireFormat,
213}
214
215impl UdpReceiver {
216 pub fn bind(
221 address: impl ToSocketAddrs,
222 width: usize,
223 height: usize,
224 format: WireFormat,
225 ) -> io::Result<Self> {
226 let socket = UdpSocket::bind(address)?;
227 Ok(Self {
228 socket,
229 width,
230 height,
231 format,
232 })
233 }
234
235 pub fn local_addr(&self) -> io::Result<std::net::SocketAddr> {
236 self.socket.local_addr()
237 }
238
239 pub fn recv_window(&self, window: Duration) -> io::Result<EventStream> {
244 let deadline = Instant::now() + window;
245 let mut builder = EventStreamBuilder::new(self.width, self.height, 0.001);
246 let mut packet = vec![0_u8; MAX_PAYLOAD_BYTES * 2];
247 let mut received = 0_i64;
248
249 loop {
250 let remaining = deadline.saturating_duration_since(Instant::now());
251 if remaining.is_zero() {
252 break;
253 }
254 self.socket.set_read_timeout(Some(remaining))?;
256 match self.socket.recv(&mut packet) {
257 Ok(size) => {
258 received += self.decode_into(&packet[..size], &mut builder, received);
259 }
260 Err(error)
261 if matches!(
262 error.kind(),
263 io::ErrorKind::WouldBlock | io::ErrorKind::TimedOut
264 ) =>
265 {
266 break
267 }
268 Err(error) => return Err(error),
269 }
270 }
271 Ok(builder.build())
272 }
273
274 fn decode_into(
276 &self,
277 payload: &[u8],
278 builder: &mut EventStreamBuilder,
279 arrival_index: i64,
280 ) -> i64 {
281 let mut words = payload
282 .as_chunks::<4>()
283 .0
284 .iter()
285 .map(|bytes| self.format.decode_word(*bytes));
286 let mut count = 0;
287 while let Some(word) = words.next() {
288 let timestamped = word & NO_TIMESTAMP_FLAG == 0;
291 let (x, y, polarity) = decode_event(word);
292 let timestamp = if timestamped {
293 match words.next() {
294 Some(t) => i64::from(t),
295 None => break,
297 }
298 } else {
299 arrival_index + count
301 };
302 builder.push(x, y, timestamp, polarity);
303 count += 1;
304 }
305 count
306 }
307}
308
309#[cfg(test)]
310mod tests {
311 use super::*;
312
313 fn sample(width: usize, height: usize, count: usize) -> EventStream {
314 let mut builder = EventStreamBuilder::new(width, height, 0.001);
315 for index in 0..count {
316 builder.push(
317 (index % width) as u16,
318 (index % height) as u16,
319 index as i64 * 100,
320 index % 2 == 0,
321 );
322 }
323 builder.build()
324 }
325
326 #[test]
327 fn encoding_round_trips_through_the_word() {
328 for (x, y, polarity) in [
329 (0, 0, false),
330 (1, 2, true),
331 (639, 479, true),
332 (0x7FFF, 0x7FFF, false),
333 ] {
334 for timestamped in [false, true] {
335 let word = encode_event(x, y, polarity, timestamped);
336 assert_eq!(decode_event(word), (x, y, polarity), "{x},{y},{polarity}");
337 assert_eq!(
339 word & NO_TIMESTAMP_FLAG == 0,
340 timestamped,
341 "mode flag wrong for timestamped={timestamped}"
342 );
343 }
344 }
345 }
346
347 #[test]
348 fn the_layout_matches_aestreams() {
349 assert_eq!(
353 encode_event(3, 5, true, false),
354 ((3 | 0x8000) << 16) | (5 | 0x8000)
355 );
356 assert_eq!(encode_event(3, 5, false, false), ((3 | 0x8000) << 16) | 5);
357 assert_eq!(encode_event(3, 5, true, true), (3 << 16) | (5 | 0x8000));
358 assert_eq!(encode_event(3, 5, false, true), (3 << 16) | 5);
359 }
360
361 #[test]
362 fn byte_order_is_selectable() {
363 let host = WireFormat::host_endian();
364 let network = WireFormat::network_endian();
365 let word = 0x1234_5678_u32;
366 assert_eq!(network.encode_word(word), word.to_be_bytes());
367 assert_eq!(host.encode_word(word), word.to_ne_bytes());
368 assert_eq!(host.decode_word(host.encode_word(word)), word);
370 assert_eq!(network.decode_word(network.encode_word(word)), word);
371 }
372
373 #[test]
374 fn a_stream_round_trips_over_loopback() {
375 let format = WireFormat::default();
376 let receiver = UdpReceiver::bind("127.0.0.1:0", 64, 64, format).unwrap();
377 let address = receiver.local_addr().unwrap();
378 let sender = UdpSender::connect(address, format).unwrap();
379
380 let original = sample(64, 64, 500);
381 let sent = sender.send(&original).unwrap();
382 assert_eq!(sent, original.len());
383
384 let received = receiver.recv_window(Duration::from_millis(300)).unwrap();
385 assert_eq!(
386 received.len(),
387 original.len(),
388 "every event should arrive on loopback"
389 );
390 assert_eq!(received.xs(), original.xs());
391 assert_eq!(received.ys(), original.ys());
392 assert_eq!(received.ps(), original.ps());
393 }
394
395 #[test]
396 fn timestamps_survive_when_the_format_carries_them() {
397 let format = WireFormat::default().with_timestamps();
398 let receiver = UdpReceiver::bind("127.0.0.1:0", 64, 64, format).unwrap();
399 let sender = UdpSender::connect(receiver.local_addr().unwrap(), format).unwrap();
400
401 let original = sample(64, 64, 200);
402 sender.send(&original).unwrap();
403 let received = receiver.recv_window(Duration::from_millis(300)).unwrap();
404 assert_eq!(received.len(), original.len());
405 assert_eq!(
406 received.ts(),
407 original.ts(),
408 "timestamps must survive the wire"
409 );
410 }
411
412 #[test]
413 fn a_large_stream_is_split_across_datagrams() {
414 let format = WireFormat::default();
420 let receiver = UdpReceiver::bind("127.0.0.1:0", 128, 128, format).unwrap();
421 let address = receiver.local_addr().unwrap();
422
423 let count = MAX_WORDS * 8;
424 let handle = std::thread::spawn(move || {
425 let mut total = 0;
426 let deadline = Instant::now() + Duration::from_secs(2);
427 while Instant::now() < deadline {
428 let received = receiver
429 .recv_window(Duration::from_millis(100))
430 .expect("receive should not error");
431 if received.is_empty() && total > 0 {
432 break;
433 }
434 total += received.len();
435 if total >= count {
436 break;
437 }
438 }
439 total
440 });
441
442 std::thread::sleep(Duration::from_millis(50));
444 let sender = UdpSender::connect(address, format).unwrap();
445 let original = sample(128, 128, count);
446 assert_eq!(sender.send(&original).unwrap(), count);
447
448 let received = handle.join().expect("receiver thread should not panic");
449 assert!(
452 received > count / 2,
453 "received {received} of {count} with a concurrent receiver"
454 );
455 }
456
457 #[test]
458 fn a_quiet_link_returns_an_empty_stream() {
459 let receiver = UdpReceiver::bind("127.0.0.1:0", 32, 32, WireFormat::default()).unwrap();
460 let events = receiver.recv_window(Duration::from_millis(20)).unwrap();
461 assert!(events.is_empty(), "silence is not an error");
462 }
463
464 #[test]
465 fn a_truncated_datagram_does_not_produce_a_bogus_event() {
466 let format = WireFormat::default().with_timestamps();
469 let receiver = UdpReceiver::bind("127.0.0.1:0", 32, 32, format).unwrap();
470 let mut builder = EventStreamBuilder::new(32, 32, 0.001);
471 let word = encode_event(4, 4, true, true);
472 let truncated = format.encode_word(word); assert_eq!(receiver.decode_into(&truncated, &mut builder, 0), 0);
474 assert!(builder.build().is_empty());
475 }
476
477 #[test]
478 fn coordinates_outside_the_sensor_are_dropped() {
479 let receiver = UdpReceiver::bind("127.0.0.1:0", 16, 16, WireFormat::default()).unwrap();
480 let mut builder = EventStreamBuilder::new(16, 16, 0.001);
481 let mut payload = Vec::new();
482 for word in [
483 encode_event(4, 4, true, false), encode_event(900, 900, true, false), ] {
486 payload.extend_from_slice(&receiver.format.encode_word(word));
487 }
488 receiver.decode_into(&payload, &mut builder, 0);
489 assert_eq!(
490 builder.build().len(),
491 1,
492 "only the in-bounds event survives"
493 );
494 }
495}