3 files changed, 71 insertions, 6 deletions
diff --git a/src/main/kotlin/events/CustomItemModelEvent.kt b/src/main/kotlin/events/CustomItemModelEvent.kt
index 21ee326..7b86980 100644
--- a/src/main/kotlin/events/CustomItemModelEvent.kt
+++ b/src/main/kotlin/events/CustomItemModelEvent.kt
@@ -1,10 +1,13 @@
 package moe.nea.firmament.events
 
+import java.util.Objects
 import java.util.Optional
 import kotlin.jvm.optionals.getOrNull
+import net.minecraft.component.DataComponentTypes
 import net.minecraft.item.ItemStack
 import net.minecraft.util.Identifier
 import moe.nea.firmament.util.collections.WeakCache
+import moe.nea.firmament.util.collections.WeakCache.CacheFunction
 import moe.nea.firmament.util.mc.IntrospectableItemModelManager
 
 // TODO: assert an order on these events
@@ -14,7 +17,36 @@ data class CustomItemModelEvent(
 	var overrideModel: Identifier? = null,
 ) : FirmamentEvent() {
 	companion object : FirmamentEventBus<CustomItemModelEvent>() {
-		val cache = WeakCache.memoize("ItemModelIdentifier", ::getModelIdentifier0)
+		val weakCache =
+			object : WeakCache<ItemStack, IntrospectableItemModelManager, Optional<Identifier>>("ItemModelIdentifier") {
+				override fun mkRef(
+					key: ItemStack,
+					extraData: IntrospectableItemModelManager
+				): WeakCache<ItemStack, IntrospectableItemModelManager, Optional<Identifier>>.Ref {
+					return IRef(key, extraData)
+				}
+
+				inner class IRef(weakInstance: ItemStack, data: IntrospectableItemModelManager) :
+					Ref(weakInstance, data) {
+					override fun shouldBeEvicted(): Boolean = false
+					val isSimpleStack = weakInstance.componentChanges.isEmpty || (weakInstance.componentChanges.size() == 1 && weakInstance.get(
+						DataComponentTypes.CUSTOM_DATA)?.isEmpty == true)
+					val item = weakInstance.item
+					override fun hashCode(): Int {
+						if (isSimpleStack)
+							return Objects.hash(item, extraData)
+						return super.hashCode()
+					}
+
+					override fun equals(other: Any?): Boolean {
+						if (other is IRef && isSimpleStack) {
+							return other.isSimpleStack && item == other.item
+						}
+						return super.equals(other)
+					}
+				}
+			}
+		val cache = CacheFunction.WithExtraData(weakCache, ::getModelIdentifier0)
 
 		@JvmStatic
 		fun getModelIdentifier(itemStack: ItemStack?, itemModelManager: IntrospectableItemModelManager): Identifier? {
diff --git a/src/main/kotlin/events/EntityUpdateEvent.kt b/src/main/kotlin/events/EntityUpdateEvent.kt
index 27a90f9..fec2fa5 100644
--- a/src/main/kotlin/events/EntityUpdateEvent.kt
+++ b/src/main/kotlin/events/EntityUpdateEvent.kt
@@ -7,6 +7,8 @@ import net.minecraft.entity.LivingEntity
 import net.minecraft.entity.data.DataTracker
 import net.minecraft.item.ItemStack
 import net.minecraft.network.packet.s2c.play.EntityAttributesS2CPacket
+import moe.nea.firmament.annotations.Subscribe
+import moe.nea.firmament.util.MC
 
 /**
  * This event is fired when some entity properties are updated.
@@ -15,7 +17,27 @@ import net.minecraft.network.packet.s2c.play.EntityAttributesS2CPacket
  * *after* the values have been applied to the entity.
  */
 sealed class EntityUpdateEvent : FirmamentEvent() {
-	companion object : FirmamentEventBus<EntityUpdateEvent>()
+	companion object : FirmamentEventBus<EntityUpdateEvent>() {
+		@Subscribe
+		fun onPlayerInventoryUpdate(event: PlayerInventoryUpdate) {
+			val p = MC.player ?: return
+			val updatedSlots = listOf(
+				EquipmentSlot.HEAD to 39,
+				EquipmentSlot.CHEST to 38,
+				EquipmentSlot.LEGS to 37,
+				EquipmentSlot.FEET to 36,
+				EquipmentSlot.OFFHAND to 40,
+				EquipmentSlot.MAINHAND to p.inventory.selectedSlot, // TODO: also equipment update when you swap your selected slot perhaps
+			).mapNotNull { (slot, stackIndex) ->
+				val slotIndex = p.playerScreenHandler.getSlotIndex(p.inventory, stackIndex).asInt
+				event.getOrNull(slotIndex)?.let {
+					Pair.of(slot, it)
+				}
+			}
+			if (updatedSlots.isNotEmpty())
+				publish(EquipmentUpdate(p, updatedSlots))
+		}
+	}
 
 	abstract val entity: Entity
 
diff --git a/src/main/kotlin/events/PlayerInventoryUpdate.kt b/src/main/kotlin/events/PlayerInventoryUpdate.kt
index 6e8203a..88439a9 100644
--- a/src/main/kotlin/events/PlayerInventoryUpdate.kt
+++ b/src/main/kotlin/events/PlayerInventoryUpdate.kt
@@ -1,11 +1,22 @@
-
 package moe.nea.firmament.events
 
 import net.minecraft.item.ItemStack
 
 sealed class PlayerInventoryUpdate : FirmamentEvent() {
-    companion object : FirmamentEventBus<PlayerInventoryUpdate>()
-    data class Single(val slot: Int, val stack: ItemStack) : PlayerInventoryUpdate()
-    data class Multi(val contents: List<ItemStack>) : PlayerInventoryUpdate()
+	companion object : FirmamentEventBus<PlayerInventoryUpdate>()
+	data class Single(val slot: Int, val stack: ItemStack) : PlayerInventoryUpdate() {
+		override fun getOrNull(slot: Int): ItemStack? {
+			if (slot == this.slot) return stack
+			return null
+		}
+
+	}
+
+	data class Multi(val contents: List<ItemStack>) : PlayerInventoryUpdate() {
+		override fun getOrNull(slot: Int): ItemStack? {
+			return contents.getOrNull(slot)
+		}
+	}
 
+	abstract fun getOrNull(slot: Int): ItemStack?
 }