Improve optifine cache eviction

2 files changed, 163 insertions, 137 deletions

diff --git a/build.gradle.kts b/build.gradle.kts
index 48d55d4..7b2378b 100644
--- a/build.gradle.kts
+++ b/build.gradle.kts
@@ -2,131 +2,135 @@ import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
 import xyz.wagyourtail.unimined.api.minecraft.task.RemapJarTask
 
 plugins {
-    kotlin("jvm") version "1.9.22"
-    `maven-publish`
-    id("com.github.johnrengelman.shadow") version "7.1.2"
-    id("xyz.wagyourtail.unimined") version "1.2.0-SNAPSHOT"
+	kotlin("jvm") version "1.9.22"
+	`maven-publish`
+	id("com.github.johnrengelman.shadow") version "7.1.2"
+	id("xyz.wagyourtail.unimined") version "1.2.0-SNAPSHOT"
 }
 
 group = "moe.nea"
 version = "1.0-SNAPSHOT"
 
 repositories {
-    maven("https://jitpack.io")
-    maven("https://repo.polyfrost.cc/releases")
-    maven("https://maven.notenoughupdates.org/releases/")
-    mavenCentral()
-    maven("https://nea.moe/redir-repo") {
-        metadataSources { artifact() }
-        content {
-            includeGroup("optifine")
-        }
-    }
+	maven("https://jitpack.io")
+	maven("https://repo.polyfrost.cc/releases")
+	maven("https://maven.notenoughupdates.org/releases/")
+	mavenCentral()
+	maven("https://nea.moe/redir-repo") {
+		metadataSources { artifact() }
+		content {
+			includeGroup("optifine")
+		}
+	}
 }
 
 val optifineConfig by configurations.creating {
 }
 configurations.compileOnly {
-    extendsFrom(optifineConfig)
+	extendsFrom(optifineConfig)
 }
 
 unimined.minecraft {
-    version("1.8.9")
-    mappings {
-        searge()
-        mcp("stable", "22-1.8.9")
-    }
-    minecraftForge {
-        loader("11.15.1.2318-1.8.9")
-        mixinConfig("veloxcaelo.mixins.json")
-    }
-    mods {
-        remap(optifineConfig) {
-            namespace("official")
-        }
-    }
-    runs {
-        this.config("client") {
-            this.args.addAll(
-                listOf(
-                    "--mods", optifineConfig.resolve().joinToString(",") { it.toRelativeString(this.workingDir) },
-                    "--tweakClass", "org.spongepowered.asm.launch.MixinTweaker",
-                    "--tweakClass", "io.github.notenoughupdates.moulconfig.tweaker.DevelopmentResourceTweaker",
-                )
-            )
-            this.env.put(
-                "LD_LIBRARY_PATH",
-                ":/nix/store/agp6lqznayysqvqkx4k1ggr8n1rsyi8c-gcc-13.2.0-lib/lib:/nix/store/ldi0rb00gmbdg6915lhch3k3b3ib460z-libXcursor-1.2.2/lib:/nix/store/8xbbv82pabjcbj30vrna4gcz4g9q97z4-libXrandr-1.5.4/lib:/nix/store/smrb2g0addhgahkfjjl3k8rfd30gdc29-libXxf86vm-1.1.5/lib:/nix/store/lpqy1z1h8li6h3cp9ax6vifl71dks1ff-libglvnd-1.7.0/lib"
-            )
-        }
-    }
+	version("1.8.9")
+	mappings {
+		searge()
+		mcp("stable", "22-1.8.9")
+	}
+	minecraftForge {
+		loader("11.15.1.2318-1.8.9")
+		mixinConfig("veloxcaelo.mixins.json")
+	}
+	mods {
+		remap(optifineConfig) {
+			namespace("official")
+		}
+	}
+	runs {
+		this.config("client") {
+			this.args.addAll(
+				listOf(
+					"--mods", optifineConfig.resolve().joinToString(",") { it.toRelativeString(this.workingDir) },
+					"--tweakClass", "org.spongepowered.asm.launch.MixinTweaker",
+					"--tweakClass", "io.github.notenoughupdates.moulconfig.tweaker.DevelopmentResourceTweaker",
+				)
+			)
+			this.env.put(
+				"LD_LIBRARY_PATH",
+				":/nix/store/agp6lqznayysqvqkx4k1ggr8n1rsyi8c-gcc-13.2.0-lib/lib:/nix/store/ldi0rb00gmbdg6915lhch3k3b3ib460z-libXcursor-1.2.2/lib:/nix/store/8xbbv82pabjcbj30vrna4gcz4g9q97z4-libXrandr-1.5.4/lib:/nix/store/smrb2g0addhgahkfjjl3k8rfd30gdc29-libXxf86vm-1.1.5/lib:/nix/store/lpqy1z1h8li6h3cp9ax6vifl71dks1ff-libglvnd-1.7.0/lib"
+			)
+		}
+	}
 }
 
 
 val shadowModImpl by configurations.creating {
-    configurations.named("modImplementation").get().extendsFrom(this)
+	configurations.named("modImplementation").get().extendsFrom(this)
 }
 val shadowImpl by configurations.creating {
-    configurations.implementation.get().extendsFrom(this)
+	configurations.implementation.get().extendsFrom(this)
 }
 
 dependencies {
-    testImplementation("org.jetbrains.kotlin:kotlin-test")
-    shadowImpl("org.spongepowered:mixin:0.7.11-SNAPSHOT")
-    shadowModImpl("org.notenoughupdates.moulconfig:legacy:3.0.0-beta.7")
-    optifineConfig("optifine:optifine:1.8.9")
-    compileOnly("org.jetbrains:annotations:24.1.0")
+	testImplementation("org.jetbrains.kotlin:kotlin-test")
+	shadowImpl("org.spongepowered:mixin:0.7.11-SNAPSHOT") {
+		isTransitive = false
+	}
+	shadowModImpl("org.notenoughupdates.moulconfig:legacy:3.0.0-beta.7") {
+		isTransitive = false
+	}
+	optifineConfig("optifine:optifine:1.8.9")
+	compileOnly("org.jetbrains:annotations:24.1.0")
 }
 
 tasks.test {
-    useJUnitPlatform()
+	useJUnitPlatform()
 }
 
 sourceSets.main {
-    output.setResourcesDir(sourceSets.main.flatMap { it.java.classesDirectory })
-    kotlin.destinationDirectory.set(java.destinationDirectory)
+	output.setResourcesDir(sourceSets.main.flatMap { it.java.classesDirectory })
+	kotlin.destinationDirectory.set(java.destinationDirectory)
 }
 tasks.processResources {
-    filesMatching("*.mixins.json") {
-        this.autoDiscoverMixins(sourceSets.main.get())
-    }
+	filesMatching("*.mixins.json") {
+		this.autoDiscoverMixins(sourceSets.main.get())
+	}
 }
 
 tasks.compileJava {
-    dependsOn(tasks.processResources)
+	dependsOn(tasks.processResources)
 }
 java {
-    targetCompatibility = JavaVersion.VERSION_1_8
-    sourceCompatibility = JavaVersion.VERSION_1_8
+	targetCompatibility = JavaVersion.VERSION_1_8
+	sourceCompatibility = JavaVersion.VERSION_1_8
 }
 tasks.withType(JavaCompile::class) {
-    this.options.encoding = "UTF-8"
+	this.options.encoding = "UTF-8"
 }
 tasks.withType(KotlinCompile::class) {
-    this.compilerOptions {
-        this.jvmTarget.set(org.jetbrains.kotlin.gradle.dsl.JvmTarget.JVM_1_8)
-    }
+	this.compilerOptions {
+		this.jvmTarget.set(org.jetbrains.kotlin.gradle.dsl.JvmTarget.JVM_1_8)
+	}
 }
 
 tasks.withType(Jar::class) {
-    destinationDirectory.set(project.layout.buildDirectory.dir("badjars"))
-    archiveBaseName.set("VeloxCaelo")
-    manifest.attributes.run {
-        this["FMLCorePluginContainsFMLMod"] = "true"
-        this["ForceLoadAsMod"] = "true"
+	destinationDirectory.set(project.layout.buildDirectory.dir("badjars"))
+	archiveBaseName.set("VeloxCaelo")
+	manifest.attributes.run {
+		this["FMLCorePluginContainsFMLMod"] = "true"
+		this["ForceLoadAsMod"] = "true"
 
-        // If you don't want mixins, remove these lines
-        this["TweakClass"] = "org.spongepowered.asm.launch.MixinTweaker"
-        this["MixinConfigs"] = "veloxcaelo.mixins.json"
-    }
+		// If you don't want mixins, remove these lines
+		this["TweakClass"] = "org.spongepowered.asm.launch.MixinTweaker"
+		this["MixinConfigs"] = "veloxcaelo.mixins.json"
+	}
 }
 tasks.shadowJar {
-    archiveClassifier.set("dep-dev")
-    configurations = listOf(shadowImpl, shadowModImpl)
-    relocate("io.github.notenoughupdates.moulconfig", "moe.nea.velox.moulconfig")
-    mergeServiceFiles()
+	archiveClassifier.set("dep-dev")
+	configurations = listOf(shadowImpl, shadowModImpl)
+	relocate("io.github.notenoughupdates.moulconfig", "moe.nea.velox.moulconfig")
+	mergeServiceFiles()
 }
 tasks.named<RemapJarTask>("remapJar") {
-    this.inputFile.set(tasks.shadowJar.flatMap { it.archiveFile })
-    dependsOn((tasks.shadowJar))
+	this.inputFile.set(tasks.shadowJar.flatMap { it.archiveFile })
+	dependsOn((tasks.shadowJar))
 }
\ No newline at end of file
diff --git a/src/main/java/moe/nea/caelo/optifine/OptifineCustomItemCache.kt b/src/main/java/moe/nea/caelo/optifine/OptifineCustomItemCache.kt
index d8021db..d0e8fc4 100644
--- a/src/main/java/moe/nea/caelo/optifine/OptifineCustomItemCache.kt
+++ b/src/main/java/moe/nea/caelo/optifine/OptifineCustomItemCache.kt
@@ -9,73 +9,95 @@ import net.minecraft.item.ItemStack
 import net.minecraftforge.fml.common.eventhandler.SubscribeEvent
 import net.optifine.CustomItemProperties
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable
+import java.lang.ref.WeakReference
 
 object OptifineCustomItemCache {
 
-    init {
-        CaeloCommand.subcommand("opticache") { args ->
-            val cache = cacheSizeHistory.lastOrNull() ?: CacheStats()
-            MC.display("OptiCache stats:")
-            MC.display("- History: §3${cacheSizeHistory.size}")
-            MC.display("- Misses: §c${cache.cacheMisses}")
-            MC.display("- Hits: §a${cache.cacheHits}")
-            MC.display("- Matches: §b${cache.uniquePropertyBearingStacks}")
-        }
-    }
+	init {
+		CaeloCommand.subcommand("opticache") { args ->
+			val cache = cacheSizeHistory.lastOrNull() ?: CacheStats()
+			MC.display("OptiCache stats:")
+			MC.display("- History: §3${cacheSizeHistory.size}")
+			MC.display("- Misses: §c${cache.cacheMisses}")
+			MC.display("- Hits: §a${cache.cacheHits}")
+			MC.display("- Insertions: §b${cache.insertions}")
+			MC.display("- Evictions: §b${cache.removals}")
+			MC.display("- Cache Size: §b${cache.size}")
+		}
+	}
 
-    class CacheKey(val itemStack: ItemStack, val type: Int) {
-        override fun equals(other: Any?): Boolean {
-            if (other !is CacheKey) return false
-            return itemStack === other.itemStack && type == other.type
-        }
+	class CacheKey(val itemStack: WeakReference<ItemStack>, val type: Int) {
+		override fun equals(other: Any?): Boolean {
+			if (other !is CacheKey) return false
+			return itemStack.get() === other.itemStack.get() && type == other.type
+		}
 
-        override fun hashCode(): Int {
-            return System.identityHashCode(itemStack) + type * 31
-        }
-    }
+		override fun hashCode(): Int {
+			return System.identityHashCode(itemStack.get()) * 31 + type
+		}
 
-    data class CacheStats(
-        var cacheHits: Int = 0,
-        var cacheMisses: Int = 0,
-        var uniquePropertyBearingStacks: Int = 0,
-    )
+		fun isPresent(): Boolean {
+			return itemStack.get() != null
+		}
+	}
 
-    private val map = mutableMapOf<CacheKey, CustomItemProperties?>()
-    private val cacheSizeHistory = Histogram<CacheStats>(1000)
-    private var cacheStats = CacheStats()
+	data class CacheStats(
+		var cacheHits: Int = 0,
+		var cacheMisses: Int = 0,
+		var insertions: Int = 0,
+		var size: Int = 0,
+		var removals: Int = 0,
+	)
 
-    @SubscribeEvent
-    fun onTick(event: NeaTickEvent) {
-        cacheSizeHistory.append(cacheStats)
-        cacheStats = CacheStats()
-        map.clear()
-    }
+	private var map = mutableMapOf<CacheKey, CustomItemProperties?>()
+	private val cacheSizeHistory = Histogram<CacheStats>(1000)
+	private var cacheStats = CacheStats()
 
-    @JvmStatic
-    fun retrieveCacheHit(
-        itemStack: ItemStack,
-        type: Int,
-        cir: CallbackInfoReturnable<CustomItemProperties?>
-    ) {
-        if (!CConfig.config.optiCache.citCache)
-            return
-        val key = CacheKey(itemStack, type)
-        if (!map.containsKey(key)) {
-            cacheStats.cacheMisses++
-            return
-        }
-        cacheStats.cacheHits++
-        cir.returnValue = map[key]
-    }
+	@SubscribeEvent
+	fun onTick(event: NeaTickEvent) {
+		var removeCount = 0
+		val nextMap = mutableMapOf<CacheKey, CustomItemProperties?>()
+		for (entry in map) {
+			if (entry.key.isPresent()) {
+				nextMap[entry.key] = entry.value
+			} else {
+				removeCount++
+			}
+		}
+		map = nextMap
+		cacheStats.size = map.size
+		cacheStats.removals = removeCount
+		cacheSizeHistory.append(cacheStats)
+		cacheStats = CacheStats()
 
-    @JvmStatic
-    fun storeCustomItemProperties(itemStack: ItemStack, type: Int, cip: CustomItemProperties) {
-        map[CacheKey(itemStack, type)] = cip
-        cacheStats.uniquePropertyBearingStacks++
-    }
+	}
 
-    @JvmStatic
-    fun storeNoCustomItemProperties(itemStack: ItemStack, type: Int) {
-        map[CacheKey(itemStack, type)] = null
-    }
+	@JvmStatic
+	fun retrieveCacheHit(
+		itemStack: ItemStack,
+		type: Int,
+		cir: CallbackInfoReturnable<CustomItemProperties?>
+	) {
+		if (!CConfig.config.optiCache.citCache)
+			return
+		val key = CacheKey(WeakReference(itemStack), type)
+		if (!map.containsKey(key)) {
+			cacheStats.cacheMisses++
+			return
+		}
+		cacheStats.cacheHits++
+		cir.returnValue = map[key]
+	}
+
+	@JvmStatic
+	fun storeCustomItemProperties(itemStack: ItemStack, type: Int, cip: CustomItemProperties) {
+		map[CacheKey(WeakReference(itemStack), type)] = cip
+		cacheStats.insertions++
+	}
+
+	@JvmStatic
+	fun storeNoCustomItemProperties(itemStack: ItemStack, type: Int) {
+		map[CacheKey(WeakReference(itemStack), type)] = null
+		cacheStats.insertions++
+	}
 }
\ No newline at end of file