1 files changed, 11 insertions, 18 deletions
diff --git a/src/main/kotlin/events/CustomItemModelEvent.kt b/src/main/kotlin/events/CustomItemModelEvent.kt
index 27524a9..e50eca4 100644
--- a/src/main/kotlin/events/CustomItemModelEvent.kt
+++ b/src/main/kotlin/events/CustomItemModelEvent.kt
@@ -1,24 +1,25 @@
-
-
 package moe.nea.firmament.events
 
-import java.util.*
+import java.util.Optional
+import kotlin.jvm.optionals.getOrNull
 import net.minecraft.client.render.model.BakedModel
 import net.minecraft.client.render.model.BakedModelManager
 import net.minecraft.client.util.ModelIdentifier
 import net.minecraft.item.ItemStack
+import moe.nea.firmament.util.collections.WeakCache
 
 data class CustomItemModelEvent(
     val itemStack: ItemStack,
     var overrideModel: ModelIdentifier? = null,
 ) : FirmamentEvent() {
     companion object : FirmamentEventBus<CustomItemModelEvent>() {
-        private val cache = IdentityHashMap<ItemStack?, Any>()
-        private val sentinelNull = Object()
-
-        fun clearCache() {
-            cache.clear()
-        }
+        val cache =
+            WeakCache.memoize<ItemStack, BakedModelManager, Optional<BakedModel>>("CustomItemModels") { stack, models ->
+                val modelId = getModelIdentifier(stack) ?: return@memoize Optional.empty()
+                val bakedModel = models.getModel(modelId)
+                if (bakedModel === models.missingModel) return@memoize Optional.empty()
+                Optional.of(bakedModel)
+            }
 
         @JvmStatic
         fun getModelIdentifier(itemStack: ItemStack?): ModelIdentifier? {
@@ -29,15 +30,7 @@ data class CustomItemModelEvent(
         @JvmStatic
         fun getModel(itemStack: ItemStack?, thing: BakedModelManager): BakedModel? {
             if (itemStack == null) return null
-            val cachedValue = cache.getOrPut(itemStack) {
-                val modelId = getModelIdentifier(itemStack) ?: return@getOrPut sentinelNull
-                val bakedModel = thing.getModel(modelId)
-                if (bakedModel === thing.missingModel) return@getOrPut sentinelNull
-                bakedModel
-            }
-            if (cachedValue === sentinelNull)
-                return null
-            return cachedValue as BakedModel
+            return cache.invoke(itemStack, thing).getOrNull()
         }
     }
 }