turbo_tasks/task/
function.rs

1//! # Function tasks
2//!
3//! This module contains the trait definitions and implementations that are
4//! necessary for accepting functions as tasks when using the
5//! `turbo_tasks::function` macro.
6//!
7//! This system is inspired by Bevy's Systems and Axum's Handlers.
8//!
9//! The original principle is somewhat simple: a function is accepted if all
10//! of its arguments implement `TaskInput` and its return type implements
11//! `TaskOutput`. There are a few hoops one needs to jump through to make this
12//! work, but they are described in this blog post:
13//! <https://blog.logrocket.com/rust-bevy-entity-component-system/>
14//!
15//! However, there is an additional complication in our case: async methods
16//! that accept a reference to the receiver as their first argument.
17//!
18//! This complication handled through our own version of the `async_trait`
19//! crate, which allows us to target `async fn` as trait bounds. The naive
20//! approach runs into many issues with lifetimes, hence the need for an
21//! intermediate trait. However, this implementation doesn't support all async
22//! methods (see commented out tests).
23
24use std::{future::Future, marker::PhantomData, pin::Pin};
25
26use anyhow::Result;
27
28use super::{TaskInput, TaskOutput};
29use crate::{RawVc, Vc, VcRead, VcValueType, magic_any::MagicAny};
30
31pub type NativeTaskFuture = Pin<Box<dyn Future<Output = Result<RawVc>> + Send>>;
32
33pub trait TaskFn: Send + Sync + 'static {
34    fn functor(&self, this: Option<RawVc>, arg: &dyn MagicAny) -> Result<NativeTaskFuture>;
35}
36
37pub trait IntoTaskFn<Mode, Inputs> {
38    type TaskFn: TaskFn;
39
40    fn into_task_fn(self) -> Self::TaskFn;
41}
42
43impl<F, Mode, Inputs> IntoTaskFn<Mode, Inputs> for F
44where
45    F: TaskFnInputFunction<Mode, Inputs>,
46    Mode: TaskFnMode,
47    Inputs: TaskInputs,
48{
49    type TaskFn = FunctionTaskFn<F, Mode, Inputs>;
50
51    fn into_task_fn(self) -> Self::TaskFn {
52        FunctionTaskFn {
53            task_fn: self,
54            mode: PhantomData,
55            inputs: PhantomData,
56        }
57    }
58}
59
60pub trait IntoTaskFnWithThis<Mode, This, Inputs> {
61    type TaskFn: TaskFn;
62
63    fn into_task_fn_with_this(self) -> Self::TaskFn;
64}
65
66impl<F, Mode, This, Inputs> IntoTaskFnWithThis<Mode, This, Inputs> for F
67where
68    F: TaskFnInputFunctionWithThis<Mode, This, Inputs>,
69    Mode: TaskFnMode,
70    This: Sync + Send + 'static,
71    Inputs: TaskInputs,
72{
73    type TaskFn = FunctionTaskFnWithThis<F, Mode, This, Inputs>;
74
75    fn into_task_fn_with_this(self) -> Self::TaskFn {
76        FunctionTaskFnWithThis {
77            task_fn: self,
78            mode: PhantomData,
79            this: PhantomData,
80            inputs: PhantomData,
81        }
82    }
83}
84
85pub struct FunctionTaskFn<F, Mode: TaskFnMode, Inputs: TaskInputs> {
86    task_fn: F,
87    mode: PhantomData<Mode>,
88    inputs: PhantomData<Inputs>,
89}
90
91impl<F, Mode, Inputs> TaskFn for FunctionTaskFn<F, Mode, Inputs>
92where
93    F: TaskFnInputFunction<Mode, Inputs>,
94    Mode: TaskFnMode,
95    Inputs: TaskInputs,
96{
97    fn functor(&self, _this: Option<RawVc>, arg: &dyn MagicAny) -> Result<NativeTaskFuture> {
98        TaskFnInputFunction::functor(&self.task_fn, arg)
99    }
100}
101
102pub struct FunctionTaskFnWithThis<
103    F,
104    Mode: TaskFnMode,
105    This: Sync + Send + 'static,
106    Inputs: TaskInputs,
107> {
108    task_fn: F,
109    mode: PhantomData<Mode>,
110    this: PhantomData<This>,
111    inputs: PhantomData<Inputs>,
112}
113
114impl<F, Mode, This, Inputs> TaskFn for FunctionTaskFnWithThis<F, Mode, This, Inputs>
115where
116    F: TaskFnInputFunctionWithThis<Mode, This, Inputs>,
117    Mode: TaskFnMode,
118    This: Sync + Send + 'static,
119    Inputs: TaskInputs,
120{
121    fn functor(&self, this: Option<RawVc>, arg: &dyn MagicAny) -> Result<NativeTaskFuture> {
122        let Some(this) = this else {
123            panic!("Method needs a `self` argument");
124        };
125        TaskFnInputFunctionWithThis::functor(&self.task_fn, this, arg)
126    }
127}
128
129trait TaskFnInputFunction<Mode: TaskFnMode, Inputs: TaskInputs>: Send + Sync + Clone + 'static {
130    fn functor(&self, arg: &dyn MagicAny) -> Result<NativeTaskFuture>;
131}
132
133trait TaskFnInputFunctionWithThis<Mode: TaskFnMode, This: Sync + Send + 'static, Inputs: TaskInputs>:
134    Send + Sync + Clone + 'static
135{
136    fn functor(&self, this: RawVc, arg: &dyn MagicAny) -> Result<NativeTaskFuture>;
137}
138
139pub trait TaskInputs: Send + Sync + 'static {}
140
141/// Modes to allow multiple `TaskFnInputFunction` blanket implementations on
142/// `Fn`s. Even though the implementations are non-conflicting in practice, they
143/// could be in theory (at least from with the compiler's current limitations).
144/// Despite this, the compiler is still able to infer the correct mode from a
145/// function.
146pub trait TaskFnMode: Send + Sync + 'static {}
147
148pub struct FunctionMode;
149impl TaskFnMode for FunctionMode {}
150
151pub struct AsyncFunctionMode;
152impl TaskFnMode for AsyncFunctionMode {}
153
154pub struct MethodMode;
155impl TaskFnMode for MethodMode {}
156
157pub struct AsyncMethodMode;
158impl TaskFnMode for AsyncMethodMode {}
159
160macro_rules! task_inputs_impl {
161    ( $( $arg:ident )* ) => {
162        impl<$($arg,)*> TaskInputs for ($($arg,)*)
163        where
164            $($arg: TaskInput + 'static,)*
165        {}
166    }
167}
168
169/// Downcast, and clone all the arguments in the singular `arg` tuple.
170///
171/// This helper function for `task_fn_impl!()` reduces the amount of code inside the macro, and
172/// gives the compiler more chances to dedupe monomorphized code across small functions with less
173/// typevars.
174fn get_args<T: MagicAny + Clone>(arg: &dyn MagicAny) -> Result<T> {
175    let value = (arg as &dyn std::any::Any).downcast_ref::<T>().cloned();
176    #[cfg(debug_assertions)]
177    return anyhow::Context::with_context(value, || {
178        crate::native_function::debug_downcast_args_error_msg(
179            std::any::type_name::<T>(),
180            arg.magic_type_name(),
181        )
182    });
183    #[cfg(not(debug_assertions))]
184    return anyhow::Context::context(value, "Invalid argument type");
185}
186
187// Helper function for `task_fn_impl!()`
188async fn output_try_into_non_local_raw_vc(output: impl TaskOutput) -> Result<RawVc> {
189    output.try_into_raw_vc()?.to_non_local().await
190}
191
192macro_rules! task_fn_impl {
193    ( $async_fn_trait:ident $arg_len:literal $( $arg:ident )* ) => {
194        impl<F, Output, $($arg,)*> TaskFnInputFunction<FunctionMode, ($($arg,)*)> for F
195        where
196            $($arg: TaskInput + 'static,)*
197            F: Fn($($arg,)*) -> Output + Send + Sync + Clone + 'static,
198            Output: TaskOutput + 'static,
199        {
200            #[allow(non_snake_case)]
201            fn functor(&self, arg: &dyn MagicAny) -> Result<NativeTaskFuture> {
202                let task_fn = self.clone();
203                let ($($arg,)*) = get_args::<($($arg,)*)>(arg)?;
204                Ok(Box::pin(async move {
205                    let output = (task_fn)($($arg,)*);
206                    output_try_into_non_local_raw_vc(output).await
207                }))
208            }
209        }
210
211        impl<F, Output, FutureOutput, $($arg,)*> TaskFnInputFunction<AsyncFunctionMode, ($($arg,)*)> for F
212        where
213            $($arg: TaskInput + 'static,)*
214            F: Fn($($arg,)*) -> FutureOutput + Send + Sync + Clone + 'static,
215            FutureOutput: Future<Output = Output> + Send + 'static,
216            Output: TaskOutput + 'static,
217        {
218            #[allow(non_snake_case)]
219            fn functor(&self, arg: &dyn MagicAny) -> Result<NativeTaskFuture> {
220                let task_fn = self.clone();
221                let ($($arg,)*) = get_args::<($($arg,)*)>(arg)?;
222                Ok(Box::pin(async move {
223                    let output = (task_fn)($($arg,)*).await;
224                    output_try_into_non_local_raw_vc(output).await
225                }))
226            }
227        }
228
229        impl<F, Output, Recv, $($arg,)*> TaskFnInputFunctionWithThis<MethodMode, Recv, ($($arg,)*)> for F
230        where
231            Recv: VcValueType,
232            $($arg: TaskInput + 'static,)*
233            F: Fn(&Recv, $($arg,)*) -> Output + Send + Sync + Clone + 'static,
234            Output: TaskOutput + 'static,
235        {
236            #[allow(non_snake_case)]
237            fn functor(&self, this: RawVc, arg: &dyn MagicAny) -> Result<NativeTaskFuture> {
238                let task_fn = self.clone();
239                let recv = Vc::<Recv>::from(this);
240                let ($($arg,)*) = get_args::<($($arg,)*)>(arg)?;
241                Ok(Box::pin(async move {
242                    let recv = recv.await?;
243                    let recv = <Recv::Read as VcRead<Recv>>::target_to_value_ref(&*recv);
244                    let output = (task_fn)(recv, $($arg,)*);
245                    output_try_into_non_local_raw_vc(output).await
246                }))
247            }
248        }
249
250        impl<F, Output, Recv, $($arg,)*> TaskFnInputFunctionWithThis<FunctionMode, Recv, ($($arg,)*)> for F
251        where
252            Recv: Sync + Send + 'static,
253            $($arg: TaskInput + 'static,)*
254            F: Fn(Vc<Recv>, $($arg,)*) -> Output + Send + Sync + Clone + 'static,
255            Output: TaskOutput + 'static,
256        {
257            #[allow(non_snake_case)]
258            fn functor(&self, this: RawVc, arg: &dyn MagicAny) -> Result<NativeTaskFuture> {
259                let task_fn = self.clone();
260                let recv = Vc::<Recv>::from(this);
261                let ($($arg,)*) = get_args::<($($arg,)*)>(arg)?;
262                Ok(Box::pin(async move {
263                    let output = (task_fn)(recv, $($arg,)*);
264                    output_try_into_non_local_raw_vc(output).await
265                }))
266            }
267        }
268
269        pub trait $async_fn_trait<A0, $($arg,)*>: Fn(A0, $($arg,)*) -> Self::OutputFuture {
270            type OutputFuture: Future<Output = <Self as $async_fn_trait<A0, $($arg,)*>>::Output> + Send;
271            type Output: TaskOutput;
272        }
273
274        impl<F: ?Sized, Fut, A0, $($arg,)*> $async_fn_trait<A0, $($arg,)*> for F
275        where
276            F: Fn(A0, $($arg,)*) -> Fut,
277            Fut: Future + Send,
278            Fut::Output: TaskOutput + 'static
279        {
280            type OutputFuture = Fut;
281            type Output = Fut::Output;
282        }
283
284        impl<F, Recv, $($arg,)*> TaskFnInputFunctionWithThis<AsyncMethodMode, Recv, ($($arg,)*)> for F
285        where
286            Recv: VcValueType,
287            $($arg: TaskInput + 'static,)*
288            F: for<'a> $async_fn_trait<&'a Recv, $($arg,)*> + Clone + Send + Sync + 'static,
289        {
290            #[allow(non_snake_case)]
291            fn functor(&self, this: RawVc, arg: &dyn MagicAny) -> Result<NativeTaskFuture> {
292                let task_fn = self.clone();
293                let recv = Vc::<Recv>::from(this);
294                let ($($arg,)*) = get_args::<($($arg,)*)>(arg)?;
295                Ok(Box::pin(async move {
296                    let recv = recv.await?;
297                    let recv = <Recv::Read as VcRead<Recv>>::target_to_value_ref(&*recv);
298                    let output = (task_fn)(recv, $($arg,)*).await;
299                    output_try_into_non_local_raw_vc(output).await
300                }))
301            }
302        }
303
304        impl<F, Recv, $($arg,)*> TaskFnInputFunctionWithThis<AsyncFunctionMode, Recv, ($($arg,)*)> for F
305        where
306            Recv: Sync + Send + 'static,
307            $($arg: TaskInput + 'static,)*
308            F: $async_fn_trait<Vc<Recv>, $($arg,)*> + Clone + Send + Sync + 'static,
309        {
310            #[allow(non_snake_case)]
311            fn functor(&self, this: RawVc, arg: &dyn MagicAny) -> Result<NativeTaskFuture> {
312                let task_fn = self.clone();
313                let recv = Vc::<Recv>::from(this);
314                let ($($arg,)*) = get_args::<($($arg,)*)>(arg)?;
315                Ok(Box::pin(async move {
316                    let output = (task_fn)(recv, $($arg,)*).await;
317                    output_try_into_non_local_raw_vc(output).await
318                }))
319            }
320        }
321    };
322}
323
324task_fn_impl! { AsyncFn0 0 }
325task_fn_impl! { AsyncFn1 1 A1 }
326task_fn_impl! { AsyncFn2 2 A1 A2 }
327task_fn_impl! { AsyncFn3 3 A1 A2 A3 }
328task_fn_impl! { AsyncFn4 4 A1 A2 A3 A4 }
329task_fn_impl! { AsyncFn5 5 A1 A2 A3 A4 A5 }
330task_fn_impl! { AsyncFn6 6 A1 A2 A3 A4 A5 A6 }
331task_fn_impl! { AsyncFn7 7 A1 A2 A3 A4 A5 A6 A7 }
332task_fn_impl! { AsyncFn8 8 A1 A2 A3 A4 A5 A6 A7 A8 }
333task_fn_impl! { AsyncFn9 9 A1 A2 A3 A4 A5 A6 A7 A8 A9 }
334task_fn_impl! { AsyncFn10 10 A1 A2 A3 A4 A5 A6 A7 A8 A9 A10 }
335task_fn_impl! { AsyncFn11 11 A1 A2 A3 A4 A5 A6 A7 A8 A9 A10 A11 }
336task_fn_impl! { AsyncFn12 12 A1 A2 A3 A4 A5 A6 A7 A8 A9 A10 A11 A12 }
337
338// There needs to be one more implementation than task_fn_impl to account for
339// the receiver.
340task_inputs_impl! {}
341task_inputs_impl! { A1 }
342task_inputs_impl! { A1 A2 }
343task_inputs_impl! { A1 A2 A3 }
344task_inputs_impl! { A1 A2 A3 A4 }
345task_inputs_impl! { A1 A2 A3 A4 A5 }
346task_inputs_impl! { A1 A2 A3 A4 A5 A6 }
347task_inputs_impl! { A1 A2 A3 A4 A5 A6 A7 }
348task_inputs_impl! { A1 A2 A3 A4 A5 A6 A7 A8 }
349task_inputs_impl! { A1 A2 A3 A4 A5 A6 A7 A8 A9 }
350task_inputs_impl! { A1 A2 A3 A4 A5 A6 A7 A8 A9 A10 }
351task_inputs_impl! { A1 A2 A3 A4 A5 A6 A7 A8 A9 A10 A11 }
352task_inputs_impl! { A1 A2 A3 A4 A5 A6 A7 A8 A9 A10 A11 A12 }
353task_inputs_impl! { A1 A2 A3 A4 A5 A6 A7 A8 A9 A10 A11 A12 A13 }
354
355#[cfg(test)]
356mod tests {
357    use turbo_rcstr::RcStr;
358
359    use super::*;
360    use crate::{ShrinkToFit, VcCellNewMode, VcDefaultRead};
361
362    #[test]
363    fn test_task_fn() {
364        fn no_args() -> crate::Vc<i32> {
365            todo!()
366        }
367
368        fn one_arg(_a: i32) -> crate::Vc<i32> {
369            todo!()
370        }
371
372        async fn async_one_arg(_a: i32) -> crate::Vc<i32> {
373            todo!()
374        }
375
376        fn with_recv(_a: &i32) -> crate::Vc<i32> {
377            todo!()
378        }
379
380        async fn async_with_recv(_a: &i32) -> crate::Vc<i32> {
381            todo!()
382        }
383
384        fn with_recv_and_str(_a: &i32, _s: RcStr) -> crate::Vc<i32> {
385            todo!()
386        }
387
388        async fn async_with_recv_and_str(_a: &i32, _s: RcStr) -> crate::Vc<i32> {
389            todo!()
390        }
391
392        async fn async_with_recv_and_str_and_result(_a: &i32, _s: RcStr) -> Result<crate::Vc<i32>> {
393            todo!()
394        }
395
396        fn accepts_task_fn<F>(_task_fn: F)
397        where
398            F: TaskFn,
399        {
400        }
401
402        struct Struct;
403        impl Struct {
404            async fn inherent_method(&self) {}
405        }
406
407        impl ShrinkToFit for Struct {
408            fn shrink_to_fit(&mut self) {}
409        }
410
411        unsafe impl VcValueType for Struct {
412            type Read = VcDefaultRead<Struct>;
413
414            type CellMode = VcCellNewMode<Struct>;
415
416            fn get_value_type_id() -> crate::ValueTypeId {
417                todo!()
418            }
419        }
420
421        trait AsyncTrait {
422            async fn async_method(&self);
423        }
424
425        impl AsyncTrait for Struct {
426            async fn async_method(&self) {
427                todo!()
428            }
429        }
430
431        /*
432        async fn async_with_recv_and_str_and_lf(
433            _a: &i32,
434            _s: String,
435        ) -> Result<crate::Vc<i32>, crate::Vc<i32>> {
436            todo!()
437        }
438
439        #[async_trait::async_trait]
440        trait BoxAsyncTrait {
441            async fn box_async_method(&self);
442        }
443
444        #[async_trait::async_trait]
445        impl BoxAsyncTrait for Struct {
446            async fn box_async_method(&self) {
447                todo!()
448            }
449        }
450        */
451
452        let _task_fn = no_args.into_task_fn();
453        accepts_task_fn(no_args.into_task_fn());
454        let _task_fn = one_arg.into_task_fn();
455        accepts_task_fn(one_arg.into_task_fn());
456        let _task_fn = async_one_arg.into_task_fn();
457        accepts_task_fn(async_one_arg.into_task_fn());
458        let task_fn = with_recv.into_task_fn_with_this();
459        accepts_task_fn(task_fn);
460        let task_fn = async_with_recv.into_task_fn_with_this();
461        accepts_task_fn(task_fn);
462        let task_fn = with_recv_and_str.into_task_fn_with_this();
463        accepts_task_fn(task_fn);
464        let task_fn = async_with_recv_and_str.into_task_fn_with_this();
465        accepts_task_fn(task_fn);
466        let task_fn = async_with_recv_and_str_and_result.into_task_fn_with_this();
467        accepts_task_fn(task_fn);
468        let task_fn = <Struct as AsyncTrait>::async_method.into_task_fn_with_this();
469        accepts_task_fn(task_fn);
470        let task_fn = Struct::inherent_method.into_task_fn_with_this();
471        accepts_task_fn(task_fn);
472
473        /*
474        let task_fn = <Struct as BoxAsyncTrait>::box_async_method.into_task_fn();
475        accepts_task_fn(task_fn);
476        let task_fn = async_with_recv_and_str_and_lf.into_task_fn();
477        accepts_task_fn(task_fn);
478        */
479    }
480}
turbo_tasks/task/function.rs

turbo_tasks/task/
function.rs
