1use std::sync::{Arc, Mutex, PoisonError};
2
3use crate::command::Command;
4use crate::error::Error;
5use crate::mirror::FirmwareState;
6use crate::operation::{Distribution, Nop, Operation};
7use crate::protocol::{Cmd, PAYLOAD_BYTES};
8
9#[derive(Clone, Debug)]
10pub(crate) enum Mirror {
11 Synced(Vec<FirmwareState>),
12 Desynced,
13}
14
15#[derive(Clone)]
16pub(crate) struct MirrorHandle {
17 pub(crate) state: Arc<Mutex<Mirror>>,
18 pub(crate) enabled: bool,
19}
20
21#[derive(Clone, Debug)]
22pub struct Datagram {
23 pub cmd: Cmd,
24 pub payload: [u8; PAYLOAD_BYTES],
25}
26
27impl Datagram {
28 #[must_use]
29 pub const fn no_payload(cmd: Cmd) -> Self {
30 Self {
31 cmd,
32 payload: [0u8; PAYLOAD_BYTES],
33 }
34 }
35}
36
37#[derive(Clone, Copy, Debug)]
38pub struct Frame<'a> {
39 dist: Distribution,
40 datagrams: &'a [Datagram],
41}
42
43impl<'a> Frame<'a> {
44 #[must_use]
45 pub fn distribution(&self) -> Distribution {
46 self.dist
47 }
48
49 #[must_use]
50 pub fn datagrams(&self) -> &'a [Datagram] {
51 self.datagrams
52 }
53}
54
55#[derive(Debug)]
56struct FrameDesc {
57 dist: Distribution,
58 start: usize,
59 len: usize,
60}
61
62#[derive(Debug, Default)]
63pub struct Frames {
64 payloads: Vec<Datagram>,
65 frames: Vec<FrameDesc>,
66}
67
68impl Frames {
69 #[must_use]
70 pub fn len(&self) -> usize {
71 self.frames.len()
72 }
73
74 #[must_use]
75 pub fn is_empty(&self) -> bool {
76 self.frames.is_empty()
77 }
78
79 #[must_use]
80 pub fn frame(&self, index: usize) -> Option<Frame<'_>> {
81 self.frames.get(index).map(|desc| Frame {
82 dist: desc.dist,
83 datagrams: &self.payloads[desc.start..desc.start + desc.len],
84 })
85 }
86
87 #[must_use]
88 pub fn iter(&self) -> FrameIter<'_> {
89 FrameIter {
90 frames: self,
91 index: 0,
92 }
93 }
94
95 fn clear(&mut self) {
96 self.payloads.clear();
97 self.frames.clear();
98 }
99
100 pub(crate) fn push_op<O: Operation + ?Sized>(
101 &mut self,
102 op: &O,
103 num_devices: usize,
104 ) -> Result<(), Error> {
105 let dist = op.distribution();
106 let encode_devices = match dist {
107 Distribution::Broadcast => 1,
108 Distribution::PerDevice => num_devices,
109 };
110 for frame in 0..op.frames() {
111 let start = self.payloads.len();
112 for device in 0..encode_devices {
113 let mut payload = [0u8; PAYLOAD_BYTES];
114 let cmd = op.encode(device, frame, &mut payload)?;
115 self.payloads.push(Datagram { cmd, payload });
116 }
117 self.frames.push(FrameDesc {
118 dist,
119 start,
120 len: encode_devices,
121 });
122 }
123 Ok(())
124 }
125
126 fn push_each_step(
127 &mut self,
128 devices: &[Vec<Box<dyn Operation + '_>>],
129 num_devices: usize,
130 ) -> Result<(), Error> {
131 let slot_frames = each_slot_frames(devices);
132 let total: usize = slot_frames.iter().sum();
133 for frame in 0..total {
134 let start = self.payloads.len();
135 for device in 0..num_devices {
136 let mut payload = [0u8; PAYLOAD_BYTES];
137 let cmd = each_encode(devices, &slot_frames, device, frame, &mut payload)?;
138 self.payloads.push(Datagram { cmd, payload });
139 }
140 self.frames.push(FrameDesc {
141 dist: Distribution::PerDevice,
142 start,
143 len: num_devices,
144 });
145 }
146 Ok(())
147 }
148}
149
150pub struct FrameIter<'a> {
151 frames: &'a Frames,
152 index: usize,
153}
154
155impl<'a> Iterator for FrameIter<'a> {
156 type Item = Frame<'a>;
157
158 fn next(&mut self) -> Option<Frame<'a>> {
159 let frame = self.frames.frame(self.index)?;
160 self.index += 1;
161 Some(frame)
162 }
163}
164
165impl<'a> IntoIterator for &'a Frames {
166 type Item = Frame<'a>;
167 type IntoIter = FrameIter<'a>;
168
169 fn into_iter(self) -> FrameIter<'a> {
170 self.iter()
171 }
172}
173
174enum Step<'a> {
175 Op(Box<dyn Operation + 'a>),
176 Each {
177 devices: Vec<Vec<Box<dyn Operation + 'a>>>,
178 },
179}
180
181pub struct DatagramBuilder<'a> {
182 num_devices: usize,
183 ops: Vec<Step<'a>>,
184 mirror: Option<MirrorHandle>,
185}
186
187impl<'a> DatagramBuilder<'a> {
188 #[must_use]
189 pub fn new(num_devices: usize) -> Self {
190 Self {
191 num_devices,
192 ops: Vec::new(),
193 mirror: None,
194 }
195 }
196
197 #[must_use]
198 pub(crate) fn with_mirror(num_devices: usize, mirror: MirrorHandle) -> Self {
199 Self {
200 num_devices,
201 ops: Vec::new(),
202 mirror: Some(mirror),
203 }
204 }
205
206 pub fn push<C: Command<'a>>(&mut self, cmd: C) -> &mut Self {
207 cmd.expand(self);
208 self
209 }
210
211 pub fn push_each<C, F>(&mut self, mut assign: F) -> &mut Self
212 where
213 C: Command<'a>,
214 F: FnMut(usize) -> Option<C>,
215 {
216 let num_devices = self.num_devices;
217 let mut new_devices: Vec<Vec<Box<dyn Operation + 'a>>> = Vec::with_capacity(num_devices);
218 for device in 0..num_devices {
219 match assign(device) {
220 Some(cmd) => {
221 let mut sub = DatagramBuilder::new(num_devices);
222 cmd.expand(&mut sub);
223 new_devices.push(sub.take_ops());
224 }
225 None => new_devices.push(Vec::new()),
226 }
227 }
228
229 let fuse = matches!(
230 self.ops.last(),
231 Some(Step::Each { devices }) if (0..num_devices)
232 .all(|d| devices[d].is_empty() || new_devices[d].is_empty())
233 );
234 if fuse {
235 if let Some(Step::Each { devices }) = self.ops.last_mut() {
236 for (device, ops) in new_devices.into_iter().enumerate() {
237 if !ops.is_empty() {
238 devices[device] = ops;
239 }
240 }
241 }
242 } else {
243 self.ops.push(Step::Each {
244 devices: new_devices,
245 });
246 }
247 self
248 }
249
250 pub(crate) fn push_op<O: Operation + 'a>(&mut self, op: O) -> &mut Self {
251 self.ops.push(Step::Op(Box::new(op)));
252 self
253 }
254
255 pub(crate) fn take_ops(self) -> Vec<Box<dyn Operation + 'a>> {
256 self.ops
257 .into_iter()
258 .map(|step| match step {
259 Step::Op(op) => op,
260 Step::Each { devices } => {
261 let slot_frames = each_slot_frames(&devices);
262 Box::new(EachOwned {
263 devices,
264 slot_frames,
265 }) as Box<dyn Operation + 'a>
266 }
267 })
268 .collect()
269 }
270
271 pub fn build(&self) -> Result<Frames, Error> {
272 let mut out = Frames::default();
273 self.build_into(&mut out)?;
274 Ok(out)
275 }
276
277 pub fn build_into(&self, out: &mut Frames) -> Result<(), Error> {
278 out.clear();
279
280 let mut guard = self
281 .mirror
282 .as_ref()
283 .filter(|handle| handle.enabled)
284 .map(|handle| handle.state.lock().unwrap_or_else(PoisonError::into_inner));
285
286 let mut work = match guard.as_deref() {
287 Some(Mirror::Synced(states)) => Some(states.clone()),
288 _ => None,
289 };
290
291 for step in &self.ops {
292 match step {
293 Step::Op(op) => {
294 out.push_op(op.as_ref(), self.num_devices)?;
295 if let Some(work) = work.as_mut() {
296 for (device, state) in work.iter_mut().enumerate() {
297 op.reflect(device, state)?;
298 }
299 }
300 }
301 Step::Each { devices } => {
302 out.push_each_step(devices, self.num_devices)?;
303 if let Some(work) = work.as_mut() {
304 for (device, state) in work.iter_mut().enumerate() {
305 each_reflect(devices, device, state)?;
306 }
307 }
308 }
309 }
310 }
311
312 if let (Some(guard), Some(work)) = (guard.as_mut(), work) {
313 **guard = Mirror::Synced(work);
314 }
315 Ok(())
316 }
317}
318
319fn each_slot_frames(devices: &[Vec<Box<dyn Operation + '_>>]) -> Vec<usize> {
320 let num_slots = devices.iter().map(Vec::len).max().unwrap_or(0);
321 let mut slot_frames = vec![0usize; num_slots];
322 for ops in devices {
323 for (slot, op) in ops.iter().enumerate() {
324 slot_frames[slot] = slot_frames[slot].max(op.frames());
325 }
326 }
327 slot_frames
328}
329
330fn each_locate(slot_frames: &[usize], frame: usize) -> Option<(usize, usize)> {
331 let mut remaining = frame;
332 for (slot, &frames) in slot_frames.iter().enumerate() {
333 if remaining < frames {
334 return Some((slot, remaining));
335 }
336 remaining -= frames;
337 }
338 None
339}
340
341fn each_encode(
342 devices: &[Vec<Box<dyn Operation + '_>>],
343 slot_frames: &[usize],
344 device: usize,
345 frame: usize,
346 out: &mut [u8; PAYLOAD_BYTES],
347) -> Result<Cmd, Error> {
348 if let Some((slot, subframe)) = each_locate(slot_frames, frame) {
349 if let Some(op) = devices.get(device).and_then(|ops| ops.get(slot))
350 && subframe < op.frames()
351 {
352 return op.encode(device, subframe, out);
353 }
354 return Nop.encode(device, subframe, out);
355 }
356 Nop.encode(device, frame, out)
357}
358
359fn each_reflect(
360 devices: &[Vec<Box<dyn Operation + '_>>],
361 device: usize,
362 state: &mut FirmwareState,
363) -> Result<(), Error> {
364 if let Some(ops) = devices.get(device) {
365 for op in ops {
366 op.reflect(device, state)?;
367 }
368 }
369 Ok(())
370}
371
372struct EachOwned<'a> {
373 devices: Vec<Vec<Box<dyn Operation + 'a>>>,
374 slot_frames: Vec<usize>,
375}
376
377impl Operation for EachOwned<'_> {
378 fn frames(&self) -> usize {
379 self.slot_frames.iter().sum()
380 }
381
382 fn distribution(&self) -> Distribution {
383 Distribution::PerDevice
384 }
385
386 fn encode(
387 &self,
388 device: usize,
389 frame: usize,
390 out: &mut [u8; PAYLOAD_BYTES],
391 ) -> Result<Cmd, Error> {
392 each_encode(&self.devices, &self.slot_frames, device, frame, out)
393 }
394
395 fn reflect(&self, device: usize, state: &mut FirmwareState) -> Result<(), Error> {
396 each_reflect(&self.devices, device, state)
397 }
398}
399
400#[cfg(test)]
401mod tests {
402 use super::*;
403 use crate::command::Pattern;
404 use crate::operation::{ConfigModulation, ConfigPattern, WritePatternBuffer};
405 use crate::value::{LoopBehavior, ModulationBank, SamplingConfig};
406
407 #[derive(Clone, Copy)]
408 struct Multi(usize);
409
410 impl Operation for Multi {
411 fn frames(&self) -> usize {
412 self.0
413 }
414
415 fn distribution(&self) -> Distribution {
416 Distribution::PerDevice
417 }
418
419 fn encode(
420 &self,
421 _device: usize,
422 frame: usize,
423 out: &mut [u8; PAYLOAD_BYTES],
424 ) -> Result<Cmd, Error> {
425 out[0] = u8::try_from(frame).unwrap();
426 Ok(Cmd::ConfigModulation)
427 }
428 }
429
430 fn cmd_at(frames: &Frames, frame: usize, device: usize) -> Cmd {
431 frames.frame(frame).unwrap().datagrams()[device].cmd
432 }
433
434 #[test]
435 fn push_each_routes_per_device() {
436 let mut b = DatagramBuilder::new(2);
437 b.push_each(|device| {
438 Some(ConfigModulation {
439 bank: if device == 0 {
440 ModulationBank::B0
441 } else {
442 ModulationBank::B1
443 },
444 config: SamplingConfig::FREQ_40K,
445 size: 1,
446 loop_behavior: LoopBehavior::Infinite,
447 })
448 });
449 let frames = b.build().unwrap();
450
451 assert_eq!(frames.len(), 1);
452 let frame = frames.frame(0).unwrap();
453 assert_eq!(frame.distribution(), Distribution::PerDevice);
454 assert_eq!(frame.datagrams()[0].payload[0], 0, "device 0 -> bank B0");
455 assert_eq!(frame.datagrams()[1].payload[0], 1, "device 1 -> bank B1");
456 }
457
458 #[test]
459 fn push_each_fills_unassigned_with_nop() {
460 let mut b = DatagramBuilder::new(2);
461 b.push_each(|device| {
462 (device == 0).then_some(ConfigModulation {
463 bank: ModulationBank::B0,
464 config: SamplingConfig::FREQ_40K,
465 size: 1,
466 loop_behavior: LoopBehavior::Infinite,
467 })
468 });
469 let frames = b.build().unwrap();
470
471 assert_eq!(cmd_at(&frames, 0, 0), Cmd::ConfigModulation);
472 assert_eq!(cmd_at(&frames, 0, 1), Cmd::Nop, "unassigned -> Nop");
473 }
474
475 #[test]
476 fn push_each_pads_shorter_device_with_nop() {
477 let mut b = DatagramBuilder::new(2);
478 b.push_each(|device| Some(if device == 0 { Multi(1) } else { Multi(3) }));
479 let frames = b.build().unwrap();
480
481 assert_eq!(frames.len(), 3, "frame count = max over devices");
482 assert_eq!(cmd_at(&frames, 0, 0), Cmd::ConfigModulation);
483 assert_eq!(cmd_at(&frames, 1, 0), Cmd::Nop);
484 assert_eq!(cmd_at(&frames, 2, 0), Cmd::Nop);
485 for frame in 0..3 {
486 assert_eq!(cmd_at(&frames, frame, 1), Cmd::ConfigModulation);
487 assert_eq!(
488 frames.frame(frame).unwrap().datagrams()[1].payload[0] as usize,
489 frame
490 );
491 }
492 }
493
494 #[test]
495 fn push_each_accepts_heterogeneous_boxed_commands() {
496 let patterns = vec![vec![crate::value::Emission::default(); Autd3::NUM_TRANSDUCERS]; 2];
497 let mut b = DatagramBuilder::new(2);
498 b.push_each(|device| {
499 Some(if device == 0 {
500 Pattern::new(&patterns).boxed()
501 } else {
502 ConfigModulation {
503 bank: ModulationBank::B0,
504 config: SamplingConfig::FREQ_40K,
505 size: 1,
506 loop_behavior: LoopBehavior::Infinite,
507 }
508 .boxed()
509 })
510 });
511 let frames = b.build().unwrap();
512
513 assert_eq!(frames.len(), 3);
514 assert_eq!(cmd_at(&frames, 0, 0), Cmd::WritePatternBuffer);
515 assert_eq!(cmd_at(&frames, 0, 1), Cmd::ConfigModulation);
516 assert_eq!(cmd_at(&frames, 1, 1), Cmd::Nop);
517 assert_eq!(cmd_at(&frames, 2, 1), Cmd::Nop);
518 }
519
520 #[test]
521 fn adjacent_disjoint_push_each_fuse_into_shared_frames() {
522 let mut b = DatagramBuilder::new(2);
523 b.push_each(|device| {
524 (device == 0).then_some(ConfigModulation {
525 bank: ModulationBank::B0,
526 config: SamplingConfig::FREQ_40K,
527 size: 1,
528 loop_behavior: LoopBehavior::Infinite,
529 })
530 });
531 b.push_each(|device| {
532 (device == 1).then_some(ConfigModulation {
533 bank: ModulationBank::B1,
534 config: SamplingConfig::FREQ_40K,
535 size: 1,
536 loop_behavior: LoopBehavior::Infinite,
537 })
538 });
539 let frames = b.build().unwrap();
540
541 assert_eq!(frames.len(), 1, "disjoint groups fuse into one frame");
542 let frame = frames.frame(0).unwrap();
543 assert_eq!(frame.datagrams()[0].payload[0], 0, "device 0 -> B0");
544 assert_eq!(frame.datagrams()[1].payload[0], 1, "device 1 -> B1");
545 }
546
547 #[test]
548 fn adjacent_overlapping_push_each_stay_sequential() {
549 let mut b = DatagramBuilder::new(2);
550 b.push_each(|_| {
551 Some(ConfigModulation {
552 bank: ModulationBank::B0,
553 config: SamplingConfig::FREQ_40K,
554 size: 1,
555 loop_behavior: LoopBehavior::Infinite,
556 })
557 });
558 b.push_each(|_| {
559 Some(ConfigModulation {
560 bank: ModulationBank::B1,
561 config: SamplingConfig::FREQ_40K,
562 size: 1,
563 loop_behavior: LoopBehavior::Infinite,
564 })
565 });
566 let frames = b.build().unwrap();
567
568 assert_eq!(frames.len(), 2, "overlapping coverage stays sequential");
569 assert_eq!(frames.frame(0).unwrap().datagrams()[0].payload[0], 0);
570 assert_eq!(frames.frame(1).unwrap().datagrams()[0].payload[0], 1);
571 }
572
573 #[test]
574 fn broadcast_push_is_a_fuse_barrier() {
575 let mut b = DatagramBuilder::new(2);
576 b.push_each(|device| {
577 (device == 0).then_some(ConfigModulation {
578 bank: ModulationBank::B0,
579 config: SamplingConfig::FREQ_40K,
580 size: 1,
581 loop_behavior: LoopBehavior::Infinite,
582 })
583 });
584 b.push(ConfigPattern {
585 bank: PatternBank::B0,
586 config: SamplingConfig::FREQ_40K,
587 size: 1,
588 loop_behavior: LoopBehavior::Infinite,
589 });
590 b.push_each(|device| {
591 (device == 1).then_some(ConfigModulation {
592 bank: ModulationBank::B1,
593 config: SamplingConfig::FREQ_40K,
594 size: 1,
595 loop_behavior: LoopBehavior::Infinite,
596 })
597 });
598 let frames = b.build().unwrap();
599
600 assert_eq!(frames.len(), 3, "broadcast between steps prevents fusion");
601 assert_eq!(
602 frames.frame(1).unwrap().distribution(),
603 Distribution::Broadcast
604 );
605 }
606 use crate::geometry::Autd3;
607 use crate::value::{Emission, PatternBank};
608
609 #[test]
610 fn broadcast_op_yields_one_frame_of_one_datagram() {
611 let op = ConfigPattern {
612 bank: PatternBank::B0,
613 config: SamplingConfig::FREQ_40K,
614 size: 1,
615 loop_behavior: LoopBehavior::Infinite,
616 };
617 let mut b = DatagramBuilder::new(4);
618 b.push(op);
619 let frames = b.build().unwrap();
620
621 assert_eq!(frames.len(), 1);
622 let frame = frames.frame(0).unwrap();
623 assert_eq!(frame.distribution(), Distribution::Broadcast);
624 assert_eq!(frame.datagrams().len(), 1);
625 assert_eq!(frame.datagrams()[0].cmd, Cmd::ConfigPattern);
626 }
627
628 #[test]
629 fn per_device_op_yields_one_datagram_per_device() {
630 let patterns = vec![vec![Emission::default(); Autd3::NUM_TRANSDUCERS]; 3];
631 let op = WritePatternBuffer {
632 bank: PatternBank::B0,
633 index: 0,
634 emissions: &patterns,
635 };
636 let mut b = DatagramBuilder::new(3);
637 b.push(op);
638 let frames = b.build().unwrap();
639
640 assert_eq!(frames.len(), 1);
641 let frame = frames.frame(0).unwrap();
642 assert_eq!(frame.distribution(), Distribution::PerDevice);
643 assert_eq!(frame.datagrams().len(), 3);
644 }
645
646 #[test]
647 fn composite_emission_orders_write_then_config() {
648 let patterns = vec![vec![Emission::default(); Autd3::NUM_TRANSDUCERS]; 2];
649 let we = WritePatternBuffer {
650 bank: PatternBank::B0,
651 index: 0,
652 emissions: &patterns,
653 };
654 let ce = ConfigPattern {
655 bank: PatternBank::B0,
656 config: SamplingConfig::FREQ_40K,
657 size: 1,
658 loop_behavior: LoopBehavior::Infinite,
659 };
660 let mut b = DatagramBuilder::new(2);
661 b.push(we).push(ce);
662 let frames = b.build().unwrap();
663
664 assert_eq!(frames.len(), 2);
665 assert_eq!(
666 frames.frame(0).unwrap().distribution(),
667 Distribution::PerDevice
668 );
669 assert_eq!(frames.frame(0).unwrap().datagrams().len(), 2);
670 assert_eq!(
671 frames.frame(1).unwrap().distribution(),
672 Distribution::Broadcast
673 );
674 assert_eq!(
675 frames.frame(1).unwrap().datagrams()[0].cmd,
676 Cmd::ConfigPattern
677 );
678 }
679
680 #[test]
681 fn build_into_reuses_buffer_without_growing() {
682 let op = ConfigPattern {
683 bank: PatternBank::B0,
684 config: SamplingConfig::FREQ_40K,
685 size: 1,
686 loop_behavior: LoopBehavior::Infinite,
687 };
688 let mut b = DatagramBuilder::new(1);
689 b.push(op);
690
691 let mut buf = Frames::default();
692 b.build_into(&mut buf).unwrap();
693 let cap_after_first = buf.payloads.capacity();
694 b.build_into(&mut buf).unwrap();
695
696 assert_eq!(buf.len(), 1);
697 assert_eq!(
698 buf.payloads.capacity(),
699 cap_after_first,
700 "second build must not reallocate"
701 );
702 }
703}