diff --git a/src/nodes/drawable/model.rs b/src/nodes/drawable/model.rs index 3e98179..0c7469d 100644 --- a/src/nodes/drawable/model.rs +++ b/src/nodes/drawable/model.rs @@ -22,11 +22,18 @@ use stereokit_rust::sk::MainThreadToken; use stereokit_rust::{material::Material, model::Model as SKModel, tex::Tex, util::Color128}; pub struct MaterialWrapper(pub Material); +impl Drop for MaterialWrapper { + fn drop(&mut self) { + MATERIAL_REGISTRY.remove(self); + } +} + impl Hash for MaterialWrapper { fn hash(&self, state: &mut H) { self.0.get_shader().0.as_ptr().hash(state); for param in self.0.get_all_param_info() { - param.to_string().hash(state) + param.name.hash(state); + param.to_string().hash(state); } self.0.get_chain().map(MaterialWrapper).hash(state) } @@ -59,10 +66,9 @@ unsafe impl Send for MaterialWrapper {} unsafe impl Sync for MaterialWrapper {} #[derive(Default)] -struct MaterialRegistry(Mutex>); +struct MaterialRegistry(Mutex>>); impl MaterialRegistry { fn add_or_get(&self, material: Arc) -> Arc { - let mut lock = self.0.lock(); let hash = { use std::hash::{Hash, Hasher}; let mut hasher = std::collections::hash_map::DefaultHasher::new(); @@ -70,15 +76,26 @@ impl MaterialRegistry { hasher.finish() }; - if let Some(id) = lock.get(&hash) { - if let Ok(existing) = Material::find(id) { - return Arc::new(MaterialWrapper(existing)); + let mut lock = self.0.lock(); + if let Some(mat) = lock.get(&hash) { + if let Some(mat) = mat.upgrade() { + return mat; } } - lock.insert(hash, material.0.get_id().to_string()); + lock.insert(hash, Arc::downgrade(&material)); material } + fn remove(&self, material: &MaterialWrapper) { + let hash = { + use std::hash::{Hash, Hasher}; + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + material.hash(&mut hasher); + hasher.finish() + }; + let mut lock = self.0.lock(); + lock.remove(&hash); + } } static MATERIAL_REGISTRY: Lazy = Lazy::new(MaterialRegistry::default); @@ -132,6 +149,7 @@ pub struct ModelPart { path: String, space: Arc, model: Weak, + material: Mutex>>, pending_material_parameters: Mutex>, pending_material_replacement: Mutex>>, aliases: AliasList, @@ -205,6 +223,7 @@ impl ModelPart { pending_material_parameters: Mutex::new(FxHashMap::default()), pending_material_replacement: Mutex::new(None), aliases: AliasList::default(), + material: Mutex::new(part.get_material().map(MaterialWrapper).map(Arc::new)), }); node.add_aspect_raw(model_part.clone()); parts.push(model_part.clone()); @@ -231,7 +250,10 @@ impl ModelPart { }; let shared_material = MATERIAL_REGISTRY.add_or_get(Arc::new(MaterialWrapper(replacement.copy()))); + + let mut lock = self.material.lock(); part.material(&shared_material.0); + lock.replace(shared_material); } fn update(&self) { @@ -258,7 +280,9 @@ impl ModelPart { }; if let Some(material_replacement) = self.pending_material_replacement.lock().take() { + let mut lock = self.material.lock(); part.material(&material_replacement.0); + lock.replace(material_replacement); } 'mat_params: { @@ -274,7 +298,9 @@ impl ModelPart { let shared_material = MATERIAL_REGISTRY.add_or_get(Arc::new(MaterialWrapper(new_material))); + let mut lock = self.material.lock(); part.material(&shared_material.0); + lock.replace(shared_material); } } }