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