Optimized ItemEntity stacking using item count categorization

CaffeineMC · Nov 8, 2023 · 8310049 · 8310049
1 parent 3343935
commit 8310049
Show file tree

Hide file tree

Showing 8 changed files with 381 additions and 39 deletions.
diff --git a/src/main/java/me/jellysquid/mods/lithium/common/entity/item/ItemEntityCategorizingList.java b/src/main/java/me/jellysquid/mods/lithium/common/entity/item/ItemEntityCategorizingList.java
diff --git a/src/main/java/me/jellysquid/mods/lithium/common/entity/item/ItemStackSubscriber.java b/src/main/java/me/jellysquid/mods/lithium/common/entity/item/ItemStackSubscriber.java
@@ -1,6 +1,11 @@
 package me.jellysquid.mods.lithium.common.entity.item;
 
+import net.minecraft.entity.ItemEntity;
+import net.minecraft.item.ItemStack;
+
 public interface ItemStackSubscriber {
 
-    void lithium$notifyBeforeCountChange(int slot, int newCount);
+    void lithium$notifyBeforeCountChange(ItemStack itemStack, int slot, int newCount);
+
+    default void lithium$notifyItemEntityStackSwap(ItemEntity itemEntity, ItemStack oldStack) {}
 }
diff --git a/src/main/java/me/jellysquid/mods/lithium/common/entity/item/ItemStackSubscriberMulti.java b/src/main/java/me/jellysquid/mods/lithium/common/entity/item/ItemStackSubscriberMulti.java
@@ -1,5 +1,8 @@
 package me.jellysquid.mods.lithium.common.entity.item;
 
+import net.minecraft.entity.ItemEntity;
+import net.minecraft.item.ItemStack;
+
 import java.util.Arrays;
 
 /**
@@ -84,11 +87,18 @@ public int getSlot(ItemStackSubscriber subscriber) {
     }
 
     @Override
-    public void lithium$notifyBeforeCountChange(int slot, int newCount) {
+    public void lithium$notifyBeforeCountChange(ItemStack itemStack, int slot, int newCount) {
         ItemStackSubscriber[] itemStackSubscribers = this.subscribers;
         for (int i = 0; i < itemStackSubscribers.length; i++) {
             ItemStackSubscriber subscriber = itemStackSubscribers[i];
-            subscriber.lithium$notifyBeforeCountChange(this.slots[i], newCount);
+            subscriber.lithium$notifyBeforeCountChange(itemStack, this.slots[i], newCount);
+        }
+    }
+
+    @Override
+    public void lithium$notifyItemEntityStackSwap(ItemEntity itemEntity, ItemStack oldStack) {
+        for (ItemStackSubscriber subscriber : this.subscribers) {
+            subscriber.lithium$notifyItemEntityStackSwap(itemEntity, oldStack);
         }
     }
 }
diff --git a/src/main/java/me/jellysquid/mods/lithium/common/hopper/LithiumStackList.java b/src/main/java/me/jellysquid/mods/lithium/common/hopper/LithiumStackList.java
@@ -95,7 +95,7 @@ public void changedALot() {
 
     }
 
-    public void lithium$notifyBeforeCountChange(int slot, int newCount) {
+    public void lithium$notifyBeforeCountChange(ItemStack itemStack, int slot, int newCount) {
         ItemStack stack = this.get(slot);
         int count = stack.getCount();
         if (newCount <= 0) {

diff --git a/src/main/java/me/jellysquid/mods/lithium/common/hopper/NotifyingItemStack.java b/src/main/java/me/jellysquid/mods/lithium/common/hopper/NotifyingItemStack.java
@@ -1,6 +1,8 @@
 package me.jellysquid.mods.lithium.common.hopper;
 
 import me.jellysquid.mods.lithium.common.entity.item.ItemStackSubscriber;
+import net.minecraft.entity.ItemEntity;
+import net.minecraft.item.ItemStack;
 
 public interface NotifyingItemStack {
     void lithium$subscribe(ItemStackSubscriber subscriber);
@@ -10,4 +12,6 @@ public interface NotifyingItemStack {
     void lithium$unsubscribe(ItemStackSubscriber subscriber);
 
     void lithium$unsubscribeWithIndex(ItemStackSubscriber subscriber, int index);
+
+    void lithium$notifyItemEntityStackSwap(ItemEntity itemEntity, ItemStack oldStack);
 }
diff --git a/src/main/java/me/jellysquid/mods/lithium/common/world/WorldHelper.java b/src/main/java/me/jellysquid/mods/lithium/common/world/WorldHelper.java
@@ -12,6 +12,7 @@
 import me.jellysquid.mods.lithium.mixin.util.accessors.ServerWorldAccessor;
 import net.minecraft.entity.Entity;
 import net.minecraft.entity.ItemEntity;
+import net.minecraft.item.Item;
 import net.minecraft.util.collection.TypeFilterableList;
 import net.minecraft.util.function.LazyIterationConsumer;
 import net.minecraft.util.math.BlockPos;
@@ -88,23 +89,6 @@ public static List<ItemEntity> getItemEntitiesForMerge(SectionedEntityCache<Enti
         return collectedEntities;
     }
 
-    public static void invalidateItemCache(ItemEntity itemEntity) {
-        SectionedEntityCache<Entity> cache = WorldHelper.getEntityCacheOrNull(itemEntity.getWorld());
-        if (cache == null) {
-            return;
-        }
-        long longPos = ChunkSectionPos.toLong(itemEntity.getBlockPos());
-
-        //noinspection unchecked
-        TypeFilterableList<Entity> allEntities = ((EntityTrackingSectionAccessor<Entity>) cache.getTrackingSection(longPos)).getCollection();
-        //noinspection unchecked
-        TypeFilterableListInternalAccess<Entity> interalEntityList = (TypeFilterableListInternalAccess<Entity>) allEntities;
-        List<ItemEntity> itemEntities = interalEntityList.lithium$getOrCreateAllOfTypeRaw(ItemEntity.class);
-        if (itemEntities instanceof ItemEntityCategorizingList categorizingList) {
-             categorizingList.invalidateCache();
-        }
-    }
-
     //Requires util.accessors
     public static SectionedEntityCache<Entity> getEntityCacheOrNull(World world) {
         if (world instanceof ClientWorldAccessor) {

diff --git a/...in/java/me/jellysquid/mods/lithium/mixin/entity/item_entity_stacking/ItemEntityMixin.java b/...in/java/me/jellysquid/mods/lithium/mixin/entity/item_entity_stacking/ItemEntityMixin.java
@@ -1,10 +1,13 @@
 package me.jellysquid.mods.lithium.mixin.entity.item_entity_stacking;
 
+import me.jellysquid.mods.lithium.common.hopper.NotifyingItemStack;
 import me.jellysquid.mods.lithium.common.world.WorldHelper;
 import me.jellysquid.mods.lithium.mixin.util.accessors.ItemStackAccessor;
 import net.minecraft.entity.Entity;
 import net.minecraft.entity.EntityType;
 import net.minecraft.entity.ItemEntity;
+import net.minecraft.entity.data.DataTracker;
+import net.minecraft.entity.data.TrackedData;
 import net.minecraft.item.Item;
 import net.minecraft.item.ItemStack;
 import net.minecraft.util.math.Box;
@@ -13,9 +16,7 @@
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Shadow;
 import org.spongepowered.asm.mixin.injection.At;
-import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.Redirect;
-import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 
 import java.util.List;
 import java.util.function.Predicate;
@@ -44,13 +45,16 @@ private List<ItemEntity> getItems(World world, Class<ItemEntity> itemEntityClass
         return world.getEntitiesByClass(itemEntityClass, box, predicate);
     }
 
-    @Inject(
-            method = "setStack", at = @At("HEAD")
+    @Redirect(
+            method = "setStack", at = @At(value = "INVOKE", target = "Lnet/minecraft/entity/data/DataTracker;set(Lnet/minecraft/entity/data/TrackedData;Ljava/lang/Object;)V")
     )
-    private void invalidateStuff(ItemStack stack, CallbackInfo ci) {
-        Item newItem = ((ItemStackAccessor) (Object) stack).lithium$getItem();
-        if (newItem != ((ItemStackAccessor) (Object) this.getStack()).lithium$getItem()) {
-            WorldHelper.invalidateItemCache((ItemEntity) (Object) this);
+    private <T> void handleItemTypeChange(DataTracker dataTracker, TrackedData<T> key, T newStack) {
+        ItemStack oldStack = this.getStack();
+        dataTracker.set(key, newStack);
+
+        Item newItem = ((ItemStackAccessor) newStack).lithium$getItem();
+        if (newItem != ((ItemStackAccessor) (Object) oldStack).lithium$getItem()) {
+            ((NotifyingItemStack) (Object) oldStack).lithium$notifyItemEntityStackSwap((ItemEntity) (Object) this, oldStack);
         }
     }
 }
diff --git a/src/main/java/me/jellysquid/mods/lithium/mixin/util/item_stack_tracking/ItemStackMixin.java b/src/main/java/me/jellysquid/mods/lithium/mixin/util/item_stack_tracking/ItemStackMixin.java
@@ -3,6 +3,7 @@
 import me.jellysquid.mods.lithium.common.entity.item.ItemStackSubscriber;
 import me.jellysquid.mods.lithium.common.entity.item.ItemStackSubscriberMulti;
 import me.jellysquid.mods.lithium.common.hopper.NotifyingItemStack;
+import net.minecraft.entity.ItemEntity;
 import net.minecraft.item.ItemStack;
 import org.jetbrains.annotations.Nullable;
 import org.spongepowered.asm.mixin.Mixin;
@@ -28,7 +29,7 @@ public abstract class ItemStackMixin implements NotifyingItemStack {
     @ModifyVariable(method = "setCount(I)V", at = @At("HEAD"), argsOnly = true)
     public int updateInventory(int count) {
         if (this.stackChangeSubscriber != null && this.count != count) {
-            this.stackChangeSubscriber.lithium$notifyBeforeCountChange(this.mySlot, count);
+            this.stackChangeSubscriber.lithium$notifyBeforeCountChange((ItemStack) (Object) this, this.mySlot, count);
         }
         return count;
     }
@@ -54,12 +55,12 @@ public int updateInventory(int count) {
     }
 
     @Override
-    public void lithium$unsubscribeWithIndex(ItemStackSubscriber myInventoryList, int index) {
-        if (this.stackChangeSubscriber == myInventoryList) {
+    public void lithium$unsubscribeWithIndex(ItemStackSubscriber subscriber, int index) {
+        if (this.stackChangeSubscriber == subscriber) {
             this.stackChangeSubscriber = null;
             this.mySlot = -1;
         } else if (this.stackChangeSubscriber instanceof ItemStackSubscriberMulti multiSubscriber) {
-            this.stackChangeSubscriber = multiSubscriber.without(myInventoryList, index);
+            this.stackChangeSubscriber = multiSubscriber.without(subscriber, index);
             this.mySlot = multiSubscriber.getSlot(this.stackChangeSubscriber);
         }  //else: no change, since the inventory wasn't subscribed
     }
@@ -72,4 +73,11 @@ public int updateInventory(int count) {
             this.mySlot = -1;
         }
     }
+
+    @Override
+    public void lithium$notifyItemEntityStackSwap(ItemEntity itemEntity, ItemStack oldStack) {
+        if (this.stackChangeSubscriber != null) {
+            this.stackChangeSubscriber.lithium$notifyItemEntityStackSwap(itemEntity, oldStack);
+        }
+    }
 }