Use weak caches for custom textures

17 files changed, 307 insertions, 78 deletions

diff --git a/src/main/kotlin/commands/rome.kt b/src/main/kotlin/commands/rome.kt
index fbcad07..3027a4b 100644
--- a/src/main/kotlin/commands/rome.kt
+++ b/src/main/kotlin/commands/rome.kt
@@ -23,6 +23,8 @@ import moe.nea.firmament.util.MC
 import moe.nea.firmament.util.SBData
 import moe.nea.firmament.util.ScreenUtil
 import moe.nea.firmament.util.SkyblockId
+import moe.nea.firmament.util.collections.InstanceList
+import moe.nea.firmament.util.collections.WeakCache
 
 
 fun firmamentCommand() = literal("firmament") {
@@ -229,6 +231,18 @@ fun firmamentCommand() = literal("firmament") {
                 }
             }
         }
+        thenLiteral("caches") {
+            thenExecute {
+                source.sendFeedback(Text.literal("Caches:"))
+                WeakCache.allInstances.getAll().forEach {
+                    source.sendFeedback(Text.literal(" - ${it.name}: ${it.size}"))
+                }
+                source.sendFeedback(Text.translatable("Instance lists:"))
+                InstanceList.allInstances.getAll().forEach {
+                    source.sendFeedback(Text.literal(" - ${it.name}: ${it.size}"))
+                }
+            }
+        }
         thenLiteral("mixins") {
             thenExecute {
                 source.sendFeedback(Text.translatable("firmament.mixins.start"))
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
+            }
+        })
+    }
 }
diff --git a/src/main/kotlin/features/debug/DebugLogger.kt b/src/main/kotlin/features/debug/DebugLogger.kt
index ab06030..69a191d 100644
--- a/src/main/kotlin/features/debug/DebugLogger.kt
+++ b/src/main/kotlin/features/debug/DebugLogger.kt
@@ -3,8 +3,15 @@ package moe.nea.firmament.features.debug
 
 import net.minecraft.text.Text
 import moe.nea.firmament.util.MC
+import moe.nea.firmament.util.collections.InstanceList
 
 class DebugLogger(val tag: String) {
+    companion object {
+        val allInstances = InstanceList<DebugLogger>("DebugLogger")
+    }
+    init {
+    	allInstances.add(this)
+    }
     fun isEnabled() = DeveloperFeatures.isEnabled // TODO: allow filtering by tag
     fun log(text: () -> String) {
         if (!isEnabled()) return
diff --git a/src/main/kotlin/features/diana/NearbyBurrowsSolver.kt b/src/main/kotlin/features/diana/NearbyBurrowsSolver.kt
index 7158bb9..ab1518a 100644
--- a/src/main/kotlin/features/diana/NearbyBurrowsSolver.kt
+++ b/src/main/kotlin/features/diana/NearbyBurrowsSolver.kt
@@ -14,7 +14,7 @@ import moe.nea.firmament.events.WorldRenderLastEvent
 import moe.nea.firmament.events.subscription.SubscriptionOwner
 import moe.nea.firmament.features.FirmamentFeature
 import moe.nea.firmament.util.TimeMark
-import moe.nea.firmament.util.mutableMapWithMaxSize
+import moe.nea.firmament.util.collections.mutableMapWithMaxSize
 import moe.nea.firmament.util.render.RenderInWorldContext.Companion.renderInWorld
 
 object NearbyBurrowsSolver : SubscriptionOwner {
diff --git a/src/main/kotlin/features/inventory/ItemRarityCosmetics.kt b/src/main/kotlin/features/inventory/ItemRarityCosmetics.kt
index 566a813..55509f5 100644
--- a/src/main/kotlin/features/inventory/ItemRarityCosmetics.kt
+++ b/src/main/kotlin/features/inventory/ItemRarityCosmetics.kt
@@ -14,9 +14,8 @@ import moe.nea.firmament.features.FirmamentFeature
 import moe.nea.firmament.gui.config.ManagedConfig
 import moe.nea.firmament.util.MC
 import moe.nea.firmament.util.item.loreAccordingToNbt
-import moe.nea.firmament.util.lastNotNullOfOrNull
-import moe.nea.firmament.util.memoize
-import moe.nea.firmament.util.memoizeIdentity
+import moe.nea.firmament.util.collections.lastNotNullOfOrNull
+import moe.nea.firmament.util.collections.memoizeIdentity
 import moe.nea.firmament.util.unformattedString
 
 object ItemRarityCosmetics : FirmamentFeature {
diff --git a/src/main/kotlin/features/inventory/buttons/InventoryButton.kt b/src/main/kotlin/features/inventory/buttons/InventoryButton.kt
index 539edf2..be173bd 100644
--- a/src/main/kotlin/features/inventory/buttons/InventoryButton.kt
+++ b/src/main/kotlin/features/inventory/buttons/InventoryButton.kt
@@ -17,7 +17,7 @@ import moe.nea.firmament.repo.ItemCache.asItemStack
 import moe.nea.firmament.repo.RepoManager
 import moe.nea.firmament.util.MC
 import moe.nea.firmament.util.SkyblockId
-import moe.nea.firmament.util.memoize
+import moe.nea.firmament.util.collections.memoize
 
 @Serializable
 data class InventoryButton(
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)
     }
 }
diff --git a/src/main/kotlin/rei/SBItemEntryDefinition.kt b/src/main/kotlin/rei/SBItemEntryDefinition.kt
index 5c6740e..077eeb1 100644
--- a/src/main/kotlin/rei/SBItemEntryDefinition.kt
+++ b/src/main/kotlin/rei/SBItemEntryDefinition.kt
@@ -204,7 +204,7 @@ object SBItemEntryDefinition : EntryDefinition<SBItemStack> {
     }
 
     override fun asFormattedText(entry: EntryStack<SBItemStack>, value: SBItemStack): Text {
-        return VanillaEntryTypes.ITEM.definition.asFormattedText(entry.asItemEntry(), value.asItemStack())
+        return VanillaEntryTypes.ITEM.definition.asFormattedText(entry.asItemEntry(), value.asImmutableItemStack())
     }
 
     override fun hash(entry: EntryStack<SBItemStack>, value: SBItemStack, context: ComparisonContext): Long {
diff --git a/src/main/kotlin/util/Optionalutil.kt b/src/main/kotlin/util/Optionalutil.kt
new file mode 100644
index 0000000..1cef2fe
--- /dev/null
+++ b/src/main/kotlin/util/Optionalutil.kt
@@ -0,0 +1,5 @@
+package moe.nea.firmament.util
+
+import java.util.Optional
+
+fun <T : Any> T?.intoOptional(): Optional<T> = Optional.ofNullable(this)
diff --git a/src/main/kotlin/util/collections/InstanceList.kt b/src/main/kotlin/util/collections/InstanceList.kt
new file mode 100644
index 0000000..fd8c786
--- /dev/null
+++ b/src/main/kotlin/util/collections/InstanceList.kt
@@ -0,0 +1,57 @@
+package moe.nea.firmament.util.collections
+
+import java.lang.ref.ReferenceQueue
+import java.lang.ref.WeakReference
+
+class InstanceList<T : Any>(val name: String) {
+    val queue = object : ReferenceQueue<T>() {}
+    val set = mutableSetOf<Ref>()
+
+    val size: Int
+        get() {
+            clearOldReferences()
+            return set.size
+        }
+
+    fun clearOldReferences() {
+        while (true) {
+            val reference = queue.poll() ?: break
+            set.remove(reference)
+        }
+    }
+
+    fun getAll(): List<T> {
+        clearOldReferences()
+        return set.mapNotNull { it.get() }
+    }
+
+    fun add(t: T) {
+        set.add(Ref(t))
+    }
+
+    init {
+        if (init)
+            allInstances.add(this)
+    }
+
+    inner class Ref(referent: T) : WeakReference<T>(referent) {
+        val hashCode = System.identityHashCode(referent)
+        override fun equals(other: Any?): Boolean {
+            return other is InstanceList<*>.Ref && hashCode == other.hashCode && get() === other.get()
+        }
+
+        override fun hashCode(): Int {
+            return hashCode
+        }
+    }
+
+    companion object {
+        private var init = false
+        val allInstances = InstanceList<InstanceList<*>>("InstanceLists")
+
+        init {
+            init = true
+            allInstances.add(allInstances)
+        }
+    }
+}
diff --git a/src/main/kotlin/util/MutableMapWithMaxSize.kt b/src/main/kotlin/util/collections/MutableMapWithMaxSize.kt
index 067e652..218bc55 100644
--- a/src/main/kotlin/util/MutableMapWithMaxSize.kt
+++ b/src/main/kotlin/util/collections/MutableMapWithMaxSize.kt
@@ -1,5 +1,7 @@
 
-package moe.nea.firmament.util
+package moe.nea.firmament.util.collections
+
+import moe.nea.firmament.util.IdentityCharacteristics
 
 fun <K, V> mutableMapWithMaxSize(maxSize: Int): MutableMap<K, V> = object : LinkedHashMap<K, V>() {
     override fun removeEldestEntry(eldest: MutableMap.MutableEntry<K, V>): Boolean {
diff --git a/src/main/kotlin/util/collections/WeakCache.kt b/src/main/kotlin/util/collections/WeakCache.kt
new file mode 100644
index 0000000..38f9886
--- /dev/null
+++ b/src/main/kotlin/util/collections/WeakCache.kt
@@ -0,0 +1,110 @@
+package moe.nea.firmament.util.collections
+
+import java.lang.ref.ReferenceQueue
+import java.lang.ref.WeakReference
+import moe.nea.firmament.features.debug.DebugLogger
+
+/**
+ * Cache class that uses [WeakReferences][WeakReference] to only cache values while there is still a life reference to
+ * the key. Each key can have additional extra data that is used to look up values. That extra data is not required to
+ * be a life reference. The main Key is compared using strict reference equality. This map is not synchronized.
+ */
+class WeakCache<Key : Any, ExtraKey : Any, Value : Any>(val name: String) {
+    private val queue = object : ReferenceQueue<Key>() {}
+    private val map = mutableMapOf<Ref, Value>()
+
+    val size: Int
+        get() {
+            clearOldReferences()
+            return map.size
+        }
+
+    fun clearOldReferences() {
+        var successCount = 0
+        var totalCount = 0
+        while (true) {
+            val reference = queue.poll() ?: break
+            totalCount++
+            if (map.remove(reference) != null)
+                successCount++
+        }
+        if (totalCount > 0)
+            logger.log { "Cleared $successCount/$totalCount references from queue" }
+    }
+
+    fun get(key: Key, extraData: ExtraKey): Value? {
+        clearOldReferences()
+        return map[Ref(key, extraData)]
+    }
+
+    fun put(key: Key, extraData: ExtraKey, value: Value) {
+        clearOldReferences()
+        map[Ref(key, extraData)] = value
+    }
+
+    fun getOrPut(key: Key, extraData: ExtraKey, value: (Key, ExtraKey) -> Value): Value {
+        clearOldReferences()
+        return map.getOrPut(Ref(key, extraData)) { value(key, extraData) }
+    }
+
+    fun clear() {
+        map.clear()
+    }
+
+    init {
+        allInstances.add(this)
+    }
+
+    companion object {
+        val allInstances = InstanceList<WeakCache<*, *, *>>("WeakCaches")
+        private val logger = DebugLogger("WeakCache")
+        fun <Key : Any, Value : Any> memoize(name: String, function: (Key) -> Value):
+            CacheFunction.NoExtraData<Key, Value> {
+            return CacheFunction.NoExtraData(WeakCache(name), function)
+        }
+
+        fun <Key : Any, ExtraKey : Any, Value : Any> memoize(name: String, function: (Key, ExtraKey) -> Value):
+            CacheFunction.WithExtraData<Key, ExtraKey, Value> {
+            return CacheFunction.WithExtraData(WeakCache(name), function)
+        }
+    }
+
+    inner class Ref(
+        weakInstance: Key,
+        val extraData: ExtraKey,
+    ) : WeakReference<Key>(weakInstance, queue) {
+        val hashCode = System.identityHashCode(weakInstance) * 31 + extraData.hashCode()
+        override fun equals(other: Any?): Boolean {
+            if (other !is WeakCache<*, *, *>.Ref) return false
+            return other.hashCode == this.hashCode
+                && other.get() === this.get()
+                && other.extraData == this.extraData
+        }
+
+        override fun hashCode(): Int {
+            return hashCode
+        }
+    }
+
+    interface CacheFunction {
+        val cache: WeakCache<*, *, *>
+
+        data class NoExtraData<Key : Any, Value : Any>(
+            override val cache: WeakCache<Key, Unit, Value>,
+            val wrapped: (Key) -> Value,
+        ) : CacheFunction, (Key) -> Value {
+            override fun invoke(p1: Key): Value {
+                return cache.getOrPut(p1, Unit, { a, _ -> wrapped(a) })
+            }
+        }
+
+        data class WithExtraData<Key : Any, ExtraKey : Any, Value : Any>(
+            override val cache: WeakCache<Key, ExtraKey, Value>,
+            val wrapped: (Key, ExtraKey) -> Value,
+        ) : CacheFunction, (Key, ExtraKey) -> Value {
+            override fun invoke(p1: Key, p2: ExtraKey): Value {
+                return cache.getOrPut(p1, p2, wrapped)
+            }
+        }
+    }
+}
diff --git a/src/main/kotlin/util/listutil.kt b/src/main/kotlin/util/collections/listutil.kt
index 73cb23e..7f85a30 100644
--- a/src/main/kotlin/util/listutil.kt
+++ b/src/main/kotlin/util/collections/listutil.kt
@@ -1,5 +1,5 @@
 
-package moe.nea.firmament.util
+package moe.nea.firmament.util.collections
 
 fun <T, R> List<T>.lastNotNullOfOrNull(func: (T) -> R?): R? {
     for (i in indices.reversed()) {