1use std::{collections::HashMap, fmt::Debug, hash::Hash, time::Duration};
2
3use crate::error::AUTDDriverError;
4
5use autd3_core::{
6 datagram::{
7 Datagram, DatagramOption, DeviceFilter, FirmwareLimits, Inspectable, InspectionResult,
8 },
9 geometry::{Device, Geometry},
10};
11use derive_more::Debug as DeriveDebug;
12use itertools::Itertools;
13
14#[derive(Default, DeriveDebug)]
37pub struct Group<K, D, F>
38where
39 K: Hash + Eq + Debug,
40 D: Datagram,
41 F: Fn(&Device) -> Option<K>,
42 AUTDDriverError: From<<D as Datagram>::Error>,
43{
44 #[debug(ignore)]
46 pub key_map: F,
47 #[debug(ignore)]
49 pub datagram_map: HashMap<K, D>,
50}
51
52impl<K, D, F> Group<K, D, F>
53where
54 K: Hash + Eq + Debug,
55 D: Datagram,
56 F: Fn(&Device) -> Option<K>,
57 AUTDDriverError: From<<D as Datagram>::Error>,
58{
59 #[must_use]
61 pub const fn new(key_map: F, datagram_map: HashMap<K, D>) -> Self {
62 Self {
63 key_map,
64 datagram_map,
65 }
66 }
67
68 fn generate_filter(key_map: &F, geometry: &Geometry) -> HashMap<K, DeviceFilter> {
69 let mut filters: HashMap<K, DeviceFilter> = HashMap::new();
70 geometry.iter().for_each(|dev| {
71 if let Some(key) = key_map(dev) {
72 if let Some(v) = filters.get_mut(&key) {
73 v.set_enable(dev.idx());
74 } else {
75 filters.insert(
76 key,
77 DeviceFilter::from_fn(geometry, |dev_| dev_.idx() == dev.idx()),
78 );
79 }
80 }
81 });
82 filters
83 }
84}
85
86pub struct GroupOpGenerator<K, F, D>
87where
88 K: Hash + Eq + Debug,
89 F: Fn(&Device) -> Option<K>,
90 D: Datagram,
91{
92 pub(crate) key_map: F,
93 pub(crate) generators: HashMap<K, D::G>,
94}
95
96impl<K, D, F> Datagram for Group<K, D, F>
97where
98 K: Hash + Eq + Debug,
99 D: Datagram,
100 F: Fn(&Device) -> Option<K>,
101 AUTDDriverError: From<<D as Datagram>::Error>,
102{
103 type G = GroupOpGenerator<K, F, D>;
104 type Error = AUTDDriverError;
105
106 fn operation_generator(
107 self,
108 geometry: &Geometry,
109 _: &DeviceFilter,
110 limits: &FirmwareLimits,
111 ) -> Result<Self::G, Self::Error> {
112 let Self {
113 key_map,
114 mut datagram_map,
115 } = self;
116
117 let filters = Self::generate_filter(&key_map, geometry);
118
119 let generators = filters
120 .into_iter()
121 .map(|(k, filter)| {
122 let datagram = datagram_map
123 .remove(&k)
124 .ok_or(AUTDDriverError::UnknownKey(format!("{:?}", k)))?;
125 Ok((
126 k,
127 datagram
128 .operation_generator(geometry, &filter, limits)
129 .map_err(AUTDDriverError::from)?,
130 ))
131 })
132 .collect::<Result<_, AUTDDriverError>>()?;
133
134 if !datagram_map.is_empty() {
135 return Err(AUTDDriverError::UnusedKey(
136 datagram_map.keys().map(|k| format!("{:?}", k)).join(", "),
137 ));
138 }
139
140 Ok(GroupOpGenerator {
141 key_map,
142 generators,
143 })
144 }
145
146 fn option(&self) -> DatagramOption {
147 self.datagram_map.values().map(|d| d.option()).fold(
148 DatagramOption {
149 timeout: Duration::ZERO,
150 parallel_threshold: usize::MAX,
151 },
152 DatagramOption::merge,
153 )
154 }
155}
156
157impl<K, D, F> Inspectable for Group<K, D, F>
158where
159 K: Hash + Eq + Debug,
160 D: Datagram + Inspectable,
161 F: Fn(&Device) -> Option<K>,
162 AUTDDriverError: From<<D as Datagram>::Error>,
163{
164 type Result = D::Result;
165
166 fn inspect(
167 self,
168 geometry: &Geometry,
169 _: &DeviceFilter,
170 limits: &FirmwareLimits,
171 ) -> Result<InspectionResult<Self::Result>, AUTDDriverError> {
172 let Self {
173 key_map,
174 mut datagram_map,
175 } = self;
176
177 let filters = Self::generate_filter(&key_map, geometry);
178
179 let results = filters
180 .into_iter()
181 .map(
182 |(k, filter)| -> Result<Vec<Option<Self::Result>>, AUTDDriverError> {
183 {
184 let datagram = datagram_map
185 .remove(&k)
186 .ok_or(AUTDDriverError::UnknownKey(format!("{:?}", k)))?;
187
188 let r = datagram
189 .inspect(geometry, &filter, limits)
190 .map_err(AUTDDriverError::from)?;
191
192 Ok(r.result)
193 }
194 },
195 )
196 .collect::<Result<Vec<_>, _>>()?;
197
198 Ok(InspectionResult {
199 result: results
200 .into_iter()
201 .reduce(|a, b| a.into_iter().zip(b).map(|(a, b)| a.or(b)).collect())
202 .unwrap(),
203 })
204 }
205}
206
207#[cfg(test)]
208mod tests {
209 use super::*;
210
211 use crate::datagram::Clear;
212
213 #[test]
214 fn unknown_key() -> anyhow::Result<()> {
215 let geometry = crate::autd3_device::tests::create_geometry(2);
216
217 assert_eq!(
218 Some(AUTDDriverError::UnknownKey("1".to_owned())),
219 Group::new(|dev| Some(dev.idx()), HashMap::from([(0, Clear {})]))
220 .operation_generator(
221 &geometry,
222 &DeviceFilter::all_enabled(),
223 &FirmwareLimits::unused()
224 )
225 .err()
226 );
227
228 Ok(())
229 }
230
231 #[test]
232 fn unused_key() -> anyhow::Result<()> {
233 let geometry = crate::autd3_device::tests::create_geometry(2);
234
235 assert_eq!(
236 Some(AUTDDriverError::UnusedKey("2".to_owned())),
237 Group::new(
238 |dev| Some(dev.idx()),
239 HashMap::from([(0, Clear {}), (1, Clear {}), (2, Clear {})])
240 )
241 .operation_generator(
242 &geometry,
243 &DeviceFilter::all_enabled(),
244 &FirmwareLimits::unused()
245 )
246 .err()
247 );
248
249 Ok(())
250 }
251}