Improve "Fix region threading with entity ai data access" patch

shiroha-server/minecraft-patches/features/0011-Fix-region-threading-with-entity-ai-data-access.patch

@@ -150,7 +150,7 @@ index 16017d819c28b077f732e2ef571eac179d24e323..2e248e4fcdec5830d83ecabf3df16311
          if (blockState == null) return false; // Paper - Prevent sync chunk loads when villagers try to find beds
          return target.pos().closerToCenterThan(body.position(), 2.0) && blockState.is(BlockTags.BEDS) && !blockState.getValue(BedBlock.OCCUPIED);
 diff --git a/net/minecraft/world/entity/ai/memory/MemorySlot.java b/net/minecraft/world/entity/ai/memory/MemorySlot.java
-index 88a89b4c72cc99dc89d3f3cc928b2dfda4125759..b4bdc7eff97b058ee8ae2c9e4321b53c824478c5 100644
+index 88a89b4c72cc99dc89d3f3cc928b2dfda4125759..9f85ba205d4e7b69e56ba538c00b1be6e1a5a617 100644
 --- a/net/minecraft/world/entity/ai/memory/MemorySlot.java
 +++ b/net/minecraft/world/entity/ai/memory/MemorySlot.java
 @@ -13,7 +13,7 @@ public class MemorySlot<T> {
@@ -162,7 +162,7 @@ index 88a89b4c72cc99dc89d3f3cc928b2dfda4125759..b4bdc7eff97b058ee8ae2c9e4321b53c
          if (this.hasValue() && this.canExpire()) {
              if (this.hasExpired()) {
                  this.clear();
-@@ -21,6 +21,41 @@ public class MemorySlot<T> {
+@@ -21,6 +21,57 @@ public class MemorySlot<T> {
                  this.timeToLive--;
              }
          }
@@ -176,6 +176,22 @@ index 88a89b4c72cc99dc89d3f3cc928b2dfda4125759..b4bdc7eff97b058ee8ae2c9e4321b53c
 +            }
 +        }
 +
++        // type: list of entity
++        if (io.nanachiyo0721.shiroha.config.modules.fixes.ForceCleanupEntityBrainMemoryConfig.enabledForEntityListed && this.value instanceof java.util.List<?> list) {
++            // check first if it's matched type
++            if (!list.isEmpty() && list.getFirst() instanceof net.minecraft.world.entity.Entity) {
++                // matched, do check
++                for (Object entityInObject : list) {
++                    final net.minecraft.world.entity.Entity entity = (net.minecraft.world.entity.Entity) entityInObject;
++
++                    if (!ca.spottedleaf.moonrise.common.util.TickThread.isTickThreadFor(entity)) {
++                        this.clear();
++                        break;
++                    }
++                }
++            }
++        }
++
 +        // type: block_pos
 +        if (io.nanachiyo0721.shiroha.config.modules.fixes.ForceCleanupEntityBrainMemoryConfig.enabledForBlockPos && this.value instanceof net.minecraft.core.BlockPos blockPos) {
 +            if (!ca.spottedleaf.moonrise.common.util.TickThread.isTickThreadFor(ownerLevel, blockPos)) {
@@ -204,6 +220,59 @@ index 88a89b4c72cc99dc89d3f3cc928b2dfda4125759..b4bdc7eff97b058ee8ae2c9e4321b53c
      }
  
      public static <T> MemorySlot<T> create() {
+diff --git a/net/minecraft/world/entity/ai/memory/NearestVisibleLivingEntities.java b/net/minecraft/world/entity/ai/memory/NearestVisibleLivingEntities.java
+index 5c967b55f6ad3b660e9cdf74fadd90c5ea67afb9..7af28b58afa6538aa47fc25da4f3b911942b7c5a 100644
+--- a/net/minecraft/world/entity/ai/memory/NearestVisibleLivingEntities.java
++++ b/net/minecraft/world/entity/ai/memory/NearestVisibleLivingEntities.java
+@@ -39,7 +39,7 @@ public class NearestVisibleLivingEntities {
+ 
+     public Optional<LivingEntity> findClosest(final Predicate<LivingEntity> filter) {
+         for (LivingEntity nearbyEntity : this.nearbyEntities) {
+-            if (filter.test(nearbyEntity) && this.lineOfSightTest.test(nearbyEntity)) {
++            if (this.getExtraFilterForThreadCheck(filter).test(nearbyEntity) && this.lineOfSightTest.test(nearbyEntity)) { // Shiroha - Fix region threading with entity ai data access (one side is enough for thread checking)
+                 return Optional.of(nearbyEntity);
+             }
+         }
+@@ -48,24 +48,35 @@ public class NearestVisibleLivingEntities {
+     }
+ 
+     public Iterable<LivingEntity> findAll(final Predicate<LivingEntity> filter) {
+-        return Iterables.filter(this.nearbyEntities, entity -> filter.test(entity) && this.lineOfSightTest.test(entity));
++        return Iterables.filter(this.nearbyEntities, entity -> this.getExtraFilterForThreadCheck(filter).test(entity) && this.lineOfSightTest.test(entity)); // Shiroha - Fix region threading with entity ai data access (one side is enough for thread checking)
+     }
+ 
+     public Stream<LivingEntity> find(final Predicate<LivingEntity> filter) {
+-        return this.nearbyEntities.stream().filter(entity -> filter.test(entity) && this.lineOfSightTest.test(entity));
++        return this.nearbyEntities.stream().filter(entity -> this.getExtraFilterForThreadCheck(filter).test(entity) && this.lineOfSightTest.test(entity)); // Shiroha - Fix region threading with entity ai data access (one side is enough for thread checking)
+     }
+ 
+     public boolean contains(final LivingEntity targetEntity) {
+-        return this.nearbyEntities.contains(targetEntity) && this.lineOfSightTest.test(targetEntity);
++        return this.nearbyEntities.contains(targetEntity) && this.getExtraFilterForThreadCheck(this.lineOfSightTest).test(targetEntity); // Shiroha - Fix region threading with entity ai data access
+     }
+ 
+     public boolean contains(final Predicate<LivingEntity> filter) {
+         for (LivingEntity nearbyEntity : this.nearbyEntities) {
+-            if (filter.test(nearbyEntity) && this.lineOfSightTest.test(nearbyEntity)) {
++            if (this.getExtraFilterForThreadCheck(filter).test(nearbyEntity) && this.lineOfSightTest.test(nearbyEntity)) { // Shiroha - Fix region threading with entity ai data access (one side is enough for thread checking)
+                 return true;
+             }
+         }
+ 
+         return false;
+     }
++    // Shiroha start - Fix region threading with entity ai data access
++    private Predicate<LivingEntity> getExtraFilterForThreadCheck(Predicate<LivingEntity> original) {
++        return io.nanachiyo0721.shiroha.config.modules.fixes.ForceCleanupEntityBrainMemoryConfig.enableForNearestLivingEntities ? ent -> {
++            if (!ca.spottedleaf.moonrise.common.util.TickThread.isTickThreadFor(ent)) {
++                return false;
++            }
++
++            return original.test(ent);
++        } : original;
++    }
++    // Shiroha end - Fix region threading with entity ai data access
+ }
 diff --git a/net/minecraft/world/entity/ai/navigation/FlyingPathNavigation.java b/net/minecraft/world/entity/ai/navigation/FlyingPathNavigation.java
 index e44814cfb6afb594456b8215bd13a92e93de2c85..d4b65f13c137495d436b7e8133e0ce09ee2e08c0 100644
 --- a/net/minecraft/world/entity/ai/navigation/FlyingPathNavigation.java

shiroha-server/src/main/java/io/nanachiyo0721/shiroha/config/modules/fixes/ForceCleanupEntityBrainMemoryConfig.java

@@ -10,9 +10,15 @@ public class ForceCleanupEntityBrainMemoryConfig implements IConfigModule {
     @ConfigInfo(name = "enabled_for_entity", comments = "When enabled, the entity's brain will clean the memory which is typed of entity and not belong to current tickregion")
     public static boolean enabledForEntity = false;
 
+    @ConfigInfo(name = "enabled_for_entity_listed", comments = "When enabled, the entity's brain will clean the memory which is typed of list of entity and not belong to current tickregion")
+    public static boolean enabledForEntityListed = false;
+
     @ConfigInfo(name = "enabled_for_block_pos", comments = "When enabled, the entity's brain will clean the memory which is typed of block_pos and not belong to current tickregion")
     public static boolean enabledForBlockPos = false;
 
     @ConfigInfo(name = "enabled_for_position_tracker", comments = "When enabled, the entity's brain will clean the memory which is typed of position_tracker and not belong to current tickregion")
     public static boolean enabledForPositionTracker = false;
+
+    @ConfigInfo(name = "enable_for_nearest_living_entities", comments = "When enabled, any entities out of current region will be ignored for memory inspection")
+    public static boolean enableForNearestLivingEntities = false;
 }