turbo_tasks/task/
local_task.rs

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

turbo_tasks/task/
local_task.rs
