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::{Datagram, DatagramOption, DeviceMask, Inspectable, InspectionResult},
7    environment::Environment,
8    geometry::{Device, Geometry},
9};
10
11/// [`Datagram`] that divide the devices into groups by given function and send different data to each group.
12///
13/// If the key is `None`, nothing is done for the devices corresponding to the key.
14///
15/// # Example
16///
17/// ```
18/// use std::collections::HashMap;
19/// # use autd3_driver::datagram::*;
20///
21/// Group {
22///     key_map: |dev| match dev.idx() {
23///         0 => Some("silencer"),
24///         1 => Some("disable"),
25///         _ => None,
26///     },
27///     datagram_map: HashMap::from([
28///         ("silencer", Silencer::default()),
29///         ("disable", Silencer::disable()),
30///     ]),
31/// };
32/// ```
33#[derive(Default)]
34pub struct Group<K, D, F>
35where
36    F: Fn(&Device) -> Option<K>,
37{
38    /// Mapping function from device to group key.
39    pub key_map: F,
40    /// Map from group key to [`Datagram`].
41    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    /// Creates a new [`Group`].
52    #[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}