turbo_tasks/

id.rs

use std::{
    fmt::{Debug, Display},
    mem::transmute_copy,
    num::{NonZero, NonZeroU64, TryFromIntError},
    ops::Deref,
};

use bincode::{
    Decode, Encode,
    de::Decoder,
    enc::Encoder,
    error::{DecodeError, EncodeError},
    impl_borrow_decode,
};
use serde::{Deserialize, Serialize, de::Visitor};

use crate::{TaskPersistence, registry};

macro_rules! define_id {
    (
        $name:ident : $primitive:ty
        $(,max = $max:expr)?
        $(,derive($($derive:ty),*))?
        $(,serde($serde:tt))?
        $(,doc = $doc:literal)*
        $(,)?
    ) => {
        $(#[doc = $doc])*
        #[derive(Hash, Clone, Copy, PartialEq, Eq, PartialOrd, Ord $($(,$derive)*)? )]
        $(#[serde($serde)])?
        pub struct $name {
            id: NonZero<$primitive>,
        }

        impl $name {
            pub const MIN: Self = Self { id: NonZero::<$primitive>::MIN };
            // `max` defaults to the primitive's max; types packed into a smaller
            // bit field (e.g. `TaskId` in `RawVc`) override it.
            pub const MAX: Self = {
                let _max: $primitive = NonZero::<$primitive>::MAX.get();
                $( let _max: $primitive = $max; )?
                // SAFETY: `_max` is either `NonZero::MAX` or a caller-provided
                // positive constant; both are non-zero.
                Self { id: unsafe { NonZero::<$primitive>::new_unchecked(_max) } }
            };

            /// Constructs a wrapper type from the numeric identifier.
            ///
            /// # Safety
            ///
            /// The passed `id` must not be zero.
            pub const unsafe fn new_unchecked(id: $primitive) -> Self {
                Self { id: unsafe { NonZero::<$primitive>::new_unchecked(id) } }
            }
            /// Constructs a wrapper type from the numeric identifier.
            ///
            /// Returns `None` if the provided `id` is zero, otherwise returns
            /// `Some(Self)` containing the wrapped non-zero identifier.
            pub fn new(id: $primitive) -> Option<Self> {
                NonZero::<$primitive>::new(id).map(|id| Self{id})
            }
            /// Allows `const` conversion to a [`NonZeroU64`], useful with
            /// [`crate::id_factory::IdFactory::new_const`].
            pub const fn to_non_zero_u64(self) -> NonZeroU64 {
                const {
                    assert!(<$primitive>::BITS <= u64::BITS);
                }
                unsafe { NonZeroU64::new_unchecked(self.id.get() as u64) }
            }
            /// Allows `const` conversion to [`NonZero<$primitive>`]
            pub const fn to_non_zero_primitive(self) -> NonZero<$primitive> {
                self.id
            }
            pub const fn to_primitive(self) -> $primitive {
                self.id.get()
            }
        }

        impl Display for $name {
            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                write!(f, concat!(stringify!($name), " {}"), self.id)
            }
        }

        impl Deref for $name {
            type Target = $primitive;

            fn deref(&self) -> &Self::Target {
                // SAFETY: `NonZero<T>` is guaranteed to have the same layout as `T`
                unsafe { transmute_copy(&&self.id) }
            }
        }

        define_id!(@impl_try_from_primitive_conversion $name $primitive);

        impl From<NonZero<$primitive>> for $name {
            fn from(id: NonZero::<$primitive>) -> Self {
                Self {
                    id,
                }
            }
        }

        impl From<$name> for NonZeroU64 {
            fn from(id: $name) -> Self {
                id.to_non_zero_u64()
            }
        }
    };
    (
        @impl_try_from_primitive_conversion $name:ident u64
    ) => {
        // we get a `TryFrom` blanket impl for free via the `From` impl
    };
    (
        @impl_try_from_primitive_conversion $name:ident $primitive:ty
    ) => {
        impl TryFrom<$primitive> for $name {
            type Error = TryFromIntError;

            fn try_from(id: $primitive) -> Result<Self, Self::Error> {
                Ok(Self {
                    id: NonZero::try_from(id)?
                })
            }
        }

        impl TryFrom<NonZeroU64> for $name {
            type Error = TryFromIntError;

            fn try_from(id: NonZeroU64) -> Result<Self, Self::Error> {
                Ok(Self { id: NonZero::try_from(id)? })
            }
        }
    };
}

define_id!(
    TaskId: u32,
    // Capped below `u32::MAX` so the id fits in 31 bits when packed into `RawVc`.
    max = TASK_ID_MAX,
    derive(Deserialize, Decode),
    serde(transparent),
);
define_id!(
    ValueTypeId: u16,
    // Capped below `u16::MAX` so the id fits in 10 bits when packed into `CellId`.
    max = crate::CellId::MAX_VALUE_TYPE_ID,
);
define_id!(FunctionId: u16);
define_id!(TraitTypeId: u16);
define_id!(
    LocalTaskId: u32,
    derive(Debug, Serialize, Deserialize, Encode, Decode),
    serde(transparent),
    doc = "Represents the nth `local` function call inside a task.",
);
define_id!(
    ExecutionId: u16,
    derive(Debug, Serialize, Deserialize, Encode, Decode),
    serde(transparent),
    doc = "An identifier for a specific task execution. Used to assert that local `Vc`s don't \
        leak. This value may overflow and re-use old values.",
);

// Preserve transparent serde and the derived NonZero<u32> bincode format, but reject references
// to transient tasks. RawVc packs TaskId into its own integer and checks that path separately.
impl Serialize for TaskId {
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        debug_assert!(
            !self.is_transient(),
            "transient TaskId must not be serialized"
        );
        self.id.serialize(serializer)
    }
}

impl Encode for TaskId {
    fn encode<E: Encoder>(&self, encoder: &mut E) -> Result<(), EncodeError> {
        debug_assert!(!self.is_transient(), "transient TaskId must not be encoded");
        self.id.encode(encoder)
    }
}

impl Debug for TaskId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TaskId").field("id", &self.id).finish()
    }
}

unsafe impl crate::NonLocalValue for TaskId {}

/// `TaskId` values are constrained to 31 bits to preserve a niche for [`crate::RawVc`]. Bit 30
/// marks transient tasks; bit 31 is always zero.
pub const TRANSIENT_TASK_BIT: u32 = 0x4000_0000;

/// The largest value a [`TaskId`] may hold (31 bits set).
pub const TASK_ID_MAX: u32 = 0x7FFF_FFFF;

impl TaskId {
    pub fn is_transient(&self) -> bool {
        **self & TRANSIENT_TASK_BIT != 0
    }
    pub fn persistence(&self) -> TaskPersistence {
        // tasks with `TaskPersistence::LocalCells` have no `TaskId`, so we can ignore that case
        if self.is_transient() {
            TaskPersistence::Transient
        } else {
            TaskPersistence::Persistent
        }
    }
}

macro_rules! make_registered_serializable {
    ($ty:ty, $primitive:ty, $get_object:path, $validate_type_id:path $(,)?) => {
        impl Serialize for $ty {
            fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
            where
                S: serde::Serializer,
            {
                serializer.serialize_u16(self.id.into())
            }
        }

        impl<'de> Deserialize<'de> for $ty {
            fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
            where
                D: serde::Deserializer<'de>,
            {
                struct DeserializeVisitor;
                impl<'de> Visitor<'de> for DeserializeVisitor {
                    type Value = $ty;

                    fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
                        formatter.write_str(concat!("an id of a registered ", stringify!($ty)))
                    }

                    fn visit_u16<E>(self, v: u16) -> Result<Self::Value, E>
                    where
                        E: serde::de::Error,
                    {
                        match Self::Value::new(v) {
                            Some(value) => {
                                if let Some(error) = $validate_type_id(value) {
                                    Err(E::custom(error))
                                } else {
                                    Ok(value)
                                }
                            }
                            None => Err(E::unknown_variant(&format!("{v}"), &["a non zero u16"])),
                        }
                    }
                }

                deserializer.deserialize_u16(DeserializeVisitor)
            }
        }

        impl Debug for $ty {
            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                f.debug_struct(stringify!($ty))
                    .field("id", &self.id)
                    .field("name", &$get_object(*self))
                    .finish()
            }
        }

        impl Encode for $ty {
            fn encode<E: Encoder>(&self, encoder: &mut E) -> Result<(), EncodeError> {
                <NonZero<$primitive> as Encode>::encode(&self.id, encoder)
            }
        }

        impl<Context> Decode<Context> for $ty {
            fn decode<D: Decoder<Context = Context>>(decoder: &mut D) -> Result<Self, DecodeError> {
                let value = Self {
                    id: NonZero::<$primitive>::decode(decoder)?,
                };
                if let Some(error) = $validate_type_id(value) {
                    Err(DecodeError::OtherString(error.to_string()))
                } else {
                    Ok(value)
                }
            }
        }

        impl_borrow_decode!($ty);
    };
}

make_registered_serializable!(
    ValueTypeId,
    u16,
    registry::get_value_type,
    registry::validate_value_type_id,
);
make_registered_serializable!(
    TraitTypeId,
    u16,
    registry::get_trait,
    registry::validate_trait_type_id,
);
make_registered_serializable!(
    FunctionId,
    u16,
    registry::get_native_function,
    registry::validate_function_id,
);