turbo_tasks/task/
local_task.rs

1use std::{fmt, sync::Arc};
2
3use anyhow::{Result, anyhow};
4
5use crate::{
6    CellId, MagicAny, OutputContent, RawVc, TaskPersistence, TraitMethod, TurboTasksBackendApi,
7    ValueTypeId, backend::Backend, event::Event, macro_helpers::NativeFunction, registry,
8};
9
10/// A potentially in-flight local task stored in `CurrentGlobalTaskState::local_tasks`.
11pub enum LocalTask {
12    Scheduled { done_event: Event },
13    Done { output: OutputContent },
14}
15
16pub struct LocalTaskSpec {
17    /// The self value, will always be present for `ResolveTrait` tasks and is optional otherwise
18    pub(crate) this: Option<RawVc>,
19    /// Function arguments
20    pub(crate) arg: Box<dyn MagicAny>,
21    pub(crate) task_type: LocalTaskType,
22}
23
24#[derive(Copy, Clone)]
25pub enum LocalTaskType {
26    /// A resolve task, which resolves arguments and calls the function with resolve arguments. The
27    /// inner function call will be a `PersistentTaskType`.
28    ResolveNative { native_fn: &'static NativeFunction },
29
30    /// A trait method resolve task. It resolves the first (`self`) argument and looks up the trait
31    /// method on that value. Then it calls that method. The method call will do a cache lookup and
32    /// might resolve arguments before.
33    ResolveTrait { trait_method: &'static TraitMethod },
34}
35
36impl fmt::Display for LocalTaskType {
37    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
38        match self {
39            LocalTaskType::ResolveNative { native_fn } => write!(f, "*{}", native_fn.name),
40            LocalTaskType::ResolveTrait { trait_method } => write!(
41                f,
42                "*{}::{}",
43                trait_method.trait_name, trait_method.method_name
44            ),
45        }
46    }
47}
48
49impl LocalTaskType {
50    /// Implementation of the LocalTaskType::ResolveNative task.
51    /// Resolves all the task inputs and then calls the given function.
52    pub(crate) async fn run_resolve_native<B: Backend + 'static>(
53        native_fn: &'static NativeFunction,
54        mut this: Option<RawVc>,
55        arg: &dyn MagicAny,
56        persistence: TaskPersistence,
57        turbo_tasks: Arc<dyn TurboTasksBackendApi<B>>,
58    ) -> Result<RawVc> {
59        if let Some(this) = this.as_mut() {
60            *this = this.resolve().await?;
61        }
62        let arg = native_fn.arg_meta.resolve(arg).await?;
63        Ok(turbo_tasks.native_call(native_fn, this, arg, persistence))
64    }
65    /// Implementation of the LocalTaskType::ResolveTrait task.
66    pub(crate) async fn run_resolve_trait<B: Backend + 'static>(
67        trait_method: &'static TraitMethod,
68        this: RawVc,
69        arg: &dyn MagicAny,
70        persistence: TaskPersistence,
71        turbo_tasks: Arc<dyn TurboTasksBackendApi<B>>,
72    ) -> Result<RawVc> {
73        let this = this.resolve().await?;
74        let RawVc::TaskCell(_, CellId { type_id, .. }) = this else {
75            return Err(anyhow!("Trait method receiver must be a cell"));
76        };
77
78        let native_fn = Self::resolve_trait_method_from_value(trait_method, type_id)?;
79        let arg = native_fn.arg_meta.filter_and_resolve(arg).await?;
80        Ok(turbo_tasks.native_call(native_fn, Some(this), arg, persistence))
81    }
82
83    fn resolve_trait_method_from_value(
84        trait_method: &'static TraitMethod,
85        value_type: ValueTypeId,
86    ) -> Result<&'static NativeFunction> {
87        match registry::get_value_type(value_type).get_trait_method(trait_method) {
88            Some(native_fn) => Ok(native_fn),
89            None => Err(anyhow!(
90                "{} doesn't implement the trait for {:?}, the compiler should have flagged this",
91                registry::get_value_type(value_type),
92                trait_method
93            )),
94        }
95    }
96}
97
98#[cfg(test)]
99pub(crate) mod tests {
100    use super::*;
101    use crate::{self as turbo_tasks, Vc};
102
103    #[turbo_tasks::function]
104    fn mock_func_task() -> Vc<()> {
105        Vc::cell(())
106    }
107
108    #[turbo_tasks::value_trait]
109    trait MockTrait {
110        #[turbo_tasks::function]
111        fn mock_method_task() -> Vc<()>;
112    }
113
114    #[test]
115    fn test_fmt() {
116        assert_eq!(
117            LocalTaskType::ResolveNative {
118                native_fn: &MOCK_FUNC_TASK_FUNCTION,
119            }
120            .to_string(),
121            "*mock_func_task",
122        );
123        assert_eq!(
124            LocalTaskType::ResolveTrait {
125                trait_method: MOCKTRAIT_TRAIT_TYPE.get("mock_method_task"),
126            }
127            .to_string(),
128            "*MockTrait::mock_method_task",
129        );
130    }
131}