crossflow/chain/
unzip.rs

1/*
2 * Copyright (C) 2024 Open Source Robotics Foundation
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 *
16*/
17
18use variadics_please::all_tuples;
19
20use itertools::Itertools;
21use smallvec::SmallVec;
22
23use crate::{
24    AddOperation, Builder, Chain, ForkTargetStorage, ForkUnzip, Input, ManageInput,
25    OperationRequest, OperationResult, OrBroken, Output, UnusedTarget,
26};
27
28/// A trait for response types that can be unzipped
29pub trait Unzippable: Sized {
30    type Unzipped;
31    fn unzip_output(output: Output<Self>, builder: &mut Builder) -> Self::Unzipped;
32
33    fn distribute_values(request: OperationRequest) -> OperationResult;
34
35    type Prepended<T>;
36    fn prepend<T>(self, value: T) -> Self::Prepended<T>;
37}
38
39macro_rules! impl_unzippable_for_tuple {
40    ($(($T:ident, $D:ident)),*) => {
41        #[allow(non_snake_case)]
42        impl<$($T: 'static + Send + Sync),*> Unzippable for ($($T,)*)
43        {
44            type Unzipped = ($(Output<$T>,)*);
45            fn unzip_output(output: Output<Self>, builder: &mut Builder) -> Self::Unzipped {
46                assert_eq!(output.scope(), builder.scope());
47                let mut targets = SmallVec::new();
48                let result =
49                    (
50                        $(
51                            {
52                                // Variable is only used to make sure this cycle is repeated once
53                                // for each instance of the $T type, but the type itself is not
54                                // used.
55                                #[allow(unused)]
56                                let $T = std::marker::PhantomData::<$T>;
57                                let target = builder.commands.spawn(UnusedTarget).id();
58                                targets.push(target);
59                                Output::new(builder.scope(), target)
60                            },
61                        )*
62                    );
63
64                builder.commands.queue(AddOperation::new(
65                    Some(output.scope()),
66                    output.id(),
67                    ForkUnzip::<Self>::new(ForkTargetStorage(targets)),
68                ));
69                result
70            }
71
72            fn distribute_values(
73                OperationRequest { source, world, roster }: OperationRequest,
74            ) -> OperationResult {
75                let Input { session, data: inputs } = world
76                    .get_entity_mut(source).or_broken()?
77                    .take_input::<Self>()?;
78
79                let ($($D,)*) = world.get::<ForkTargetStorage>(source).or_broken()?.0.iter().copied().next_tuple().or_broken()?;
80                let ($($T,)*) = inputs;
81                $(
82                    if let Ok(mut t_mut) = world.get_entity_mut($D) {
83                        t_mut.give_input(session, $T, roster)?;
84                    }
85                )*
86                Ok(())
87            }
88
89            type Prepended<T> = (T, $($T,)*);
90            fn prepend<T>(self, value: T) -> Self::Prepended<T> {
91                let ($($T,)*) = self;
92                (value, $($T,)*)
93            }
94        }
95    }
96}
97
98// Implements the `Unzippable` trait for all tuples between size 1 and 12
99// (inclusive) made of 'static lifetime types that are `Send` and `Sync`
100// D is a dummy type
101all_tuples!(impl_unzippable_for_tuple, 1, 12, T, D);
102
103/// A trait for constructs that are able to perform a forking unzip of an
104/// unzippable chain. An unzippable chain is one whose response type contains a
105/// tuple.
106pub trait UnzipBuilder<Z> {
107    type ReturnType;
108    fn fork_unzip(self, output: Output<Z>, builder: &mut Builder) -> Self::ReturnType;
109}
110
111macro_rules! impl_unzipbuilder_for_tuple {
112    ($(($A:ident, $F:ident, $U:ident)),*) => {
113        #[allow(non_snake_case)]
114        impl<$($A: 'static + Send + Sync),*, $($F: FnOnce(Chain<$A>) -> $U),*, $($U),*> UnzipBuilder<($($A,)*)> for ($($F,)*)
115        {
116            type ReturnType = ($($U),*);
117            fn fork_unzip(self, output: Output<($($A,)*)>, builder: &mut Builder) -> Self::ReturnType {
118                let outputs = <($($A),*)>::unzip_output(output, builder);
119                let ($($A,)*) = outputs;
120                let ($($F,)*) = self;
121                (
122                    $(
123                        ($F)($A.chain(builder)),
124                    )*
125                )
126            }
127        }
128    }
129}
130
131// Implements the `UnzipBuilder` trait for all tuples between size 1 and 12
132all_tuples!(impl_unzipbuilder_for_tuple, 2, 12, A, F, U);
crossflow/chain/unzip.rs

crossflow/chain/
unzip.rs
