Skip to main content

turbo_tasks/
value_type.rs

1use std::{
2    any::TypeId,
3    cell::SyncUnsafeCell,
4    fmt::{self, Debug, Display, Formatter},
5    hash::Hash,
6};
7
8use bincode::{Decode, Encode};
9use tracing::Span;
10use turbo_bincode::{AnyDecodeFn, AnyEncodeFn};
11
12use crate::{
13    RawVc, SharedReference, TaskPriority, VcValueType,
14    id::TraitTypeId,
15    macro_helpers::{CollectableTraitMethods, NativeFunction},
16    magic_any::any_as_encode,
17    registry::{RegistryType, get_trait_type_id, trait_type_count, turbo_registry},
18    task::shared_reference::TypedSharedReference,
19    vc::VcCellMode,
20};
21
22type RawCellFactoryFn = fn(TypedSharedReference) -> RawVc;
23type Vtable = &'static [&'static NativeFunction];
24
25// TODO this type need some refactoring when multiple languages are added to
26// turbo-task In this case a trait_method might be of a different function type.
27// It probably need to be a Vc<Function>.
28// That's also needed in a distributed world, where the function might be only
29// available on a remote instance.
30
31/// A definition of a type of data.
32///
33/// Contains a list of traits and trait methods that are available on that type.
34pub struct ValueType {
35    pub ty: RegistryType,
36
37    /// Functions to convert to write the type to a buffer or read it from a buffer.
38    pub bincode: Option<(AnyEncodeFn, AnyDecodeFn<SharedReference>)>,
39
40    /// An implementation of
41    /// [`VcCellMode::raw_cell`][crate::vc::VcCellMode::raw_cell].
42    ///
43    /// Allows dynamically constructing a cell using the type id. Used inside of
44    /// [`TraitRef`][crate::TraitRef] where we have a type id, but not the concrete type `T` of
45    /// `Vc<T>`.
46    ///
47    /// Because we allow resolving `Vc<dyn Trait>`, it's otherwise not possible
48    /// for `RawVc` to know what the appropriate `VcCellMode` is.
49    pub(crate) raw_cell: RawCellFactoryFn,
50
51    traits: SyncUnsafeCell<ValueTypeTraits>,
52}
53
54impl Debug for ValueType {
55    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
56        f.debug_struct("ValueType")
57            .field("name", &self.ty.name)
58            .finish()
59    }
60}
61
62impl Display for ValueType {
63    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
64        f.write_str(self.ty.name)
65    }
66}
67
68struct ValueTypeTraits {
69    /// Flat array indexed by TraitTypeId (1-based, so index 0 = TraitTypeId 1).
70    /// `None` means this value type does not implement that trait.
71    /// The outer Option is None before init, Some after.
72    traits: Option<Box<[Option<Vtable>]>>,
73}
74
75pub trait ManualEncodeWrapper: Encode {
76    type Value;
77
78    // this uses RPIT to avoid some lifetime problems
79    fn new<'a>(value: &'a Self::Value) -> impl Encode + 'a;
80}
81
82pub trait ManualDecodeWrapper: Decode<()> {
83    type Value;
84
85    fn inner(self) -> Self::Value;
86}
87
88impl ValueType {
89    /// This is internally used by [`#[turbo_tasks::value]`][crate::value].
90    pub const fn new<T: VcValueType>(global_name: &'static str) -> Self {
91        Self::new_inner::<T>(global_name, None)
92    }
93
94    /// This is internally used by [`#[turbo_tasks::value]`][crate::value].
95    pub const fn new_with_bincode<T: VcValueType + Encode + Decode<()>>(
96        global_name: &'static str,
97    ) -> Self {
98        Self::new_inner::<T>(
99            global_name,
100            Some((
101                |this, enc| {
102                    T::encode(any_as_encode::<T>(this), enc)?;
103                    Ok(())
104                },
105                |dec| {
106                    let val = T::decode(dec)?;
107                    Ok(SharedReference::new(triomphe::Arc::new(val)))
108                },
109            )),
110        )
111    }
112
113    /// This is used internally by [`turbo_tasks_macros::primitive`] to encode/decode foreign types
114    /// that cannot implement the [`bincode`] traits due to the [orphan rules].
115    ///
116    /// This is done by constructing wrapper types that implement the bincode traits on behalf of
117    /// the wrapped type.
118    ///
119    /// [orphan rules]: https://doc.rust-lang.org/reference/items/implementations.html#orphan-rules
120    pub const fn new_with_bincode_wrappers<
121        T: VcValueType,
122        E: ManualEncodeWrapper<Value = T>,
123        D: ManualDecodeWrapper<Value = T>,
124    >(
125        global_name: &'static str,
126    ) -> Self {
127        Self::new_inner::<T>(
128            global_name,
129            Some((
130                |this, enc| {
131                    E::new(any_as_encode::<T>(this)).encode(enc)?;
132                    Ok(())
133                },
134                |dec| {
135                    let val = D::inner(D::decode(dec)?);
136                    Ok(SharedReference::new(triomphe::Arc::new(val)))
137                },
138            )),
139        )
140    }
141
142    // Helper for other constructor functions
143    const fn new_inner<T: VcValueType>(
144        global_name: &'static str,
145        bincode: Option<(AnyEncodeFn, AnyDecodeFn<SharedReference>)>,
146    ) -> Self {
147        Self {
148            ty: RegistryType::new::<T>(std::any::type_name::<T>(), global_name),
149            bincode,
150            raw_cell: <T::CellMode as VcCellMode<T>>::raw_cell,
151            traits: SyncUnsafeCell::new(ValueTypeTraits { traits: None }),
152        }
153    }
154
155    /// Returns the TypeId of the concrete type this ValueType represents.
156    pub fn type_id(&self) -> TypeId {
157        self.ty.type_id
158    }
159
160    #[inline]
161    fn trait_info(&self) -> &ValueTypeTraits {
162        // SAFETY: Written during single-threaded Lazy init, read-only after.
163        unsafe { &*self.traits.get() }
164    }
165
166    #[inline]
167    pub fn get_trait_method(
168        &self,
169        trait_method: &'static TraitMethod,
170    ) -> Option<&'static NativeFunction> {
171        let trait_type_id = trait_method.trait_type_id();
172        let vtable = self.trait_info().traits.as_ref()?[*trait_type_id as usize - 1]?;
173        Some(vtable[trait_method.index as usize])
174    }
175
176    fn register_trait(&self, trait_type: &'static TraitType, trait_methods: Vtable) {
177        // SAFETY: Called only during single-threaded registry init
178        let traits = unsafe { &mut *self.traits.get() };
179        let trait_type_id = get_trait_type_id(trait_type);
180        let array = traits
181            .traits
182            .get_or_insert_with(|| vec![None; trait_type_count()].into_boxed_slice());
183        array[*trait_type_id as usize - 1] = Some(trait_methods);
184    }
185
186    #[inline]
187    pub fn has_trait(&self, trait_type: &TraitTypeId) -> bool {
188        self.trait_info()
189            .traits
190            .as_ref()
191            .is_some_and(|t| t[**trait_type as usize - 1].is_some())
192    }
193}
194
195turbo_registry!("Value", ValueType);
196
197// Called during ValueType registry post_init to register all trait methods.
198// Single-threaded during Lazy init.
199pub(crate) fn register_all_trait_methods(_: &[&'static ValueType]) {
200    for entry in inventory::iter::<CollectableTraitMethods> {
201        entry
202            .value_type
203            .register_trait(entry.trait_type, entry.methods)
204    }
205}
206
207pub struct TraitMethod {
208    pub trait_type: &'static TraitType,
209    pub index: u8,
210    pub trait_name: &'static str,
211    pub method_name: &'static str,
212    pub default_method: Option<&'static NativeFunction>,
213}
214impl Hash for TraitMethod {
215    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
216        (self as *const TraitMethod).hash(state);
217    }
218}
219
220impl Eq for TraitMethod {}
221
222impl PartialEq for TraitMethod {
223    fn eq(&self, other: &Self) -> bool {
224        std::ptr::eq(self, other)
225    }
226}
227impl Debug for TraitMethod {
228    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
229        f.debug_struct("TraitMethod")
230            .field("trait_name", &self.trait_name)
231            .field("name", &self.method_name)
232            .field("default_method", &self.default_method)
233            .finish()
234    }
235}
236impl TraitMethod {
237    /// Returns the TraitTypeId by reading directly from the trait type's registry entry.
238    /// Must only be called after registry init.
239    #[inline]
240    fn trait_type_id(&self) -> TraitTypeId {
241        // SAFETY: Written during single-threaded Lazy init. Lazy provides acquire barrier.
242        let raw = unsafe { std::ptr::read(self.trait_type.ty.id.get()) };
243        debug_assert!(raw != 0, "TraitMethod::trait_type_id not initialized");
244        unsafe { TraitTypeId::new_unchecked(raw) }
245    }
246
247    pub(crate) fn resolve_span(&self, priority: TaskPriority) -> Span {
248        tracing::trace_span!(
249            "turbo_tasks::resolve_trait_call",
250            name = format_args!("{}::{}", &self.trait_name, &self.method_name),
251            priority = %priority,
252        )
253    }
254}
255
256pub struct TraitType {
257    pub ty: RegistryType,
258    pub methods: phf::Map<&'static str, TraitMethod>,
259    pub method_names: &'static [&'static str],
260    pub default_methods: &'static [Option<&'static NativeFunction>],
261}
262
263impl Debug for TraitType {
264    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
265        let mut d = f.debug_struct("TraitType");
266        d.field("name", &self.ty.name);
267        for (name, method) in self.methods.entries() {
268            d.field(name, method);
269        }
270        d.finish()
271    }
272}
273
274impl Display for TraitType {
275    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
276        write!(f, "trait {}", self.ty.name)
277    }
278}
279
280impl TraitType {
281    pub const fn new<T: 'static>(
282        name: &'static str,
283        global_name: &'static str,
284        methods: phf::Map<&'static str, TraitMethod>,
285        method_names: &'static [&'static str],
286        default_methods: &'static [Option<&'static NativeFunction>],
287    ) -> Self {
288        Self {
289            ty: RegistryType::new::<T>(name, global_name),
290            methods,
291            method_names,
292            default_methods,
293        }
294    }
295
296    pub fn get(&self, name: &str) -> &TraitMethod {
297        self.methods.get(name).unwrap()
298    }
299}
300
301turbo_registry!("Trait", TraitType);
302
303pub trait TraitVtablePrototype {
304    const LEN: usize;
305    const NAMES: &'static [&'static str];
306    const DEFAULTS: &'static [Option<&'static NativeFunction>];
307}
308
309pub(crate) const fn index_of_name(array: &'static [&'static str], name: &'static str) -> usize {
310    let mut i = 0;
311    'outer: while i < array.len() {
312        if array[i].len() == name.len() {
313            let mut j = 0;
314            while j < name.len() {
315                if array[i].as_bytes()[j] != name.as_bytes()[j] {
316                    i += 1;
317                    continue 'outer;
318                }
319                j += 1;
320            }
321            return i;
322        }
323        i += 1;
324    }
325    panic!("Method not found!")
326}
327
328pub const fn build_trait_vtable<B: TraitVtablePrototype, const LEN: usize>(
329    overrides: &[(&'static str, &'static NativeFunction)],
330) -> [&'static NativeFunction; LEN] {
331    let mut methods = [&crate::native_function::VTABLE_DEFAULT; LEN];
332    let mut i = 0;
333    while i < LEN {
334        if let Some(default) = B::DEFAULTS[i] {
335            methods[i] = default;
336        }
337        i += 1;
338    }
339    // N*M scan where N = overrides, M = method names. Both are small (single digits).
340    let mut i = 0;
341    while i < overrides.len() {
342        let (name, f) = overrides[i];
343        methods[index_of_name(B::NAMES, name)] = f;
344        i += 1;
345    }
346    methods
347}