turbo_tasks/
registry.rs

1use std::{fmt::Debug, hash::Hash, num::NonZeroU64, ops::Deref, sync::RwLock};
2
3use dashmap::mapref::entry::Entry;
4use once_cell::sync::Lazy;
5use rustc_hash::FxHashMap;
6
7use crate::{
8    FxDashMap, TraitType, ValueType,
9    id::{TraitTypeId, ValueTypeId},
10    id_factory::IdFactory,
11    native_function::NativeFunction,
12    no_move_vec::NoMoveVec,
13};
14
15static NAME_TO_FUNCTION: Lazy<RwLock<FxHashMap<&'static str, &'static NativeFunction>>> =
16    Lazy::new(RwLock::default);
17
18static VALUE_TYPE_ID_FACTORY: IdFactory<ValueTypeId> = IdFactory::new_const(
19    ValueTypeId::MIN.to_non_zero_u64(),
20    ValueTypeId::MAX.to_non_zero_u64(),
21);
22static VALUE_TYPES_BY_NAME: Lazy<FxDashMap<&'static str, ValueTypeId>> =
23    Lazy::new(FxDashMap::default);
24static VALUE_TYPES_BY_VALUE: Lazy<FxDashMap<&'static ValueType, ValueTypeId>> =
25    Lazy::new(FxDashMap::default);
26static VALUE_TYPES: Lazy<NoMoveVec<(&'static ValueType, &'static str)>> = Lazy::new(NoMoveVec::new);
27
28static TRAIT_TYPE_ID_FACTORY: IdFactory<TraitTypeId> = IdFactory::new_const(
29    TraitTypeId::MIN.to_non_zero_u64(),
30    TraitTypeId::MAX.to_non_zero_u64(),
31);
32static TRAIT_TYPES_BY_NAME: Lazy<FxDashMap<&'static str, TraitTypeId>> =
33    Lazy::new(FxDashMap::default);
34static TRAIT_TYPES_BY_VALUE: Lazy<FxDashMap<&'static TraitType, TraitTypeId>> =
35    Lazy::new(FxDashMap::default);
36static TRAIT_TYPES: Lazy<NoMoveVec<(&'static TraitType, &'static str)>> = Lazy::new(NoMoveVec::new);
37
38/// Registers the value and returns its id if this is the initial
39fn register_thing<
40    K: Copy + Deref<Target = u32> + TryFrom<NonZeroU64>,
41    V: Copy + Hash + Eq,
42    const INITIAL_CAPACITY_BITS: u32,
43>(
44    global_name: &'static str,
45    value: V,
46    id_factory: &IdFactory<K>,
47    store: &NoMoveVec<(V, &'static str), INITIAL_CAPACITY_BITS>,
48    map_by_name: &FxDashMap<&'static str, K>,
49    map_by_value: &FxDashMap<V, K>,
50) -> Option<K> {
51    if let Entry::Vacant(e) = map_by_value.entry(value) {
52        let new_id = id_factory.get();
53        // SAFETY: this is a fresh id
54        unsafe {
55            store.insert(*new_id as usize, (value, global_name));
56        }
57        map_by_name.insert(global_name, new_id);
58        e.insert(new_id);
59        Some(new_id)
60    } else {
61        None
62    }
63}
64
65fn get_thing_id<K, V>(value: V, map_by_value: &FxDashMap<V, K>) -> K
66where
67    V: Hash + Eq + Debug,
68    K: Clone,
69{
70    if let Some(id) = map_by_value.get(&value) {
71        id.clone()
72    } else {
73        panic!("Use of unregistered {value:?}");
74    }
75}
76
77/// Registers a function so it is available for persistence
78pub fn register_function(global_name: &'static str, func: &'static NativeFunction) {
79    let prev = NAME_TO_FUNCTION.write().unwrap().insert(global_name, func);
80    debug_assert!(
81        prev.is_none(),
82        "registration mappings for {global_name} are inconsistent!"
83    );
84}
85
86pub fn get_function_by_global_name(global_name: &str) -> &'static NativeFunction {
87    NAME_TO_FUNCTION.read().unwrap().get(global_name).unwrap()
88}
89
90pub fn register_value_type(
91    global_name: &'static str,
92    ty: &'static ValueType,
93) -> Option<ValueTypeId> {
94    register_thing(
95        global_name,
96        ty,
97        &VALUE_TYPE_ID_FACTORY,
98        &VALUE_TYPES,
99        &VALUE_TYPES_BY_NAME,
100        &VALUE_TYPES_BY_VALUE,
101    )
102}
103
104pub fn get_value_type_id(func: &'static ValueType) -> ValueTypeId {
105    get_thing_id(func, &VALUE_TYPES_BY_VALUE)
106}
107
108pub fn get_value_type_id_by_global_name(global_name: &str) -> Option<ValueTypeId> {
109    VALUE_TYPES_BY_NAME.get(global_name).map(|x| *x)
110}
111
112pub fn get_value_type(id: ValueTypeId) -> &'static ValueType {
113    VALUE_TYPES.get(*id as usize).unwrap().0
114}
115
116pub fn get_value_type_global_name(id: ValueTypeId) -> &'static str {
117    VALUE_TYPES.get(*id as usize).unwrap().1
118}
119
120pub fn register_trait_type(global_name: &'static str, ty: &'static TraitType) {
121    register_thing(
122        global_name,
123        ty,
124        &TRAIT_TYPE_ID_FACTORY,
125        &TRAIT_TYPES,
126        &TRAIT_TYPES_BY_NAME,
127        &TRAIT_TYPES_BY_VALUE,
128    );
129}
130
131pub fn get_trait_type_id(func: &'static TraitType) -> TraitTypeId {
132    get_thing_id(func, &TRAIT_TYPES_BY_VALUE)
133}
134
135pub fn get_trait_type_id_by_global_name(global_name: &str) -> Option<TraitTypeId> {
136    TRAIT_TYPES_BY_NAME.get(global_name).map(|x| *x)
137}
138
139pub fn get_trait(id: TraitTypeId) -> &'static TraitType {
140    TRAIT_TYPES.get(*id as usize).unwrap().0
141}
142
143pub fn get_trait_type_global_name(id: TraitTypeId) -> &'static str {
144    TRAIT_TYPES.get(*id as usize).unwrap().1
145}
turbo_tasks/registry.rs

turbo_tasks/
registry.rs
