2 files changed, 37 insertions, 19 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()
         }
     }
 }
diff --git a/src/main/kotlin/events/FinalizeResourceManagerEvent.kt b/src/main/kotlin/events/FinalizeResourceManagerEvent.kt
index c43ad3b..0d411f1 100644
--- a/src/main/kotlin/events/FinalizeResourceManagerEvent.kt
+++ b/src/main/kotlin/events/FinalizeResourceManagerEvent.kt
@@ -1,10 +1,35 @@
-
 package moe.nea.firmament.events
 
+import java.util.concurrent.CompletableFuture
+import java.util.concurrent.Executor
 import net.minecraft.resource.ReloadableResourceManagerImpl
+import net.minecraft.resource.ResourceManager
+import net.minecraft.resource.ResourceReloader
+import net.minecraft.util.profiler.Profiler
 
 data class FinalizeResourceManagerEvent(
     val resourceManager: ReloadableResourceManagerImpl,
 ) : FirmamentEvent() {
     companion object : FirmamentEventBus<FinalizeResourceManagerEvent>()
+
+    inline fun registerOnApply(name: String, crossinline function: () -> Unit) {
+        resourceManager.registerReloader(object : ResourceReloader {
+            override fun reload(
+                synchronizer: ResourceReloader.Synchronizer,
+                manager: ResourceManager?,
+                prepareProfiler: Profiler?,
+                applyProfiler: Profiler?,
+                prepareExecutor: Executor?,
+                applyExecutor: Executor
+            ): CompletableFuture<Void> {
+                return CompletableFuture.completedFuture(Unit)
+                    .thenCompose(synchronizer::whenPrepared)
+                    .thenAcceptAsync({ function() }, applyExecutor)
+            }
+
+            override fun getName(): String {
+                return name
+            }
+        })
+    }
 }