Skip to main content

rama_core/combinators/
either.rs

1#![expect(
2    clippy::allow_attributes,
3    reason = "macro-emitted `#[allow(clippy::multiple_unsafe_ops_per_block)]`: older rustc versions emitted the lint, newer versions don't, so `#[expect]` can warn unfulfilled"
4)]
5
6use core::fmt;
7use core::pin::Pin;
8use core::task::Poll;
9
10#[cfg(feature = "std")]
11use core::task::Context as TaskContext;
12#[cfg(feature = "std")]
13use std::io::IoSlice;
14#[cfg(feature = "std")]
15use tokio::io::{AsyncRead, AsyncWrite, Error as IoError, ReadBuf, Result as IoResult};
16
17#[macro_export]
18#[doc(hidden)]
19/// Implement the `Either` type for the available number of type parameters,
20/// using the given macro to define each variant.
21macro_rules! __impl_either {
22    ($macro:ident) => {
23        $macro!(Either, A, B,);
24        $macro!(Either3, A, B, C,);
25        $macro!(Either4, A, B, C, D,);
26        $macro!(Either5, A, B, C, D, E,);
27        $macro!(Either6, A, B, C, D, E, F,);
28        $macro!(Either7, A, B, C, D, E, F, G,);
29        $macro!(Either8, A, B, C, D, E, F, G, H,);
30        $macro!(Either9, A, B, C, D, E, F, G, H, I,);
31    };
32}
33
34#[doc(inline)]
35pub use crate::__impl_either as impl_either;
36
37#[macro_export]
38#[doc(hidden)]
39macro_rules! __define_either {
40    ($id:ident, $($param:ident),+ $(,)?) => {
41        /// A type to allow you to use multiple types as a single type.
42        ///
43        /// and will delegate the functionality to the type that is wrapped in the `Either` type.
44        /// To keep it easy all wrapped types are expected to work with the same inputs and outputs.
45        ///
46        /// You can use [`crate::combinators::impl_either`] to
47        /// implement the `Either` type for the available number of type parameters
48        /// on your own Trait implementations.
49        #[derive(Debug, Clone)]
50        pub enum $id<$($param),+> {
51            $(
52                /// one of the Either variants
53                $param($param),
54            )+
55        }
56
57        impl<$($param),+> fmt::Display for $id<$($param),+>
58        where
59            $($param: fmt::Display),+
60        {
61            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
62                match self {
63                    $(
64                        $id::$param(s) => write!(f, "{}", s),
65                    )+
66                }
67            }
68        }
69
70
71        impl<$($param),+> $id<$($param),+> {
72            /// Convert `Pin<&mut Either<A, B>>` to `Either<Pin<&mut A>, Pin<&mut B>>`,
73            /// pinned projections of the inner variants.
74            ///
75            /// This is the pin-projection used to implement pinned traits (e.g.
76            /// `Future`/`AsyncRead`) for `Either`; it is public so downstream
77            /// crates can implement their own pinned traits over `Either` too.
78            #[allow(clippy::multiple_unsafe_ops_per_block, reason = "macro-generated pin projection: get_unchecked_mut + NĂ—Pin::new_unchecked is one logical projection sequence; SAFETY comment below covers the entire block")]
79            pub fn as_pin_mut(self: Pin<&mut Self>) -> $id<$(Pin<&mut $param>),+> {
80                // SAFETY: `get_unchecked_mut` is fine because we don't move anything.
81                // We can use `new_unchecked` because the `inner` parts are guaranteed
82                // to be pinned, as they come from `self` which is pinned, and we never
83                // offer an unpinned `&mut A` or `&mut B` through `Pin<&mut Self>`. We
84                // also don't have an implementation of `Drop`, nor manual `Unpin`.
85                unsafe {
86                    match self.get_unchecked_mut() {
87                        $(
88                            Self::$param(inner) => $id::$param(Pin::new_unchecked(inner)),
89                        )+
90                    }
91                }
92            }
93        }
94
95        impl<$($param),+, Output> Future for $id<$($param),+>
96        where
97            $($param: Future<Output = Output>),+
98        {
99            type Output = Output;
100
101            fn poll(self: Pin<&mut Self>, cx: &mut core::task::Context<'_>) -> Poll<Self::Output> {
102                match self.as_pin_mut() {
103                    $(
104                        $id::$param(fut) => fut.poll(cx),
105                    )+
106                }
107            }
108        }
109    };
110}
111
112#[doc(inline)]
113pub use crate::__define_either as define_either;
114
115impl_either!(define_either);
116
117#[macro_export]
118#[doc(hidden)]
119macro_rules! __impl_iterator_either {
120    ($id:ident, $($param:ident),+ $(,)?) => {
121        impl<$($param),+, Item> Iterator for $id<$($param),+>
122        where
123            $($param: Iterator<Item = Item>),+,
124        {
125            type Item = Item;
126
127            fn next(&mut self) -> Option<Item> {
128                match self {
129                    $(
130                        $id::$param(iter) => iter.next(),
131                    )+
132                }
133            }
134
135            fn size_hint(&self) -> (usize, Option<usize>) {
136                match self {
137                    $(
138                        $id::$param(iter) => iter.size_hint(),
139                    )+
140                }
141            }
142        }
143    };
144}
145
146#[doc(inline)]
147pub use crate::__impl_iterator_either as impl_iterator_either;
148
149impl_either!(impl_iterator_either);
150
151#[macro_export]
152#[doc(hidden)]
153#[cfg(feature = "std")]
154macro_rules! __impl_async_read_write_either {
155    ($id:ident, $($param:ident),+ $(,)?) => {
156        #[warn(clippy::missing_trait_methods)]
157        impl<$($param),+> AsyncRead for $id<$($param),+>
158        where
159            $($param: AsyncRead),+,
160        {
161            fn poll_read(
162                self: Pin<&mut Self>,
163                cx: &mut TaskContext<'_>,
164                buf: &mut ReadBuf<'_>,
165            ) -> Poll<IoResult<()>> {
166                match self.as_pin_mut() {
167                    $(
168                        $id::$param(reader) => reader.poll_read(cx, buf),
169                    )+
170                }
171            }
172        }
173
174        #[warn(clippy::missing_trait_methods)]
175        impl<$($param),+> AsyncWrite for $id<$($param),+>
176        where
177            $($param: AsyncWrite),+,
178        {
179            fn poll_write(
180                self: Pin<&mut Self>,
181                cx: &mut TaskContext<'_>,
182                buf: &[u8],
183            ) -> Poll<Result<usize, IoError>> {
184                match self.as_pin_mut() {
185                    $(
186                        $id::$param(writer) => writer.poll_write(cx, buf),
187                    )+
188                }
189            }
190
191            fn poll_flush(self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<Result<(), IoError>> {
192                match self.as_pin_mut() {
193                    $(
194                        $id::$param(writer) => writer.poll_flush(cx),
195                    )+
196                }
197            }
198
199            fn poll_shutdown(self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<Result<(), IoError>> {
200                match self.as_pin_mut() {
201                    $(
202                        $id::$param(writer) => writer.poll_shutdown(cx),
203                    )+
204                }
205            }
206
207            fn poll_write_vectored(
208                self: Pin<&mut Self>,
209                cx: &mut TaskContext<'_>,
210                bufs: &[IoSlice<'_>],
211            ) -> Poll<Result<usize, IoError>> {
212                match self.as_pin_mut() {
213                    $(
214                        $id::$param(writer) => writer.poll_write_vectored(cx, bufs),
215                    )+
216                }
217            }
218
219            fn is_write_vectored(&self) -> bool {
220                match self {
221                    $(
222                        $id::$param(writer) => writer.is_write_vectored(),
223                    )+
224                }
225            }
226        }
227    };
228}
229
230#[doc(inline)]
231#[cfg(feature = "std")]
232pub use crate::__impl_async_read_write_either as impl_async_read_write_either;
233
234#[cfg(feature = "std")]
235impl_either!(impl_async_read_write_either);