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.downcast_ref::<T>().cloned();
176    #[cfg(debug_assertions)]
177    return anyhow::Context::with_context(value, || {
178        crate::native_function::debug_downcast_args_error_msg(std::any::type_name::<T>(), arg)
179    });
180    #[cfg(not(debug_assertions))]
181    return anyhow::Context::context(value, "Invalid argument type");
182}
183
184// Helper function for `task_fn_impl!()`
185async fn output_try_into_non_local_raw_vc(output: impl TaskOutput) -> Result<RawVc> {
186    // TODO: Potential future optimization: If we know we're inside a local task, we can avoid
187    // calling `to_non_local()` here, which might let us avoid constructing a non-local cell for the
188    // local task's return value. Flattening chains of `RawVc::LocalOutput` may still be useful to
189    // reduce traversal later.
190    output.try_into_raw_vc()?.to_non_local().await
191}
192
193macro_rules! task_fn_impl {
194    ( $async_fn_trait:ident $arg_len:literal $( $arg:ident )* ) => {
195        impl<F, Output, $($arg,)*> TaskFnInputFunction<FunctionMode, ($($arg,)*)> for F
196        where
197            $($arg: TaskInput + 'static,)*
198            F: Fn($($arg,)*) -> Output + Send + Sync + Clone + 'static,
199            Output: TaskOutput + 'static,
200        {
201            #[allow(non_snake_case)]
202            fn functor(&self, arg: &dyn MagicAny) -> Result<NativeTaskFuture> {
203                let task_fn = self.clone();
204                let ($($arg,)*) = get_args::<($($arg,)*)>(arg)?;
205                Ok(Box::pin(async move {
206                    let output = (task_fn)($($arg,)*);
207                    output_try_into_non_local_raw_vc(output).await
208                }))
209            }
210        }
211
212        impl<F, Output, FutureOutput, $($arg,)*> TaskFnInputFunction<AsyncFunctionMode, ($($arg,)*)> for F
213        where
214            $($arg: TaskInput + 'static,)*
215            F: Fn($($arg,)*) -> FutureOutput + Send + Sync + Clone + 'static,
216            FutureOutput: Future<Output = Output> + Send + 'static,
217            Output: TaskOutput + 'static,
218        {
219            #[allow(non_snake_case)]
220            fn functor(&self, arg: &dyn MagicAny) -> Result<NativeTaskFuture> {
221                let task_fn = self.clone();
222                let ($($arg,)*) = get_args::<($($arg,)*)>(arg)?;
223                Ok(Box::pin(async move {
224                    let output = (task_fn)($($arg,)*).await;
225                    output_try_into_non_local_raw_vc(output).await
226                }))
227            }
228        }
229
230        impl<F, Output, Recv, $($arg,)*> TaskFnInputFunctionWithThis<MethodMode, Recv, ($($arg,)*)> for F
231        where
232            Recv: VcValueType,
233            $($arg: TaskInput + 'static,)*
234            F: Fn(&Recv, $($arg,)*) -> Output + Send + Sync + Clone + 'static,
235            Output: TaskOutput + 'static,
236        {
237            #[allow(non_snake_case)]
238            fn functor(&self, this: RawVc, arg: &dyn MagicAny) -> Result<NativeTaskFuture> {
239                let task_fn = self.clone();
240                let recv = Vc::<Recv>::from(this);
241                let ($($arg,)*) = get_args::<($($arg,)*)>(arg)?;
242                Ok(Box::pin(async move {
243                    let recv = recv.await?;
244                    let recv = <Recv::Read as VcRead<Recv>>::target_to_value_ref(&*recv);
245                    let output = (task_fn)(recv, $($arg,)*);
246                    output_try_into_non_local_raw_vc(output).await
247                }))
248            }
249        }
250
251        impl<F, Output, Recv, $($arg,)*> TaskFnInputFunctionWithThis<FunctionMode, Recv, ($($arg,)*)> for F
252        where
253            Recv: Sync + Send + 'static,
254            $($arg: TaskInput + 'static,)*
255            F: Fn(Vc<Recv>, $($arg,)*) -> Output + Send + Sync + Clone + 'static,
256            Output: TaskOutput + 'static,
257        {
258            #[allow(non_snake_case)]
259            fn functor(&self, this: RawVc, arg: &dyn MagicAny) -> Result<NativeTaskFuture> {
260                let task_fn = self.clone();
261                let recv = Vc::<Recv>::from(this);
262                let ($($arg,)*) = get_args::<($($arg,)*)>(arg)?;
263                Ok(Box::pin(async move {
264                    let output = (task_fn)(recv, $($arg,)*);
265                    output_try_into_non_local_raw_vc(output).await
266                }))
267            }
268        }
269
270        pub trait $async_fn_trait<A0, $($arg,)*>: Fn(A0, $($arg,)*) -> Self::OutputFuture {
271            type OutputFuture: Future<Output = <Self as $async_fn_trait<A0, $($arg,)*>>::Output> + Send;
272            type Output: TaskOutput;
273        }
274
275        impl<F: ?Sized, Fut, A0, $($arg,)*> $async_fn_trait<A0, $($arg,)*> for F
276        where
277            F: Fn(A0, $($arg,)*) -> Fut,
278            Fut: Future + Send,
279            Fut::Output: TaskOutput + 'static
280        {
281            type OutputFuture = Fut;
282            type Output = Fut::Output;
283        }
284
285        impl<F, Recv, $($arg,)*> TaskFnInputFunctionWithThis<AsyncMethodMode, Recv, ($($arg,)*)> for F
286        where
287            Recv: VcValueType,
288            $($arg: TaskInput + 'static,)*
289            F: for<'a> $async_fn_trait<&'a Recv, $($arg,)*> + Clone + Send + Sync + 'static,
290        {
291            #[allow(non_snake_case)]
292            fn functor(&self, this: RawVc, arg: &dyn MagicAny) -> Result<NativeTaskFuture> {
293                let task_fn = self.clone();
294                let recv = Vc::<Recv>::from(this);
295                let ($($arg,)*) = get_args::<($($arg,)*)>(arg)?;
296                Ok(Box::pin(async move {
297                    let recv = recv.await?;
298                    let recv = <Recv::Read as VcRead<Recv>>::target_to_value_ref(&*recv);
299                    let output = (task_fn)(recv, $($arg,)*).await;
300                    output_try_into_non_local_raw_vc(output).await
301                }))
302            }
303        }
304
305        impl<F, Recv, $($arg,)*> TaskFnInputFunctionWithThis<AsyncFunctionMode, Recv, ($($arg,)*)> for F
306        where
307            Recv: Sync + Send + 'static,
308            $($arg: TaskInput + 'static,)*
309            F: $async_fn_trait<Vc<Recv>, $($arg,)*> + Clone + Send + Sync + 'static,
310        {
311            #[allow(non_snake_case)]
312            fn functor(&self, this: RawVc, arg: &dyn MagicAny) -> Result<NativeTaskFuture> {
313                let task_fn = self.clone();
314                let recv = Vc::<Recv>::from(this);
315                let ($($arg,)*) = get_args::<($($arg,)*)>(arg)?;
316                Ok(Box::pin(async move {
317                    let output = (task_fn)(recv, $($arg,)*).await;
318                    output_try_into_non_local_raw_vc(output).await
319                }))
320            }
321        }
322    };
323}
324
325task_fn_impl! { AsyncFn0 0 }
326task_fn_impl! { AsyncFn1 1 A1 }
327task_fn_impl! { AsyncFn2 2 A1 A2 }
328task_fn_impl! { AsyncFn3 3 A1 A2 A3 }
329task_fn_impl! { AsyncFn4 4 A1 A2 A3 A4 }
330task_fn_impl! { AsyncFn5 5 A1 A2 A3 A4 A5 }
331task_fn_impl! { AsyncFn6 6 A1 A2 A3 A4 A5 A6 }
332task_fn_impl! { AsyncFn7 7 A1 A2 A3 A4 A5 A6 A7 }
333task_fn_impl! { AsyncFn8 8 A1 A2 A3 A4 A5 A6 A7 A8 }
334task_fn_impl! { AsyncFn9 9 A1 A2 A3 A4 A5 A6 A7 A8 A9 }
335task_fn_impl! { AsyncFn10 10 A1 A2 A3 A4 A5 A6 A7 A8 A9 A10 }
336task_fn_impl! { AsyncFn11 11 A1 A2 A3 A4 A5 A6 A7 A8 A9 A10 A11 }
337task_fn_impl! { AsyncFn12 12 A1 A2 A3 A4 A5 A6 A7 A8 A9 A10 A11 A12 }
338
339// There needs to be one more implementation than task_fn_impl to account for
340// the receiver.
341task_inputs_impl! {}
342task_inputs_impl! { A1 }
343task_inputs_impl! { A1 A2 }
344task_inputs_impl! { A1 A2 A3 }
345task_inputs_impl! { A1 A2 A3 A4 }
346task_inputs_impl! { A1 A2 A3 A4 A5 }
347task_inputs_impl! { A1 A2 A3 A4 A5 A6 }
348task_inputs_impl! { A1 A2 A3 A4 A5 A6 A7 }
349task_inputs_impl! { A1 A2 A3 A4 A5 A6 A7 A8 }
350task_inputs_impl! { A1 A2 A3 A4 A5 A6 A7 A8 A9 }
351task_inputs_impl! { A1 A2 A3 A4 A5 A6 A7 A8 A9 A10 }
352task_inputs_impl! { A1 A2 A3 A4 A5 A6 A7 A8 A9 A10 A11 }
353task_inputs_impl! { A1 A2 A3 A4 A5 A6 A7 A8 A9 A10 A11 A12 }
354task_inputs_impl! { A1 A2 A3 A4 A5 A6 A7 A8 A9 A10 A11 A12 A13 }
355
356#[cfg(test)]
357mod tests {
358    use turbo_rcstr::RcStr;
359
360    use super::*;
361    use crate::{ShrinkToFit, VcCellNewMode, VcDefaultRead};
362
363    #[test]
364    fn test_task_fn() {
365        fn no_args() -> crate::Vc<i32> {
366            todo!()
367        }
368
369        fn one_arg(_a: i32) -> crate::Vc<i32> {
370            todo!()
371        }
372
373        async fn async_one_arg(_a: i32) -> crate::Vc<i32> {
374            todo!()
375        }
376
377        fn with_recv(_a: &i32) -> crate::Vc<i32> {
378            todo!()
379        }
380
381        async fn async_with_recv(_a: &i32) -> crate::Vc<i32> {
382            todo!()
383        }
384
385        fn with_recv_and_str(_a: &i32, _s: RcStr) -> crate::Vc<i32> {
386            todo!()
387        }
388
389        async fn async_with_recv_and_str(_a: &i32, _s: RcStr) -> crate::Vc<i32> {
390            todo!()
391        }
392
393        async fn async_with_recv_and_str_and_result(_a: &i32, _s: RcStr) -> Result<crate::Vc<i32>> {
394            todo!()
395        }
396
397        fn accepts_task_fn<F>(_task_fn: F)
398        where
399            F: TaskFn,
400        {
401        }
402
403        struct Struct;
404        impl Struct {
405            async fn inherent_method(&self) {}
406        }
407
408        impl ShrinkToFit for Struct {
409            fn shrink_to_fit(&mut self) {}
410        }
411
412        unsafe impl VcValueType for Struct {
413            type Read = VcDefaultRead<Struct>;
414
415            type CellMode = VcCellNewMode<Struct>;
416
417            fn get_value_type_id() -> crate::ValueTypeId {
418                todo!()
419            }
420        }
421
422        trait AsyncTrait {
423            async fn async_method(&self);
424        }
425
426        impl AsyncTrait for Struct {
427            async fn async_method(&self) {
428                todo!()
429            }
430        }
431
432        /*
433        async fn async_with_recv_and_str_and_lf(
434            _a: &i32,
435            _s: String,
436        ) -> Result<crate::Vc<i32>, crate::Vc<i32>> {
437            todo!()
438        }
439
440        #[async_trait::async_trait]
441        trait BoxAsyncTrait {
442            async fn box_async_method(&self);
443        }
444
445        #[async_trait::async_trait]
446        impl BoxAsyncTrait for Struct {
447            async fn box_async_method(&self) {
448                todo!()
449            }
450        }
451        */
452
453        let _task_fn = no_args.into_task_fn();
454        accepts_task_fn(no_args.into_task_fn());
455        let _task_fn = one_arg.into_task_fn();
456        accepts_task_fn(one_arg.into_task_fn());
457        let _task_fn = async_one_arg.into_task_fn();
458        accepts_task_fn(async_one_arg.into_task_fn());
459        let task_fn = with_recv.into_task_fn_with_this();
460        accepts_task_fn(task_fn);
461        let task_fn = async_with_recv.into_task_fn_with_this();
462        accepts_task_fn(task_fn);
463        let task_fn = with_recv_and_str.into_task_fn_with_this();
464        accepts_task_fn(task_fn);
465        let task_fn = async_with_recv_and_str.into_task_fn_with_this();
466        accepts_task_fn(task_fn);
467        let task_fn = async_with_recv_and_str_and_result.into_task_fn_with_this();
468        accepts_task_fn(task_fn);
469        let task_fn = <Struct as AsyncTrait>::async_method.into_task_fn_with_this();
470        accepts_task_fn(task_fn);
471        let task_fn = Struct::inherent_method.into_task_fn_with_this();
472        accepts_task_fn(task_fn);
473
474        /*
475        let task_fn = <Struct as BoxAsyncTrait>::box_async_method.into_task_fn();
476        accepts_task_fn(task_fn);
477        let task_fn = async_with_recv_and_str_and_lf.into_task_fn();
478        accepts_task_fn(task_fn);
479        */
480    }
481}
turbo_tasks/task/function.rs

turbo_tasks/task/
function.rs
