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

turbo_tasks/
registry.rs
