Make use of a reference queue

2 files changed, 29 insertions, 26 deletions

diff --git a/src/main/java/moe/nea/caelo/mixin/PatchCustomItemModelCache.java b/src/main/java/moe/nea/caelo/mixin/PatchCustomItemModelCache.java
index afeafef..3374768 100644
--- a/src/main/java/moe/nea/caelo/mixin/PatchCustomItemModelCache.java
+++ b/src/main/java/moe/nea/caelo/mixin/PatchCustomItemModelCache.java
@@ -16,27 +16,26 @@ import org.spongepowered.asm.mixin.injection.callback.LocalCapture;
 @Mixin(value = CustomItems.class, remap = false)
 public class PatchCustomItemModelCache {
 
-    @SuppressWarnings("InvalidInjectorMethodSignature")
-    @Inject(method = "getCustomItemProperties", at = @At("HEAD"), cancellable = true)
-    private static void overrideCustomItemProperties(
-            ItemStack itemStack, int type,
-            CallbackInfoReturnable<CustomItemProperties> cir) {
-        OptifineCustomItemCache.retrieveCacheHit(itemStack, type, cir);
-    }
+	@Inject(method = "getCustomItemProperties", at = @At("HEAD"), cancellable = true)
+	private static void overrideCustomItemProperties(
+			ItemStack itemStack, int type,
+			CallbackInfoReturnable<CustomItemProperties> cir) {
+		OptifineCustomItemCache.retrieveCacheHit(itemStack, type, cir);
+	}
 
-    @Inject(method = "getCustomItemProperties", at = @At(value = "RETURN", ordinal = 2),
-            locals = LocalCapture.CAPTURE_FAILHARD)
-    private static void storeCustomItemProperties(
-            ItemStack itemStack, int type, CallbackInfoReturnable<CustomItemProperties> cir,
-            Item item, int itemId, CustomItemProperties[] cips, int i, CustomItemProperties cip) {
-        OptifineCustomItemCache.storeCustomItemProperties(itemStack, type, cip);
-    }
+	@Inject(method = "getCustomItemProperties", at = @At(value = "RETURN", ordinal = 2),
+			locals = LocalCapture.CAPTURE_FAILHARD)
+	private static void storeCustomItemProperties(
+			ItemStack itemStack, int type, CallbackInfoReturnable<CustomItemProperties> cir,
+			Item item, int itemId, CustomItemProperties[] cips, int i, CustomItemProperties cip) {
+		OptifineCustomItemCache.storeCustomItemProperties(itemStack, type, cip);
+	}
 
-    @Inject(method = "getCustomItemProperties", at = @At(value = "RETURN", ordinal = 3))
-    private static void storeCustomItemProperties(
-            ItemStack itemStack, int type, CallbackInfoReturnable<CustomItemProperties> cir) {
-        OptifineCustomItemCache.storeNoCustomItemProperties(itemStack, type);
-    }
+	@Inject(method = "getCustomItemProperties", at = @At(value = "RETURN", ordinal = 3))
+	private static void storeCustomItemProperties(
+			ItemStack itemStack, int type, CallbackInfoReturnable<CustomItemProperties> cir) {
+		OptifineCustomItemCache.storeNoCustomItemProperties(itemStack, type);
+	}
 
 
 }
\ 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 60ef545..244eff3 100644
--- a/src/main/java/moe/nea/caelo/optifine/OptifineCustomItemCache.kt
+++ b/src/main/java/moe/nea/caelo/optifine/OptifineCustomItemCache.kt
@@ -9,6 +9,7 @@ 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.ReferenceQueue
 import java.lang.ref.WeakReference
 
 object OptifineCustomItemCache {
@@ -26,9 +27,14 @@ object OptifineCustomItemCache {
 		}
 	}
 
+	val referenceQueue = ReferenceQueue<ItemStack>()
+
+	class CacheKeyReference(val cacheKey: CacheKey, itemStack: ItemStack) :
+		WeakReference<ItemStack>(itemStack, referenceQueue)
+
 	class CacheKey(itemStack: ItemStack, val type: Int) {
 		val hashCode = System.identityHashCode(itemStack) * 31 + type
-		val ref = WeakReference(itemStack)
+		val ref = CacheKeyReference(this, itemStack)
 
 		override fun equals(other: Any?): Boolean {
 			if (other === this) return true
@@ -60,12 +66,10 @@ object OptifineCustomItemCache {
 	@SubscribeEvent
 	fun onTick(event: NeaTickEvent) {
 		var removeCount = 0
-		val it = map.iterator()
-		while (it.hasNext()) {
-			if (!it.next().key.isPresent()) {
-				it.remove()
-				removeCount++
-			}
+		while (true) {
+			val ref = referenceQueue.poll() as CacheKeyReference? ?: break
+			removeCount++
+			map.remove(ref.cacheKey)
 		}
 		cacheStats.size = map.size
 		cacheStats.removals = removeCount