Use weak caches for custom textures

4 files changed, 68 insertions, 51 deletions

diff --git a/src/main/kotlin/features/texturepack/AndPredicate.kt b/src/main/kotlin/features/texturepack/AndPredicate.kt
index 55a4f32..dc8e852 100644
--- a/src/main/kotlin/features/texturepack/AndPredicate.kt
+++ b/src/main/kotlin/features/texturepack/AndPredicate.kt
@@ -1,4 +1,3 @@
-
 package moe.nea.firmament.features.texturepack
 
 import com.google.gson.JsonArray
diff --git a/src/main/kotlin/features/texturepack/CustomGlobalArmorOverrides.kt b/src/main/kotlin/features/texturepack/CustomGlobalArmorOverrides.kt
index 23577ee..7b6e62b 100644
--- a/src/main/kotlin/features/texturepack/CustomGlobalArmorOverrides.kt
+++ b/src/main/kotlin/features/texturepack/CustomGlobalArmorOverrides.kt
@@ -1,12 +1,13 @@
-
 @file:UseSerializers(IdentifierSerializer::class)
 
 package moe.nea.firmament.features.texturepack
 
+import java.util.Optional
 import kotlinx.serialization.SerialName
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.Transient
 import kotlinx.serialization.UseSerializers
+import kotlin.jvm.optionals.getOrNull
 import net.minecraft.item.ArmorMaterial
 import net.minecraft.item.ItemStack
 import net.minecraft.resource.ResourceManager
@@ -20,8 +21,7 @@ import moe.nea.firmament.events.subscription.SubscriptionOwner
 import moe.nea.firmament.features.FirmamentFeature
 import moe.nea.firmament.features.texturepack.CustomGlobalTextures.logger
 import moe.nea.firmament.util.IdentifierSerializer
-import moe.nea.firmament.util.IdentityCharacteristics
-import moe.nea.firmament.util.computeNullableFunction
+import moe.nea.firmament.util.collections.WeakCache
 import moe.nea.firmament.util.skyBlockId
 
 object CustomGlobalArmorOverrides : SubscriptionOwner {
@@ -59,21 +59,21 @@ object CustomGlobalArmorOverrides : SubscriptionOwner {
     override val delegateFeature: FirmamentFeature
         get() = CustomSkyBlockTextures
 
-    val overrideCache = mutableMapOf<IdentityCharacteristics<ItemStack>, Any>()
+    val overrideCache = WeakCache.memoize<ItemStack, Optional<List<ArmorMaterial.Layer>>>("ArmorOverrides") { stack ->
+        val id = stack.skyBlockId ?: return@memoize Optional.empty()
+        val override = overrides[id.neuItem] ?: return@memoize Optional.empty()
+        for (suboverride in override.overrides) {
+            if (suboverride.predicate.test(stack)) {
+                return@memoize Optional.of(suboverride.bakedLayers)
+            }
+        }
+        return@memoize Optional.of(override.bakedLayers)
+    }
 
     @JvmStatic
     fun overrideArmor(stack: ItemStack): List<ArmorMaterial.Layer>? {
         if (!CustomSkyBlockTextures.TConfig.enableArmorOverrides) return null
-        return overrideCache.computeNullableFunction(IdentityCharacteristics(stack)) {
-            val id = stack.skyBlockId ?: return@computeNullableFunction null
-            val override = overrides[id.neuItem] ?: return@computeNullableFunction null
-            for (suboverride in override.overrides) {
-                if (suboverride.predicate.test(stack)) {
-                    return@computeNullableFunction suboverride.bakedLayers
-                }
-            }
-            return@computeNullableFunction override.bakedLayers
-        }
+        return overrideCache.invoke(stack).getOrNull()
     }
 
     var overrides: Map<String, ArmorOverride> = mapOf()
diff --git a/src/main/kotlin/features/texturepack/CustomGlobalTextures.kt b/src/main/kotlin/features/texturepack/CustomGlobalTextures.kt
index d64c844..2771699 100644
--- a/src/main/kotlin/features/texturepack/CustomGlobalTextures.kt
+++ b/src/main/kotlin/features/texturepack/CustomGlobalTextures.kt
@@ -1,9 +1,9 @@
-
 @file:UseSerializers(IdentifierSerializer::class, CustomModelOverrideParser.FirmamentRootPredicateSerializer::class)
 
 package moe.nea.firmament.features.texturepack
 
 
+import java.util.Optional
 import java.util.concurrent.CompletableFuture
 import org.slf4j.LoggerFactory
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable
@@ -28,9 +28,9 @@ import moe.nea.firmament.events.ScreenChangeEvent
 import moe.nea.firmament.events.subscription.SubscriptionOwner
 import moe.nea.firmament.features.FirmamentFeature
 import moe.nea.firmament.util.IdentifierSerializer
-import moe.nea.firmament.util.IdentityCharacteristics
 import moe.nea.firmament.util.MC
-import moe.nea.firmament.util.computeNullableFunction
+import moe.nea.firmament.util.collections.WeakCache
+import moe.nea.firmament.util.intoOptional
 import moe.nea.firmament.util.json.SingletonSerializableList
 import moe.nea.firmament.util.runNull
 
@@ -140,7 +140,17 @@ object CustomGlobalTextures : SinglePreparationResourceReloader<CustomGlobalText
             .filterTo(mutableSetOf()) { it.screenFilter.title.matches(newTitle) }
     }
 
-    val overrideCache = mutableMapOf<IdentityCharacteristics<ItemStack>, Any>()
+    val overrideCache = WeakCache.memoize<ItemStack, ItemModels, Optional<BakedModel>>("CustomGlobalTextureModelOverrides") { stack, models ->
+        matchingOverrides
+            .firstNotNullOfOrNull {
+                it.overrides
+                    .asSequence()
+                    .filter { it.predicate.test(stack) }
+                    .map { models.modelManager.getModel(ModelIdentifier(it.model, "inventory")) }
+                    .firstOrNull()
+            }
+            .intoOptional()
+    }
 
     @JvmStatic
     fun replaceGlobalModel(
@@ -148,19 +158,8 @@ object CustomGlobalTextures : SinglePreparationResourceReloader<CustomGlobalText
         stack: ItemStack,
         cir: CallbackInfoReturnable<BakedModel>
     ) {
-        val value = overrideCache.computeNullableFunction(IdentityCharacteristics(stack)) {
-            for (guiClassOverride in matchingOverrides) {
-                for (override in guiClassOverride.overrides) {
-                    if (override.predicate.test(stack)) {
-                        return@computeNullableFunction models.modelManager.getModel(
-                            ModelIdentifier(override.model, "inventory"))
-                    }
-                }
-            }
-            null
-        }
-        if (value != null)
-            cir.returnValue = value
+        overrideCache.invoke(stack, models)
+            .ifPresent(cir::setReturnValue)
     }
 
 
diff --git a/src/main/kotlin/features/texturepack/CustomSkyBlockTextures.kt b/src/main/kotlin/features/texturepack/CustomSkyBlockTextures.kt
index 692f73b..3b7cb96 100644
--- a/src/main/kotlin/features/texturepack/CustomSkyBlockTextures.kt
+++ b/src/main/kotlin/features/texturepack/CustomSkyBlockTextures.kt
@@ -2,7 +2,9 @@ package moe.nea.firmament.features.texturepack
 
 import com.mojang.authlib.minecraft.MinecraftProfileTexture
 import com.mojang.authlib.properties.Property
+import java.util.Optional
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable
+import kotlin.jvm.optionals.getOrNull
 import net.minecraft.block.SkullBlock
 import net.minecraft.client.MinecraftClient
 import net.minecraft.client.render.RenderLayer
@@ -12,10 +14,11 @@ import net.minecraft.util.Identifier
 import moe.nea.firmament.annotations.Subscribe
 import moe.nea.firmament.events.BakeExtraModelsEvent
 import moe.nea.firmament.events.CustomItemModelEvent
+import moe.nea.firmament.events.FinalizeResourceManagerEvent
 import moe.nea.firmament.events.TickEvent
 import moe.nea.firmament.features.FirmamentFeature
 import moe.nea.firmament.gui.config.ManagedConfig
-import moe.nea.firmament.util.IdentityCharacteristics
+import moe.nea.firmament.util.collections.WeakCache
 import moe.nea.firmament.util.item.decodeProfileTextureProperty
 import moe.nea.firmament.util.skyBlockId
 
@@ -26,7 +29,8 @@ object CustomSkyBlockTextures : FirmamentFeature {
     object TConfig : ManagedConfig(identifier) {
         val enabled by toggle("enabled") { true }
         val skullsEnabled by toggle("skulls-enabled") { true }
-        val cacheDuration by integer("cache-duration", 0, 20) { 1 }
+        val cacheForever by toggle("cache-forever") { true }
+        val cacheDuration by integer("cache-duration", 0, 100) { 1 }
         val enableModelOverrides by toggle("model-overrides") { true }
         val enableArmorOverrides by toggle("armor-overrides") { true }
         val enableBlockOverrides by toggle("block-overrides") { true }
@@ -36,14 +40,31 @@ object CustomSkyBlockTextures : FirmamentFeature {
     override val config: ManagedConfig
         get() = TConfig
 
+    val allItemCaches by lazy {
+        listOf(
+            CustomItemModelEvent.cache.cache,
+            skullTextureCache.cache,
+            CustomGlobalTextures.overrideCache.cache,
+            CustomGlobalArmorOverrides.overrideCache.cache
+        )
+    }
+
+    fun clearAllCaches() {
+        allItemCaches.forEach(WeakCache<*, *, *>::clear)
+    }
+
     @Subscribe
     fun onTick(it: TickEvent) {
+        if (TConfig.cacheForever) return
         if (TConfig.cacheDuration < 1 || it.tickCount % TConfig.cacheDuration == 0) {
-            // TODO: unify all of those caches somehow
-            CustomItemModelEvent.clearCache()
-            skullTextureCache.clear()
-            CustomGlobalTextures.overrideCache.clear()
-            CustomGlobalArmorOverrides.overrideCache.clear()
+            clearAllCaches()
+        }
+    }
+
+    @Subscribe
+    fun onStart(event: FinalizeResourceManagerEvent) {
+        event.registerOnApply("Clear firmament CIT caches") {
+            clearAllCaches()
         }
     }
 
@@ -74,8 +95,14 @@ object CustomSkyBlockTextures : FirmamentFeature {
         it.overrideModel = ModelIdentifier.ofInventoryVariant(Identifier.of("firmskyblock", id.identifier.path))
     }
 
-    private val skullTextureCache = mutableMapOf<IdentityCharacteristics<ProfileComponent>, Any>()
-    private val sentinelPresentInvalid = Object()
+    private val skullTextureCache =
+        WeakCache.memoize<ProfileComponent, Optional<Identifier>>("SkullTextureCache") { component ->
+            val id = getSkullTexture(component) ?: return@memoize Optional.empty()
+            if (!MinecraftClient.getInstance().resourceManager.getResource(id).isPresent) {
+                return@memoize Optional.empty()
+            }
+            return@memoize Optional.of(id)
+        }
 
     private val mcUrlRegex = "https?://textures.minecraft.net/texture/([a-fA-F0-9]+)".toRegex()
 
@@ -100,16 +127,8 @@ object CustomSkyBlockTextures : FirmamentFeature {
         if (type != SkullBlock.Type.PLAYER) return
         if (!TConfig.skullsEnabled) return
         if (component == null) return
-        val ic = IdentityCharacteristics(component)
 
-        val n = skullTextureCache.getOrPut(ic) {
-            val id = getSkullTexture(component) ?: return@getOrPut sentinelPresentInvalid
-            if (!MinecraftClient.getInstance().resourceManager.getResource(id).isPresent) {
-                return@getOrPut sentinelPresentInvalid
-            }
-            return@getOrPut id
-        }
-        if (n === sentinelPresentInvalid) return
-        cir.returnValue = RenderLayer.getEntityTranslucent(n as Identifier)
+        val n = skullTextureCache.invoke(component).getOrNull() ?: return
+        cir.returnValue = RenderLayer.getEntityTranslucent(n)
     }
 }