Use a registry abstraction to reduce duplication

Author: lukesandberg

turbopack/crates/turbo-tasks/src/native_function.rs

@@ -273,18 +273,16 @@ impl NativeFunction {
         tracing::trace_span!("turbo_tasks::resolve_call", name = self.name, flags = flags)
     }
 }
-
-impl PartialEq for &'static NativeFunction {
+impl PartialEq for NativeFunction {
     fn eq(&self, other: &Self) -> bool {
-        std::ptr::eq(*self, *other)
+        std::ptr::eq(self, other)
     }
 }
 
-impl Eq for &'static NativeFunction {}
-
-impl Hash for &'static NativeFunction {
+impl Eq for NativeFunction {}
+impl Hash for NativeFunction {
     fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
-        Hash::hash(&(*self as *const NativeFunction), state);
+        (self as *const NativeFunction).hash(state);
     }
 }
 

turbopack/crates/turbo-tasks/src/registry.rs

@@ -12,168 +12,175 @@ use crate::{
     value_type::{CollectableTrait, CollectableValueType},
 };
 
-struct Functions {
-    id_to_value: Box<[&'static NativeFunction]>,
-    value_to_id: FxHashMap<&'static NativeFunction, FunctionId>,
+/// A trait for types that can be registered in a registry.
+///
+/// This allows the generic registry to work with different types
+/// while maintaining their specific requirements.
+trait RegistryItem: 'static + Eq + std::hash::Hash {
+    /// The ID type used for this registry item
+    type Id: Copy + From<NonZeroU16> + std::ops::Deref<Target = u16> + std::fmt::Display;
+    const TYPE_NAME: &'static str;
+
+    /// Get the global name used for sorting and uniqueness validation
+    fn global_name(&self) -> &'static str;
 }
-static FUNCTIONS: Lazy<Functions> = Lazy::new(|| {
-    let mut functions = inventory::iter::<CollectableFunction>
-        .into_iter()
-        .map(|c| &**c.0)
-        .collect::<Vec<_>>();
-    functions.sort_unstable_by_key(|f| f.global_name);
-    let mut value_to_id = FxHashMap::with_capacity_and_hasher(functions.len(), Default::default());
-    let mut names = FxHashSet::with_capacity_and_hasher(functions.len(), Default::default());
-
-    let mut id = NonZeroU16::MIN;
-    for &native_function in functions.iter() {
-        value_to_id.insert(native_function, id.into());
-        let global_name = native_function.global_name;
-        assert!(
-            names.insert(global_name),
-            "multiple functions registered with name: {global_name}!"
-        );
-        id = id.checked_add(1).expect("overflowing function ids");
+
+impl RegistryItem for NativeFunction {
+    type Id = FunctionId;
+    const TYPE_NAME: &'static str = "Function";
+
+    fn global_name(&self) -> &'static str {
+        self.global_name
+    }
+}
+
+impl RegistryItem for ValueType {
+    type Id = ValueTypeId;
+    const TYPE_NAME: &'static str = "Value";
+
+    fn global_name(&self) -> &'static str {
+        self.global_name
+    }
+}
+
+impl RegistryItem for TraitType {
+    type Id = TraitTypeId;
+    const TYPE_NAME: &'static str = "Trait";
+    fn global_name(&self) -> &'static str {
+        self.global_name
+    }
+}
+
+/// A generic registry that maps between IDs and static references to items.
+///
+/// This eliminates the code duplication between Functions, Values, and Traits registries.
+struct Registry<T: RegistryItem> {
+    id_to_item: Box<[&'static T]>,
+    item_to_id: FxHashMap<&'static T, T::Id>,
+}
+
+impl<T: RegistryItem> Registry<T> {
+    /// Create a new registry from a collection of items.
+    ///
+    /// Items are sorted by global_name to ensure stable ID assignment.
+    fn new_from_items(mut items: Vec<&'static T>) -> Self {
+        // Sort by global name to get stable order
+        items.sort_unstable_by_key(|item| item.global_name());
+
+        let mut item_to_id = FxHashMap::with_capacity_and_hasher(items.len(), Default::default());
+        let mut names = FxHashSet::with_capacity_and_hasher(items.len(), Default::default());
+
+        let mut id = NonZeroU16::MIN;
+        for &item in items.iter() {
+            item_to_id.insert(item, id.into());
+            let global_name = item.global_name();
+            assert!(
+                names.insert(global_name),
+                "multiple {ty} items registered with name: {global_name}!",
+                ty = T::TYPE_NAME
+            );
+            id = id.checked_add(1).expect("overflowing item ids");
+        }
+
+        Self {
+            id_to_item: items.into_boxed_slice(),
+            item_to_id,
+        }
+    }
+
+    /// Get an item by its ID
+    fn get_item(&self, id: T::Id) -> &'static T {
+        self.id_to_item[*id as usize - 1]
+    }
+
+    /// Get the ID for an item
+    fn get_id(&self, item: &'static T) -> T::Id {
+        match self.item_to_id.get(&item) {
+            Some(id) => *id,
+            None => panic!(
+                "{ty} isn't registered: {item}",
+                ty = T::TYPE_NAME,
+                item = item.global_name()
+            ),
+        }
     }
 
-    Functions {
-        id_to_value: functions.into_boxed_slice(),
-        value_to_id,
+    /// Validate that an ID is within the valid range
+    fn validate_id(&self, id: T::Id) -> Option<Error> {
+        let len = self.id_to_item.len();
+        if *id as usize <= len {
+            None
+        } else {
+            Some(anyhow::anyhow!(
+                "Invalid {ty} id, {id} expected a value <= {len}",
+                ty = T::TYPE_NAME
+            ))
+        }
     }
+}
+
+static FUNCTIONS: Lazy<Registry<NativeFunction>> = Lazy::new(|| {
+    let functions = inventory::iter::<CollectableFunction>
+        .into_iter()
+        .map(|c| &**c.0)
+        .collect::<Vec<_>>();
+    Registry::new_from_items(functions)
 });
 
 pub fn get_native_function(id: FunctionId) -> &'static NativeFunction {
-    FUNCTIONS.id_to_value[*id as usize - 1]
+    FUNCTIONS.get_item(id)
 }
 
 pub fn get_function_id(func: &'static NativeFunction) -> FunctionId {
-    *FUNCTIONS
-        .value_to_id
-        .get(&func)
-        .expect("function isn't registered")
+    FUNCTIONS.get_id(func)
 }
 
 pub fn validate_function_id(id: FunctionId) -> Option<Error> {
-    let len = FUNCTIONS.id_to_value.len();
-    if *id as usize <= len {
-        None
-    } else {
-        Some(anyhow::anyhow!(
-            "Invalid function type id, {id} expected a value <= {len}"
-        ))
-    }
-}
-
-struct Values {
-    id_to_value: Box<[&'static ValueType]>,
-    value_to_id: FxHashMap<&'static ValueType, ValueTypeId>,
+    FUNCTIONS.validate_id(id)
 }
 
-static VALUES: Lazy<Values> = Lazy::new(|| {
+static VALUES: Lazy<Registry<ValueType>> = Lazy::new(|| {
     // Inventory does not guarantee an order. So we sort by the global name to get a stable order
     // This ensures that assigned ids are also stable.
     // We don't currently take advantage of this but we could in the future.  The remaining issue is
     // ensuring the set of values is the same across runs.
-    let mut all_values = inventory::iter::<CollectableValueType>
+    let all_values = inventory::iter::<CollectableValueType>
         .into_iter()
         .map(|t| &**t.0)
         .collect::<Vec<_>>();
-    all_values.sort_unstable_by_key(|t| t.global_name);
-
-    let mut value_to_id = FxHashMap::with_capacity_and_hasher(all_values.len(), Default::default());
-    // Our sort above is non-sensical if names are not unique
-    let mut names = FxHashSet::with_capacity_and_hasher(all_values.len(), Default::default());
-
-    let mut id = NonZeroU16::MIN;
-    for &value_type in all_values.iter() {
-        value_to_id.insert(value_type, id.into());
-        let global_name = value_type.global_name;
-        assert!(
-            names.insert(global_name),
-            "two values registered with the same name: {global_name}"
-        );
-        id = id.checked_add(1).expect("overflowing value type ids");
-    }
-
-    Values {
-        value_to_id,
-        id_to_value: all_values.into_boxed_slice(),
-    }
+    Registry::new_from_items(all_values)
 });
 
 pub fn get_value_type_id(value: &'static ValueType) -> ValueTypeId {
-    match VALUES.value_to_id.get(value) {
-        Some(id) => *id,
-        None => panic!("Use of unregistered trait {value:?}"),
-    }
+    VALUES.get_id(value)
 }
 
 pub fn get_value_type(id: ValueTypeId) -> &'static ValueType {
-    VALUES.id_to_value[*id as usize - 1]
+    VALUES.get_item(id)
 }
 
 pub fn validate_value_type_id(id: ValueTypeId) -> Option<Error> {
-    let len = VALUES.id_to_value.len();
-    if *id as usize <= len {
-        None
-    } else {
-        Some(anyhow::anyhow!(
-            "Invalid value type id, {id} expected a value <= {len}"
-        ))
-    }
-}
-
-struct Traits {
-    id_to_trait: Box<[&'static TraitType]>,
-    trait_to_id: FxHashMap<&'static TraitType, TraitTypeId>,
+    VALUES.validate_id(id)
 }
 
-static TRAITS: Lazy<Traits> = Lazy::new(|| {
+static TRAITS: Lazy<Registry<TraitType>> = Lazy::new(|| {
     // Inventory does not guarantee an order. So we sort by the global name to get a stable order
     // This ensures that assigned ids are also stable.
-    let mut all_traits = inventory::iter::<CollectableTrait>
+    let all_traits = inventory::iter::<CollectableTrait>
         .into_iter()
         .map(|t| &**t.0)
         .collect::<Vec<_>>();
-    all_traits.sort_unstable_by_key(|t| t.global_name);
-
-    let mut trait_to_id = FxHashMap::with_capacity_and_hasher(all_traits.len(), Default::default());
-    // Our sort above is non-sensical if names are not unique
-    let mut names = FxHashSet::with_capacity_and_hasher(all_traits.len(), Default::default());
-    let mut id = NonZeroU16::MIN;
-    for &trait_type in all_traits.iter() {
-        trait_to_id.insert(trait_type, id.into());
-
-        let global_name = trait_type.global_name;
-        assert!(
-            names.insert(global_name),
-            "two traits registered with the same name: {global_name}"
-        );
-        id = id.checked_add(1).expect("overflowing trait type ids");
-    }
-    Traits {
-        trait_to_id,
-        id_to_trait: all_traits.into_boxed_slice(),
-    }
+    Registry::new_from_items(all_traits)
 });
 
 pub fn get_trait_type_id(trait_type: &'static TraitType) -> TraitTypeId {
-    match TRAITS.trait_to_id.get(trait_type) {
-        Some(id) => *id,
-        None => panic!("Use of unregistered trait {trait_type:?}"),
-    }
+    TRAITS.get_id(trait_type)
 }
 
 pub fn get_trait(id: TraitTypeId) -> &'static TraitType {
-    TRAITS.id_to_trait[*id as usize - 1]
+    TRAITS.get_item(id)
 }
+
 pub fn validate_trait_type_id(id: TraitTypeId) -> Option<Error> {
-    let len = TRAITS.id_to_trait.len();
-    if *id as usize <= len {
-        None
-    } else {
-        Some(anyhow::anyhow!(
-            "Invalid trait type id, {id} expected a value <= {len}"
-        ))
-    }
+    TRAITS.validate_id(id)
 }