1use crate::{ComponentMap, Keyed, WithArgs};
2use futures::future::join_all;
3
4impl<Key, Args, Comp, FnInit> ComponentMap<Key, Args, Comp, FnInit> {
5 pub async fn init_async(args: impl IntoIterator<Item = (Key, Args)>, init: FnInit) -> Self
6 where
7 Key: Eq + std::hash::Hash,
8 FnInit: AsyncFn(&Key, &Args) -> Comp + Clone,
9 {
10 let components_fut = args.into_iter().map(|(key, args)| {
11 let init = init.clone();
12 async move {
13 let component = (init)(&key, &args).await;
14 (key, WithArgs { component, args })
15 }
16 });
17
18 let map = join_all(components_fut).await.into_iter().collect();
19
20 Self { map: map, init }
21 }
22
23 pub async fn reinit_all_async(&mut self) -> impl Iterator<Item = Keyed<&Key, Comp>>
24 where
25 FnInit: AsyncFn(&Key, &Args) -> Comp + Clone,
26 {
27 let next_components_fut = self
28 .map
29 .iter()
30 .map(|(key, component)| (self.init)(key, &component.args));
31
32 let next_components = join_all(next_components_fut).await;
33
34 self.map
35 .iter_mut()
36 .zip(next_components)
37 .map(|((key, prev), next)| {
38 let prev = std::mem::replace(&mut prev.component, next);
39 Keyed::new(key, prev)
40 })
41 }
42
43 pub async fn reinit_async(
44 &mut self,
45 keys: impl IntoIterator<Item = Key>,
46 ) -> impl Iterator<Item = Keyed<Key, Option<Comp>>>
47 where
48 Key: Eq + std::hash::Hash + Clone,
49 FnInit: AsyncFn(&Key, &Args) -> Comp + Clone,
50 {
51 let next_components_fut = keys.into_iter().map(|key| {
52 let init = self.init.clone();
53 let args = self.map.get(&key).map(|component| &component.args);
54 async move {
55 let next = match args {
56 Some(args) => Some((init)(&key, args).await),
57 None => None,
58 };
59 Keyed::new(key, next)
60 }
61 });
62
63 let results = join_all(next_components_fut).await;
64
65 results.into_iter().map(|Keyed { key, value: next }| {
66 let prev = next.and_then(|next| {
67 self.map
68 .get_mut(&key)
69 .map(|component| std::mem::replace(&mut component.component, next))
70 });
71 Keyed::new(key, prev)
72 })
73 }
74
75 pub async fn update_async(
76 &mut self,
77 updates: impl IntoIterator<Item = (Key, Args)>,
78 ) -> impl Iterator<Item = Keyed<Key, Option<WithArgs<Args, Comp>>>>
79 where
80 Key: Clone + Eq + std::hash::Hash,
81 FnInit: AsyncFn(&Key, &Args) -> Comp + Clone,
82 {
83 let updated_components_fut = updates.into_iter().map(|(key, args)| {
84 let init = self.init.clone();
85 async move {
86 let component = (init)(&key, &args).await;
87 (key, WithArgs { component, args })
88 }
89 });
90
91 join_all(updated_components_fut)
92 .await
93 .into_iter()
94 .map(|(key, component)| {
95 let prev = self.map.insert(key.clone(), component);
96 Keyed::new(key, prev)
97 })
98 }
99}
100
101#[cfg(test)]
102mod tests {
103 use super::*;
104 use std::sync::{Arc, Mutex};
105
106 #[derive(Debug, Clone, PartialEq, Eq)]
107 struct Counter(usize);
108
109 #[derive(Debug, Clone, PartialEq, Eq)]
110 struct Args {
111 value: usize,
112 }
113
114 #[tokio::test]
115 async fn test_init_async() {
116 let init = |_key: &&str, args: &Args| {
117 let value = args.value;
118 async move { Counter(value) }
119 };
120 let manager = ComponentMap::init_async(
121 [("key1", Args { value: 1 }), ("key2", Args { value: 2 })],
122 init,
123 )
124 .await;
125
126 assert_eq!(manager.map.len(), 2);
127 assert_eq!(manager.map.get("key1").unwrap().component, Counter(1));
128 assert_eq!(manager.map.get("key2").unwrap().component, Counter(2));
129 assert_eq!(manager.map.get("key1").unwrap().args.value, 1);
130 }
131
132 #[tokio::test]
133 async fn test_init_async_empty() {
134 let init = |_key: &&str, args: &Args| {
135 let value = args.value;
136 async move { Counter(value) }
137 };
138 let manager: ComponentMap<&str, Args, Counter, _> =
139 ComponentMap::init_async([], init).await;
140
141 assert_eq!(manager.map.len(), 0);
142 }
143
144 #[tokio::test]
145 async fn test_reinit_all_async() {
146 let call_count = Arc::new(Mutex::new(0));
147 let call_count_clone = call_count.clone();
148
149 let init = move |_key: &&str, args: &Args| {
150 let call_count = call_count_clone.clone();
151 let value = args.value;
152 async move {
153 *call_count.lock().unwrap() += 1;
154 Counter(value * 2)
155 }
156 };
157
158 let mut manager = ComponentMap::init_async(
159 [("key1", Args { value: 1 }), ("key2", Args { value: 2 })],
160 init,
161 )
162 .await;
163
164 let prev_components: Vec<_> = manager.reinit_all_async().await.collect();
165
166 assert_eq!(prev_components.len(), 2);
167 assert_eq!(manager.map.get("key1").unwrap().component, Counter(2));
168 assert_eq!(manager.map.get("key2").unwrap().component, Counter(4));
169
170 assert_eq!(*call_count.lock().unwrap(), 4);
172 }
173
174 #[tokio::test]
175 async fn test_reinit_all_async_empty() {
176 let init = |_key: &&str, args: &Args| {
177 let value = args.value;
178 async move { Counter(value) }
179 };
180 let mut manager: ComponentMap<&str, Args, Counter, _> =
181 ComponentMap::init_async([], init).await;
182
183 let results: Vec<_> = manager.reinit_all_async().await.collect();
184 assert_eq!(results.len(), 0);
185 }
186
187 #[tokio::test]
188 async fn test_reinit_async_existing_key() {
189 let init = |_key: &&str, args: &Args| {
190 let value = args.value;
191 async move { Counter(value * 2) }
192 };
193
194 let mut manager = ComponentMap::init_async(
195 [("key1", Args { value: 1 }), ("key2", Args { value: 2 })],
196 init,
197 )
198 .await;
199
200 let results: Vec<_> = manager.reinit_async(["key1"]).await.collect();
201
202 assert_eq!(results.len(), 1);
203 assert_eq!(results[0].key, "key1");
204 assert_eq!(results[0].value, Some(Counter(2)));
205
206 assert_eq!(manager.map.get("key1").unwrap().component, Counter(2));
207 assert_eq!(manager.map.get("key2").unwrap().component, Counter(4));
208 }
209
210 #[tokio::test]
211 async fn test_reinit_async_multiple_keys() {
212 let init = |_key: &&str, args: &Args| {
213 let value = args.value;
214 async move { Counter(value * 3) }
215 };
216
217 let mut manager = ComponentMap::init_async(
218 [
219 ("key1", Args { value: 1 }),
220 ("key2", Args { value: 2 }),
221 ("key3", Args { value: 3 }),
222 ],
223 init,
224 )
225 .await;
226
227 let results: Vec<_> = manager.reinit_async(["key1", "key3"]).await.collect();
228
229 assert_eq!(results.len(), 2);
230 assert_eq!(manager.map.get("key1").unwrap().component, Counter(3));
231 assert_eq!(manager.map.get("key2").unwrap().component, Counter(6));
232 assert_eq!(manager.map.get("key3").unwrap().component, Counter(9));
233 }
234
235 #[tokio::test]
236 async fn test_reinit_async_nonexistent_key() {
237 let init = |_key: &&str, args: &Args| {
238 let value = args.value;
239 async move { Counter(value) }
240 };
241
242 let mut manager = ComponentMap::init_async([("key1", Args { value: 1 })], init).await;
243
244 let results: Vec<_> = manager.reinit_async(["nonexistent"]).await.collect();
245
246 assert_eq!(results.len(), 1);
247 assert_eq!(results[0].key, "nonexistent");
248 assert_eq!(results[0].value, None);
249 assert_eq!(manager.map.len(), 1);
250 }
251
252 #[tokio::test]
253 async fn test_update_async_existing() {
254 let init = |_key: &&str, args: &Args| {
255 let value = args.value;
256 async move { Counter(value) }
257 };
258
259 let mut manager = ComponentMap::init_async([("key1", Args { value: 1 })], init).await;
260
261 let results: Vec<_> = manager
262 .update_async([("key1", Args { value: 10 })])
263 .await
264 .collect();
265
266 assert_eq!(results.len(), 1);
267 assert_eq!(results[0].key, "key1");
268 assert!(results[0].value.is_some());
269 assert_eq!(results[0].value.as_ref().unwrap().component, Counter(1));
270
271 assert_eq!(manager.map.get("key1").unwrap().component, Counter(10));
272 assert_eq!(manager.map.get("key1").unwrap().args.value, 10);
273 }
274
275 #[tokio::test]
276 async fn test_update_async_new_key() {
277 let init = |_key: &&str, args: &Args| {
278 let value = args.value;
279 async move { Counter(value) }
280 };
281
282 let mut manager = ComponentMap::init_async([("key1", Args { value: 1 })], init).await;
283
284 let results: Vec<_> = manager
285 .update_async([("key2", Args { value: 20 })])
286 .await
287 .collect();
288
289 assert_eq!(results.len(), 1);
290 assert_eq!(results[0].key, "key2");
291 assert!(results[0].value.is_none());
292
293 assert_eq!(manager.map.len(), 2);
294 assert_eq!(manager.map.get("key2").unwrap().component, Counter(20));
295 }
296
297 #[tokio::test]
298 async fn test_update_async_multiple() {
299 let init = |_key: &&str, args: &Args| {
300 let value = args.value;
301 async move { Counter(value) }
302 };
303
304 let mut manager = ComponentMap::init_async([("key1", Args { value: 1 })], init).await;
305
306 let results: Vec<_> = manager
307 .update_async([
308 ("key1", Args { value: 10 }),
309 ("key2", Args { value: 20 }),
310 ("key3", Args { value: 30 }),
311 ])
312 .await
313 .collect();
314
315 assert_eq!(results.len(), 3);
316 assert_eq!(manager.map.len(), 3);
317 assert_eq!(manager.map.get("key1").unwrap().component, Counter(10));
318 assert_eq!(manager.map.get("key2").unwrap().component, Counter(20));
319 assert_eq!(manager.map.get("key3").unwrap().component, Counter(30));
320 }
321}