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