autd3_driver/datagram/
group.rs

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/// [`Datagram`] that divide the devices into groups by given function and send different data to each group.
16///
17/// If the key is `None`, nothing is done for the devices corresponding to the key.
18///
19/// # Example
20///
21/// ```
22/// use std::collections::HashMap;
23/// # use autd3_driver::datagram::*;
24///
25/// Group {
26///     key_map: |dev| match dev.idx() {
27///         0 => Some("silencer"),
28///         1 => Some("disable"),
29///         _ => None,
30///     },
31///     datagram_map: HashMap::from([
32///         ("silencer", Silencer::default()),
33///         ("disable", Silencer::disable()),
34///     ]),
35/// };
36/// ```
37#[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    /// Mapping function from device to group key.
46    #[debug(ignore)]
47    pub key_map: F,
48    /// Map from group key to [`Datagram`].
49    #[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    /// Creates a new [`Group`].
61    #[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}