turbo_tasks/
registry.rs

1use std::num::NonZeroU16;
2
3use anyhow::Error;
4use once_cell::sync::Lazy;
5use rustc_hash::{FxHashMap, FxHashSet};
6
7use crate::{
8    TraitType, ValueType,
9    id::{FunctionId, TraitTypeId, ValueTypeId},
10    macro_helpers::CollectableFunction,
11    native_function::NativeFunction,
12    value_type::{CollectableTrait, CollectableValueType},
13};
14
15/// A trait for types that can be registered in a registry.
16///
17/// This allows the generic registry to work with different types
18/// while maintaining their specific requirements.
19trait RegistryItem: 'static + Eq + std::hash::Hash {
20    /// The ID type used for this registry item
21    type Id: Copy + From<NonZeroU16> + std::ops::Deref<Target = u16> + std::fmt::Display;
22    const TYPE_NAME: &'static str;
23
24    /// Get the global name used for sorting and uniqueness validation
25    fn global_name(&self) -> &'static str;
26}
27
28impl RegistryItem for NativeFunction {
29    type Id = FunctionId;
30    const TYPE_NAME: &'static str = "Function";
31
32    fn global_name(&self) -> &'static str {
33        self.global_name
34    }
35}
36
37impl RegistryItem for ValueType {
38    type Id = ValueTypeId;
39    const TYPE_NAME: &'static str = "Value";
40
41    fn global_name(&self) -> &'static str {
42        self.global_name
43    }
44}
45
46impl RegistryItem for TraitType {
47    type Id = TraitTypeId;
48    const TYPE_NAME: &'static str = "Trait";
49    fn global_name(&self) -> &'static str {
50        self.global_name
51    }
52}
53
54/// A generic registry that maps between IDs and static references to items.
55///
56/// This eliminates the code duplication between Functions, Values, and Traits registries.
57struct Registry<T: RegistryItem> {
58    id_to_item: Box<[&'static T]>,
59    item_to_id: FxHashMap<&'static T, T::Id>,
60}
61
62impl<T: RegistryItem> Registry<T> {
63    /// Create a new registry from a collection of items.
64    ///
65    /// Items are sorted by global_name to ensure stable ID assignment.
66    fn new_from_items(mut items: Vec<&'static T>) -> Self {
67        // Sort by global name to get stable order
68        items.sort_unstable_by_key(|item| item.global_name());
69
70        let mut item_to_id = FxHashMap::with_capacity_and_hasher(items.len(), Default::default());
71        let mut names = FxHashSet::with_capacity_and_hasher(items.len(), Default::default());
72
73        let mut id = NonZeroU16::MIN;
74        for &item in items.iter() {
75            item_to_id.insert(item, id.into());
76            let global_name = item.global_name();
77            assert!(
78                names.insert(global_name),
79                "multiple {ty} items registered with name: {global_name}!",
80                ty = T::TYPE_NAME
81            );
82            id = id.checked_add(1).expect("overflowing item ids");
83        }
84
85        Self {
86            id_to_item: items.into_boxed_slice(),
87            item_to_id,
88        }
89    }
90
91    /// Get an item by its ID
92    fn get_item(&self, id: T::Id) -> &'static T {
93        self.id_to_item[*id as usize - 1]
94    }
95
96    /// Get the ID for an item
97    fn get_id(&self, item: &'static T) -> T::Id {
98        match self.item_to_id.get(&item) {
99            Some(id) => *id,
100            None => panic!(
101                "{ty} isn't registered: {item}",
102                ty = T::TYPE_NAME,
103                item = item.global_name()
104            ),
105        }
106    }
107
108    /// Validate that an ID is within the valid range
109    fn validate_id(&self, id: T::Id) -> Option<Error> {
110        let len = self.id_to_item.len();
111        if *id as usize <= len {
112            None
113        } else {
114            Some(anyhow::anyhow!(
115                "Invalid {ty} id, {id} expected a value <= {len}",
116                ty = T::TYPE_NAME
117            ))
118        }
119    }
120}
121
122static FUNCTIONS: Lazy<Registry<NativeFunction>> = Lazy::new(|| {
123    let functions = inventory::iter::<CollectableFunction>
124        .into_iter()
125        .map(|c| &**c.0)
126        .collect::<Vec<_>>();
127    Registry::new_from_items(functions)
128});
129
130pub fn get_native_function(id: FunctionId) -> &'static NativeFunction {
131    FUNCTIONS.get_item(id)
132}
133
134pub fn get_function_id(func: &'static NativeFunction) -> FunctionId {
135    FUNCTIONS.get_id(func)
136}
137
138pub fn validate_function_id(id: FunctionId) -> Option<Error> {
139    FUNCTIONS.validate_id(id)
140}
141
142static VALUES: Lazy<Registry<ValueType>> = Lazy::new(|| {
143    // Inventory does not guarantee an order. So we sort by the global name to get a stable order
144    // This ensures that assigned ids are also stable which is important since they are serialized.
145    let all_values = inventory::iter::<CollectableValueType>
146        .into_iter()
147        .map(|t| &**t.0)
148        .collect::<Vec<_>>();
149    Registry::new_from_items(all_values)
150});
151
152pub fn get_value_type_id(value: &'static ValueType) -> ValueTypeId {
153    VALUES.get_id(value)
154}
155
156pub fn get_value_type(id: ValueTypeId) -> &'static ValueType {
157    VALUES.get_item(id)
158}
159
160pub fn validate_value_type_id(id: ValueTypeId) -> Option<Error> {
161    VALUES.validate_id(id)
162}
163
164static TRAITS: Lazy<Registry<TraitType>> = Lazy::new(|| {
165    // Inventory does not guarantee an order. So we sort by the global name to get a stable order
166    // This ensures that assigned ids are also stable.
167    let all_traits = inventory::iter::<CollectableTrait>
168        .into_iter()
169        .map(|t| &**t.0)
170        .collect::<Vec<_>>();
171    Registry::new_from_items(all_traits)
172});
173
174pub fn get_trait_type_id(trait_type: &'static TraitType) -> TraitTypeId {
175    TRAITS.get_id(trait_type)
176}
177
178pub fn get_trait(id: TraitTypeId) -> &'static TraitType {
179    TRAITS.get_item(id)
180}
181
182pub fn validate_trait_type_id(id: TraitTypeId) -> Option<Error> {
183    TRAITS.validate_id(id)
184}
turbo_tasks/registry.rs

turbo_tasks/
registry.rs
