proper prehashing

Author: cart

crates/bevy_render/src/mesh/mesh/mod.rs

@@ -10,11 +10,8 @@ use bevy_core::cast_slice;
 use bevy_ecs::system::{lifetimeless::SRes, SystemParamItem};
 use bevy_math::*;
 use bevy_reflect::TypeUuid;
-use bevy_utils::{AHasher, EnumVariantMeta};
-use std::{
-    collections::BTreeMap,
-    hash::{Hash, Hasher},
-};
+use bevy_utils::{EnumVariantMeta, Hashed};
+use std::{collections::BTreeMap, hash::Hash};
 use wgpu::{
     util::BufferInitDescriptor, BufferUsages, IndexFormat, PrimitiveTopology, VertexAttribute,
     VertexFormat, VertexStepMode,
@@ -186,14 +183,14 @@ impl Mesh {
             accumulated_offset += data.attribute.format.get_size();
         }
 
-        MeshVertexBufferLayout::new(
-            VertexBufferLayout {
+        MeshVertexBufferLayout::new(InnerMeshVertexBufferLayout {
+            layout: VertexBufferLayout {
                 array_stride: accumulated_offset,
                 step_mode: VertexStepMode::Vertex,
                 attributes,
             },
             attribute_ids,
-        )
+        })
     }
 
     /// Counts all vertices of the mesh.
@@ -390,26 +387,15 @@ impl From<MeshVertexAttribute> for MeshVertexAttributeId {
     }
 }
 
-#[derive(Debug, Clone)]
-pub struct MeshVertexBufferLayout {
+pub type MeshVertexBufferLayout = Hashed<InnerMeshVertexBufferLayout>;
+
+#[derive(Debug, Clone, Hash, Eq, PartialEq)]
+pub struct InnerMeshVertexBufferLayout {
     attribute_ids: Vec<MeshVertexAttributeId>,
     layout: VertexBufferLayout,
-    hash: MeshVertexBufferLayoutHash,
 }
 
-impl MeshVertexBufferLayout {
-    pub fn new(layout: VertexBufferLayout, attribute_ids: Vec<MeshVertexAttributeId>) -> Self {
-        let mut hasher = AHasher::default();
-        let key = (layout, attribute_ids);
-        key.hash(&mut hasher);
-
-        Self {
-            layout: key.0,
-            attribute_ids: key.1,
-            hash: MeshVertexBufferLayoutHash(hasher.finish()),
-        }
-    }
-
+impl InnerMeshVertexBufferLayout {
     #[inline]
     pub fn contains(&self, attribute_id: impl Into<MeshVertexAttributeId>) -> bool {
         self.attribute_ids.contains(&attribute_id.into())
@@ -454,16 +440,8 @@ impl MeshVertexBufferLayout {
             attributes,
         })
     }
-
-    #[inline]
-    pub(crate) fn hash(&self) -> MeshVertexBufferLayoutHash {
-        self.hash
-    }
 }
 
-#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)]
-pub(crate) struct MeshVertexBufferLayoutHash(u64);
-
 pub struct VertexAttributeDescriptor {
     pub shader_location: u32,
     pub id: MeshVertexAttributeId,

crates/bevy_render/src/render_resource/pipeline_specializer.rs

@@ -1,8 +1,8 @@
 use crate::{
-    mesh::{MeshVertexBufferLayout, MeshVertexBufferLayoutHash},
+    mesh::{InnerMeshVertexBufferLayout, MeshVertexBufferLayout},
     render_resource::{CachedPipelineId, RenderPipelineCache, RenderPipelineDescriptor},
 };
-use bevy_utils::HashMap;
+use bevy_utils::{HashMap, PreHashMap, PreHashMapExt};
 use std::hash::Hash;
 
 pub struct SpecializedPipelines<S: SpecializedPipeline> {
@@ -35,8 +35,9 @@ pub trait SpecializedPipeline {
     type Key: Clone + Hash + PartialEq + Eq;
     fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor;
 }
+
 pub struct SpecializedMeshPipelines<S: SpecializedMeshPipeline> {
-    cache: HashMap<(S::Key, MeshVertexBufferLayoutHash), CachedPipelineId>,
+    cache: PreHashMap<InnerMeshVertexBufferLayout, HashMap<S::Key, CachedPipelineId>>,
 }
 
 impl<S: SpecializedMeshPipeline> Default for SpecializedMeshPipelines<S> {
@@ -48,20 +49,19 @@ impl<S: SpecializedMeshPipeline> Default for SpecializedMeshPipelines<S> {
 }
 
 impl<S: SpecializedMeshPipeline> SpecializedMeshPipelines<S> {
+    #[inline]
     pub fn specialize(
         &mut self,
         cache: &mut RenderPipelineCache,
         specialize_pipeline: &S,
         key: S::Key,
         layout: &MeshVertexBufferLayout,
     ) -> CachedPipelineId {
-        *self
-            .cache
-            .entry((key.clone(), layout.hash()))
-            .or_insert_with(|| {
-                let descriptor = specialize_pipeline.specialize(key, layout);
-                cache.queue(descriptor)
-            })
+        let map = self.cache.get_or_insert_with(layout, Default::default);
+        *map.entry(key.clone()).or_insert_with(|| {
+            let descriptor = specialize_pipeline.specialize(key, layout);
+            cache.queue(descriptor)
+        })
     }
 }
 

crates/bevy_utils/src/lib.rs

@@ -4,12 +4,21 @@ pub mod label;
 pub use ahash::AHasher;
 pub use enum_variant_meta::*;
 pub type Entry<'a, K, V> = hashbrown::hash_map::Entry<'a, K, V, RandomState>;
+pub use hashbrown;
+use hashbrown::hash_map::RawEntryMut;
 pub use instant::{Duration, Instant};
 pub use tracing;
 pub use uuid::Uuid;
 
 use ahash::RandomState;
-use std::{future::Future, pin::Pin};
+use std::{
+    fmt::Debug,
+    future::Future,
+    hash::{BuildHasher, Hash, Hasher},
+    marker::PhantomData,
+    ops::Deref,
+    pin::Pin,
+};
 
 #[cfg(not(target_arch = "wasm32"))]
 pub type BoxedFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
@@ -208,3 +217,138 @@ pub type StableHashSet<K> = hashbrown::HashSet<K, FixedState>;
 //         StableHashSet::with_capacity_and_hasher(capacity, FixedState::default())
 //     }
 // }
+
+pub struct Hashed<V, H = FixedState> {
+    hash: u64,
+    value: V,
+    marker: PhantomData<H>,
+}
+
+impl<V: Hash, H: BuildHasher + Default> Hashed<V, H> {
+    pub fn new(value: V) -> Self {
+        let builder = H::default();
+        let mut hasher = builder.build_hasher();
+        value.hash(&mut hasher);
+        Self {
+            hash: hasher.finish(),
+            value,
+            marker: PhantomData,
+        }
+    }
+
+    #[inline]
+    pub fn hash(&self) -> u64 {
+        self.hash
+    }
+}
+
+impl<V, H> Hash for Hashed<V, H> {
+    #[inline]
+    fn hash<R: Hasher>(&self, state: &mut R) {
+        state.write_u64(self.hash);
+    }
+}
+
+impl<V, H> Deref for Hashed<V, H> {
+    type Target = V;
+
+    #[inline]
+    fn deref(&self) -> &Self::Target {
+        &self.value
+    }
+}
+
+impl<V: PartialEq, H> Hashed<V, H> {
+    #[inline]
+    pub fn fast_eq(&self, other: &Hashed<V, H>) -> bool {
+        // Makes the common case of two values not been equal very fast
+        self.hash == other.hash && self.value.eq(&other.value)
+    }
+}
+
+impl<V: PartialEq, H> PartialEq for Hashed<V, H> {
+    #[inline]
+    fn eq(&self, other: &Self) -> bool {
+        self.value.eq(&other.value)
+    }
+}
+
+impl<V: Debug, H> Debug for Hashed<V, H> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Hashed")
+            .field("hash", &self.hash)
+            .field("value", &self.value)
+            .finish()
+    }
+}
+
+impl<V: Clone, H> Clone for Hashed<V, H> {
+    #[inline]
+    fn clone(&self) -> Self {
+        Self {
+            hash: self.hash.clone(),
+            value: self.value.clone(),
+            marker: PhantomData,
+        }
+    }
+}
+
+impl<V: Eq, H> Eq for Hashed<V, H> {}
+
+#[derive(Default)]
+pub struct PassHash;
+
+impl BuildHasher for PassHash {
+    type Hasher = PassHasher;
+
+    fn build_hasher(&self) -> Self::Hasher {
+        PassHasher::default()
+    }
+}
+
+#[derive(Debug)]
+pub struct PassHasher {
+    hash: u64,
+}
+
+impl Default for PassHasher {
+    fn default() -> Self {
+        Self { hash: 0 }
+    }
+}
+
+impl Hasher for PassHasher {
+    fn write(&mut self, _bytes: &[u8]) {
+        panic!("cannot hash byte arrays using PassHasher");
+    }
+
+    fn write_u64(&mut self, i: u64) {
+        self.hash = i;
+    }
+
+    fn finish(&self) -> u64 {
+        self.hash
+    }
+}
+
+pub type PreHashMap<K, V> = hashbrown::HashMap<Hashed<K>, V, PassHash>;
+
+pub trait PreHashMapExt<K, V> {
+    fn get_or_insert_with<F: FnOnce() -> V>(&mut self, key: &Hashed<K>, func: F) -> &mut V;
+}
+
+impl<K: Hash + Eq + PartialEq + Clone, V> PreHashMapExt<K, V> for PreHashMap<K, V> {
+    #[inline]
+    fn get_or_insert_with<F: FnOnce() -> V>(&mut self, key: &Hashed<K>, func: F) -> &mut V {
+        let entry = self
+            .raw_entry_mut()
+            .from_key_hashed_nocheck(key.hash(), key);
+        match entry {
+            RawEntryMut::Occupied(entry) => entry.into_mut(),
+            RawEntryMut::Vacant(entry) => {
+                let (_, value) = entry.insert_hashed_nocheck(key.hash(), key.clone(), func());
+                value
+            }
+        }
+    }
+}