Skip to main content

turbo_tasks/task/
task_input.rs

1use std::{
2    collections::{BTreeMap, BTreeSet},
3    fmt::Debug,
4    future::Future,
5    hash::Hash,
6    ops::{Deref, DerefMut},
7    pin::Pin,
8    sync::Arc,
9    task::{Context, Poll},
10    time::Duration,
11};
12
13use anyhow::Result;
14use bincode::{
15    Decode, Encode,
16    de::Decoder,
17    enc::Encoder,
18    error::{DecodeError, EncodeError},
19};
20use either::Either;
21use smallvec::{Array, SmallVec};
22use turbo_frozenmap::{FrozenMap, FrozenSet};
23use turbo_rcstr::RcStr;
24use turbo_tasks_hash::HashAlgorithm;
25
26// The task_input macro expansion refers to the crate name directly in these unit tests.
27#[cfg(test)]
28use crate as turbo_tasks;
29use crate::{
30    DynTaskInputs, OrdResolvedVc, ReadRef, ResolvedVc, TaskId, TransientInstance, TransientValue,
31    ValueTypeId, Vc,
32};
33
34/// An 8-byte hand-rolled [`Future`] that immediately resolves to `Ok(self.clone())` of the
35/// referenced value.
36///
37/// Used by the [`TaskInput::resolve_input`] default implementation
38struct CloneReady<'a, T> {
39    pub inner: Option<&'a T>,
40}
41
42impl<'a, T: Clone> Future for CloneReady<'a, T> {
43    type Output = Result<T>;
44
45    fn poll(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Self::Output> {
46        Poll::Ready(Ok(self
47            .inner
48            .take()
49            .expect("future already polled to completion")
50            .clone()))
51    }
52}
53
54// `CloneReady` holds only a shared reference; it has no self-referential state.
55impl<'a, T> Unpin for CloneReady<'a, T> {}
56
57/// Trait to implement in order for a type to be accepted as a
58/// [`#[turbo_tasks::function]`][crate::function] argument.
59///
60/// ## Serialization
61///
62/// For persistent caching of a task, arguments must be serializable. All `TaskInput`s must
63/// implement the bincode [`Encode`] and [`Decode`] traits.
64///
65/// Transient task inputs are required to implement [`Encode`] and [`Decode`], but are allowed to
66/// panic at runtime. This requirement could be lifted in the future.
67///
68/// Bincode encoding must be deterministic and compatible with [`Eq`] comparisons. If two
69/// `TaskInput`s compare equal they must also encode to the same bytes.
70///
71/// ## Hash and Eq
72///
73/// Arguments are used as part of keys in a `HashMap`, so they must implement of [`PartialEq`],
74/// [`Eq`], and [`Hash`] traits.
75///
76/// ## [`Vc<T>`][Vc]
77///
78/// A [`Vc`] is a pointer to a cell. It implements `TaskInput` and serves as a "pass by reference"
79/// argument:
80///
81/// - **Memoization**: [`Vc`] is keyed by pointer for memoization purposes. Identical values in
82///   different cells are treated as distinct.
83/// - **Singleton Pattern**: To ensure memoization efficiency, the singleton pattern can be employed
84///   to guarantee that identical values yield the same `Vc`. For more info see [Singleton Pattern
85///   Guide][singleton].
86///
87/// [singleton]: https://turbopack-rust-docs.vercel.sh/turbo-engine/singleton.html
88///
89/// ## Deriving `TaskInput`
90///
91/// Structs or enums can be made into task inputs by deriving `TaskInput`:
92///
93/// ```rust
94/// # use turbo_tasks::macro_helpers::bincode::{Decode, Encode};
95/// #[turbo_tasks::task_input]
96/// #[derive(Clone, Debug, PartialEq, Eq, Hash, Encode, Decode)]
97/// struct MyStruct {
98///     // Fields go here...
99/// }
100/// ```
101///
102/// Derived `TaskInput` types **passed by value**. When called, arguments are moved into a `Box`,
103/// and then cloned before being passed into the function. If the task is invalidated, the
104/// `TaskInput` is cloned again to allow the function to be re-executed. It's recommended to ensure
105/// that these types are cheap to clone.
106///
107/// Reference-counted types like [`Arc`] are cheap to clone, but each reference contained in a
108/// `TaskInput` will be serialized independently in the persistent cache, and may consume extra disk
109/// space. If an [`Arc`] points to a large type, consider wrapping that type in [`Vc`], so that only
110/// one copy of the value will be serialized.
111pub trait TaskInput:
112    Send + Sync + Clone + Debug + PartialEq + Eq + Hash + Encode + Decode<()>
113{
114    /// This method should resolve any [`Vc`]s nested inside of this object, cloning the object in
115    /// the process. If the input is unresolved ([`TaskInput::is_resolved`]) a "local" resolution
116    /// task is created that runs this method.
117    fn resolve_input(&self) -> impl Future<Output = Result<Self>> + Send + '_ {
118        CloneReady { inner: Some(self) }
119    }
120
121    /// This should return `true` if there are any unresolved [`Vc`]s in the type.
122    ///
123    /// Note that [`Vc`]s can sometimes be internally resolved, so you should call
124    /// [`Vc::is_resolved`] (or rely on the derive macro for this trait) instead of returning `true`
125    /// for any [`Vc`]. [`ResolvedVc::is_resolved`] always returns `true`.
126    ///
127    /// If this returns `true`, a "local" resolution task calling [`TaskInput::resolve_input`] will
128    /// be spawned before the function accepting the arguments is run.
129    ///
130    /// If this returns `false`, the `TaskInput` will be [cloned][Clone] instead of resolved, and
131    /// the function's task will be spawned directly without a resolution step.
132    fn is_resolved(&self) -> bool {
133        true
134    }
135
136    /// This should return true if this object contains a [`Vc`] (or any subtype of [`Vc`]) pointing
137    /// to a cell owned by a transient task.
138    ///
139    /// Any function called with a transient `TaskInput` will be transient. Any [`Vc`] constructed
140    /// in a transient task or in a top-level [`run_once`][crate::run_once] closure will be
141    /// transient.
142    ///
143    /// Internally, a [`Vc`] can be determined to be transient by comparing the owning task's id
144    /// with the [`TRANSIENT_TASK_BIT`][crate::TRANSIENT_TASK_BIT] mask.
145    fn is_transient(&self) -> bool;
146}
147
148macro_rules! impl_task_input {
149    ($($t:ty),*) => {
150        $(
151            impl TaskInput for $t {
152                fn is_transient(&self) -> bool {
153                    false
154                }
155            }
156        )*
157    };
158}
159
160impl_task_input! {
161    (),
162    bool,
163    u8,
164    u16,
165    u32,
166    i32,
167    u64,
168    u128,
169    usize,
170    RcStr,
171    TaskId,
172    ValueTypeId,
173    Duration,
174    String,
175    HashAlgorithm
176}
177
178impl<T> TaskInput for Vec<T>
179where
180    T: TaskInput,
181{
182    fn is_resolved(&self) -> bool {
183        self.iter().all(TaskInput::is_resolved)
184    }
185
186    fn is_transient(&self) -> bool {
187        self.iter().any(TaskInput::is_transient)
188    }
189
190    async fn resolve_input(&self) -> Result<Self> {
191        let mut resolved = Vec::with_capacity(self.len());
192        for value in self {
193            resolved.push(value.resolve_input().await?);
194        }
195        Ok(resolved)
196    }
197}
198
199impl<A> TaskInput for SmallVec<A>
200where
201    A: Array + Send + Sync,
202    A::Item: TaskInput,
203    SmallVec<A>: Encode + Decode<()>,
204{
205    fn is_resolved(&self) -> bool {
206        self.iter().all(TaskInput::is_resolved)
207    }
208
209    fn is_transient(&self) -> bool {
210        self.iter().any(TaskInput::is_transient)
211    }
212
213    async fn resolve_input(&self) -> Result<Self> {
214        let mut resolved = SmallVec::with_capacity(self.len());
215        for value in self {
216            resolved.push(value.resolve_input().await?);
217        }
218        Ok(resolved)
219    }
220}
221
222impl<T> TaskInput for Box<T>
223where
224    T: TaskInput,
225{
226    fn is_resolved(&self) -> bool {
227        self.as_ref().is_resolved()
228    }
229
230    fn is_transient(&self) -> bool {
231        self.as_ref().is_transient()
232    }
233
234    async fn resolve_input(&self) -> Result<Self> {
235        Ok(Box::new(Box::pin(self.as_ref().resolve_input()).await?))
236    }
237}
238
239impl<T> TaskInput for Arc<T>
240where
241    T: TaskInput,
242{
243    fn is_resolved(&self) -> bool {
244        self.as_ref().is_resolved()
245    }
246
247    fn is_transient(&self) -> bool {
248        self.as_ref().is_transient()
249    }
250
251    async fn resolve_input(&self) -> Result<Self> {
252        Ok(Arc::new(Box::pin(self.as_ref().resolve_input()).await?))
253    }
254}
255
256impl<T> TaskInput for ReadRef<T>
257where
258    T: TaskInput,
259{
260    fn is_resolved(&self) -> bool {
261        Self::as_raw_ref(self).is_resolved()
262    }
263
264    fn is_transient(&self) -> bool {
265        Self::as_raw_ref(self).is_transient()
266    }
267
268    async fn resolve_input(&self) -> Result<Self> {
269        Ok(ReadRef::new_owned(
270            Box::pin(Self::as_raw_ref(self).resolve_input()).await?,
271        ))
272    }
273}
274
275impl<T> TaskInput for Option<T>
276where
277    T: TaskInput,
278{
279    fn is_resolved(&self) -> bool {
280        match self {
281            Some(value) => value.is_resolved(),
282            None => true,
283        }
284    }
285
286    fn is_transient(&self) -> bool {
287        match self {
288            Some(value) => value.is_transient(),
289            None => false,
290        }
291    }
292
293    async fn resolve_input(&self) -> Result<Self> {
294        match self {
295            Some(value) => Ok(Some(value.resolve_input().await?)),
296            None => Ok(None),
297        }
298    }
299}
300
301impl<T> TaskInput for Vc<T>
302where
303    T: Send + Sync + ?Sized,
304{
305    fn is_resolved(&self) -> bool {
306        Vc::is_resolved(*self)
307    }
308
309    fn is_transient(&self) -> bool {
310        self.node.is_transient()
311    }
312
313    fn resolve_input(&self) -> impl Future<Output = Result<Self>> + Send + '_ {
314        // It isn't ideal to use this function but it exactly matches this usecase (resolved but
315        // still a Vc)
316        (*self).resolve()
317    }
318}
319
320// `TaskInput` isn't needed/used for a bare `ResolvedVc`, as we'll expose `ResolvedVc` arguments as
321// `Vc`, but it is useful for structs that contain `ResolvedVc` and want to derive `TaskInput`.
322impl<T> TaskInput for ResolvedVc<T>
323where
324    T: Send + Sync + ?Sized,
325{
326    fn is_resolved(&self) -> bool {
327        true
328    }
329
330    fn is_transient(&self) -> bool {
331        self.node.is_transient()
332    }
333}
334
335impl<T> TaskInput for OrdResolvedVc<T>
336where
337    T: Send + Sync + ?Sized,
338{
339    fn is_resolved(&self) -> bool {
340        true
341    }
342
343    fn is_transient(&self) -> bool {
344        self.node.is_transient()
345    }
346}
347
348impl<T> TaskInput for TransientValue<T>
349where
350    T: DynTaskInputs + Clone + Debug + Hash + Eq + 'static,
351{
352    fn is_transient(&self) -> bool {
353        true
354    }
355}
356
357impl<T> Encode for TransientValue<T> {
358    fn encode<E: Encoder>(&self, _encoder: &mut E) -> Result<(), EncodeError> {
359        Err(EncodeError::Other("cannot encode transient task inputs"))
360    }
361}
362
363impl<Context, T> Decode<Context> for TransientValue<T> {
364    fn decode<D: Decoder<Context = Context>>(_decoder: &mut D) -> Result<Self, DecodeError> {
365        Err(DecodeError::Other("cannot decode transient task inputs"))
366    }
367}
368
369impl<T> TaskInput for TransientInstance<T>
370where
371    T: Sync + Send + 'static,
372{
373    fn is_transient(&self) -> bool {
374        true
375    }
376}
377
378impl<T> Encode for TransientInstance<T> {
379    fn encode<E: Encoder>(&self, _encoder: &mut E) -> Result<(), EncodeError> {
380        Err(EncodeError::Other("cannot encode transient task inputs"))
381    }
382}
383
384impl<Context, T> Decode<Context> for TransientInstance<T> {
385    fn decode<D: Decoder<Context = Context>>(_decoder: &mut D) -> Result<Self, DecodeError> {
386        Err(DecodeError::Other("cannot decode transient task inputs"))
387    }
388}
389
390impl<K, V> TaskInput for BTreeMap<K, V>
391where
392    K: TaskInput + Ord,
393    V: TaskInput,
394{
395    async fn resolve_input(&self) -> Result<Self> {
396        let mut new_map = BTreeMap::new();
397        for (k, v) in self {
398            new_map.insert(
399                TaskInput::resolve_input(k).await?,
400                TaskInput::resolve_input(v).await?,
401            );
402        }
403        Ok(new_map)
404    }
405
406    fn is_resolved(&self) -> bool {
407        self.iter()
408            .all(|(k, v)| TaskInput::is_resolved(k) && TaskInput::is_resolved(v))
409    }
410
411    fn is_transient(&self) -> bool {
412        self.iter()
413            .any(|(k, v)| TaskInput::is_transient(k) || TaskInput::is_transient(v))
414    }
415}
416
417impl<T> TaskInput for BTreeSet<T>
418where
419    T: TaskInput + Ord,
420{
421    async fn resolve_input(&self) -> Result<Self> {
422        let mut new_set = BTreeSet::new();
423        for value in self {
424            new_set.insert(TaskInput::resolve_input(value).await?);
425        }
426        Ok(new_set)
427    }
428
429    fn is_resolved(&self) -> bool {
430        self.iter().all(TaskInput::is_resolved)
431    }
432
433    fn is_transient(&self) -> bool {
434        self.iter().any(TaskInput::is_transient)
435    }
436}
437
438impl<K, V> TaskInput for FrozenMap<K, V>
439where
440    K: TaskInput + Ord + 'static,
441    V: TaskInput + 'static,
442{
443    async fn resolve_input(&self) -> Result<Self> {
444        let mut new_entries = Vec::with_capacity(self.len());
445        for (k, v) in self {
446            new_entries.push((
447                TaskInput::resolve_input(k).await?,
448                TaskInput::resolve_input(v).await?,
449            ));
450        }
451        // note: resolving might deduplicate `Vc`s in keys
452        Ok(Self::from(new_entries))
453    }
454
455    fn is_resolved(&self) -> bool {
456        self.iter()
457            .all(|(k, v)| TaskInput::is_resolved(k) && TaskInput::is_resolved(v))
458    }
459
460    fn is_transient(&self) -> bool {
461        self.iter()
462            .any(|(k, v)| TaskInput::is_transient(k) || TaskInput::is_transient(v))
463    }
464}
465
466impl<T> TaskInput for FrozenSet<T>
467where
468    T: TaskInput + Ord + 'static,
469{
470    async fn resolve_input(&self) -> Result<Self> {
471        let mut new_set = Vec::with_capacity(self.len());
472        for value in self {
473            new_set.push(TaskInput::resolve_input(value).await?);
474        }
475        Ok(Self::from_iter(new_set))
476    }
477
478    fn is_resolved(&self) -> bool {
479        self.iter().all(TaskInput::is_resolved)
480    }
481
482    fn is_transient(&self) -> bool {
483        self.iter().any(TaskInput::is_transient)
484    }
485}
486
487/// A thin wrapper around [`Either`] that implements the traits required by [`TaskInput`], notably
488/// [`Encode`] and [`Decode`].
489#[derive(Clone, Debug, PartialEq, Eq, Hash)]
490pub struct EitherTaskInput<L, R>(pub Either<L, R>);
491
492impl<L, R> Deref for EitherTaskInput<L, R> {
493    type Target = Either<L, R>;
494
495    fn deref(&self) -> &Self::Target {
496        &self.0
497    }
498}
499
500impl<L, R> DerefMut for EitherTaskInput<L, R> {
501    fn deref_mut(&mut self) -> &mut Self::Target {
502        &mut self.0
503    }
504}
505
506impl<L, R> Encode for EitherTaskInput<L, R>
507where
508    L: Encode,
509    R: Encode,
510{
511    fn encode<E: Encoder>(&self, encoder: &mut E) -> Result<(), EncodeError> {
512        turbo_bincode::either::encode(self, encoder)
513    }
514}
515
516impl<Context, L, R> Decode<Context> for EitherTaskInput<L, R>
517where
518    L: Decode<Context>,
519    R: Decode<Context>,
520{
521    fn decode<D: Decoder<Context = Context>>(decoder: &mut D) -> Result<Self, DecodeError> {
522        turbo_bincode::either::decode(decoder).map(Self)
523    }
524}
525
526impl<L, R> TaskInput for EitherTaskInput<L, R>
527where
528    L: TaskInput,
529    R: TaskInput,
530{
531    fn resolve_input(&self) -> impl Future<Output = Result<Self>> + Send + '_ {
532        self.as_ref().map_either(
533            async |l| anyhow::Ok(Self(Either::Left(l.resolve_input().await?))),
534            async |r| anyhow::Ok(Self(Either::Right(r.resolve_input().await?))),
535        )
536    }
537
538    fn is_resolved(&self) -> bool {
539        self.as_ref()
540            .either(TaskInput::is_resolved, TaskInput::is_resolved)
541    }
542
543    fn is_transient(&self) -> bool {
544        self.as_ref()
545            .either(TaskInput::is_transient, TaskInput::is_transient)
546    }
547}
548
549macro_rules! tuple_impls {
550    ( $( $name:ident )+ ) => {
551        impl<$($name: TaskInput),+> TaskInput for ($($name,)+)
552        where $($name: TaskInput),+
553        {
554            #[allow(non_snake_case)]
555            fn is_resolved(&self) -> bool {
556                let ($($name,)+) = self;
557                $($name.is_resolved() &&)+ true
558            }
559
560            #[allow(non_snake_case)]
561            fn is_transient(&self) -> bool {
562                let ($($name,)+) = self;
563                $($name.is_transient() ||)+ false
564            }
565
566            #[allow(non_snake_case)]
567            async fn resolve_input(&self) -> Result<Self> {
568                let ($($name,)+) = self;
569                Ok(($($name.resolve_input().await?,)+))
570            }
571        }
572    };
573}
574
575// Implement `TaskInput` for all tuples of 1 to 12 elements.
576tuple_impls! { A }
577tuple_impls! { A B }
578tuple_impls! { A B C }
579tuple_impls! { A B C D }
580tuple_impls! { A B C D E }
581tuple_impls! { A B C D E F }
582tuple_impls! { A B C D E F G }
583tuple_impls! { A B C D E F G H }
584tuple_impls! { A B C D E F G H I }
585tuple_impls! { A B C D E F G H I J }
586tuple_impls! { A B C D E F G H I J K }
587tuple_impls! { A B C D E F G H I J K L }
588
589#[cfg(test)]
590mod tests {
591    use turbo_rcstr::rcstr;
592
593    use super::*;
594
595    fn assert_task_input<T>(_: T)
596    where
597        T: TaskInput,
598    {
599    }
600
601    #[test]
602    fn test_no_fields() -> Result<()> {
603        #[turbo_tasks::task_input]
604        #[derive(Clone, Eq, PartialEq, Hash, Debug, Encode, Decode)]
605        struct NoFields;
606
607        assert_task_input(NoFields);
608        Ok(())
609    }
610
611    #[tokio::test]
612    async fn test_small_vec() -> Result<()> {
613        let inline: SmallVec<[RcStr; 1]> = smallvec::smallvec![rcstr!("a")];
614        assert_task_input(inline.clone());
615        assert!(inline.is_resolved());
616        assert!(!inline.is_transient());
617        let resolved = inline.resolve_input().await?;
618        assert_eq!(resolved, inline);
619        assert!(!resolved.spilled());
620
621        let spilled: SmallVec<[u32; 1]> = smallvec::smallvec![1, 2, 3];
622        assert!(spilled.is_resolved());
623        assert!(!spilled.is_transient());
624        assert_eq!(spilled.resolve_input().await?, spilled);
625        Ok(())
626    }
627
628    #[test]
629    fn test_one_unnamed_field() -> Result<()> {
630        #[turbo_tasks::task_input]
631        #[derive(Clone, Eq, PartialEq, Hash, Debug, Encode, Decode)]
632        struct OneUnnamedField(u32);
633
634        assert_task_input(OneUnnamedField(42));
635        Ok(())
636    }
637
638    #[test]
639    fn test_multiple_unnamed_fields() -> Result<()> {
640        #[turbo_tasks::task_input]
641        #[derive(Clone, Eq, PartialEq, Hash, Debug, Encode, Decode)]
642        struct MultipleUnnamedFields(u32, RcStr);
643
644        assert_task_input(MultipleUnnamedFields(42, rcstr!("42")));
645        Ok(())
646    }
647
648    #[test]
649    fn test_one_named_field() -> Result<()> {
650        #[turbo_tasks::task_input]
651        #[derive(Clone, Eq, PartialEq, Hash, Debug, Encode, Decode)]
652        struct OneNamedField {
653            named: u32,
654        }
655
656        assert_task_input(OneNamedField { named: 42 });
657        Ok(())
658    }
659
660    #[test]
661    fn test_multiple_named_fields() -> Result<()> {
662        #[turbo_tasks::task_input]
663        #[derive(Clone, Eq, PartialEq, Hash, Debug, Encode, Decode)]
664        struct MultipleNamedFields {
665            named: u32,
666            other: RcStr,
667        }
668
669        assert_task_input(MultipleNamedFields {
670            named: 42,
671            other: rcstr!("42"),
672        });
673        Ok(())
674    }
675
676    #[test]
677    fn test_generic_field() -> Result<()> {
678        #[turbo_tasks::task_input]
679        #[derive(Clone, Eq, PartialEq, Hash, Debug, Encode, Decode)]
680        struct GenericField<T>(T);
681
682        assert_task_input(GenericField(42));
683        assert_task_input(GenericField(rcstr!("42")));
684        Ok(())
685    }
686
687    #[turbo_tasks::task_input]
688    #[derive(Clone, Eq, PartialEq, Hash, Debug, Encode, Decode)]
689    enum OneVariant {
690        Variant,
691    }
692
693    #[test]
694    fn test_one_variant() -> Result<()> {
695        assert_task_input(OneVariant::Variant);
696        Ok(())
697    }
698
699    #[test]
700    fn test_multiple_variants() -> Result<()> {
701        #[turbo_tasks::task_input]
702        #[derive(Clone, PartialEq, Eq, Hash, Debug, Encode, Decode)]
703        enum MultipleVariants {
704            Variant1,
705            Variant2,
706        }
707
708        assert_task_input(MultipleVariants::Variant2);
709        Ok(())
710    }
711
712    #[turbo_tasks::task_input]
713    #[derive(Clone, Eq, PartialEq, Hash, Debug, Encode, Decode)]
714    enum MultipleVariantsAndHeterogeneousFields {
715        Variant1,
716        Variant2(u32),
717        Variant3 { named: u32 },
718        Variant4(u32, RcStr),
719        Variant5 { named: u32, other: RcStr },
720    }
721
722    #[test]
723    fn test_multiple_variants_and_heterogeneous_fields() -> Result<()> {
724        assert_task_input(MultipleVariantsAndHeterogeneousFields::Variant5 {
725            named: 42,
726            other: rcstr!("42"),
727        });
728        Ok(())
729    }
730
731    #[test]
732    fn test_nested_variants() -> Result<()> {
733        #[turbo_tasks::task_input]
734        #[derive(Clone, Eq, PartialEq, Hash, Debug, Encode, Decode)]
735        enum NestedVariants {
736            Variant1,
737            Variant2(MultipleVariantsAndHeterogeneousFields),
738            Variant3 { named: OneVariant },
739            Variant4(OneVariant, RcStr),
740            Variant5 { named: OneVariant, other: RcStr },
741        }
742
743        assert_task_input(NestedVariants::Variant5 {
744            named: OneVariant::Variant,
745            other: rcstr!("42"),
746        });
747        assert_task_input(NestedVariants::Variant2(
748            MultipleVariantsAndHeterogeneousFields::Variant5 {
749                named: 42,
750                other: rcstr!("42"),
751            },
752        ));
753        Ok(())
754    }
755}