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