Fix raid banner caching incompatibility with minimal modded worlds (e…

….g. SchematicWorld)
CaffeineMC · Jun 16, 2024 · 6f0bdce · 6f0bdce
1 parent 33ea83f
commit 6f0bdce
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 5 deletions.
diff --git a/src/main/java/me/jellysquid/mods/lithium/common/world/LithiumData.java b/src/main/java/me/jellysquid/mods/lithium/common/world/LithiumData.java
@@ -37,7 +37,7 @@ record Data(
         public Data(World world) {
             this(
                     new Long2ReferenceOpenHashMap<>(),
-                    Raid.getOminousBanner(world.getRegistryManager().getWrapperOrThrow(RegistryKeys.BANNER_PATTERN)),
+                    world.getRegistryManager().getOptionalWrapper(RegistryKeys.BANNER_PATTERN).map(Raid::getOminousBanner).orElse(null),
                     new ReferenceOpenHashSet<>(),
                     new LithiumInterner<>(),
                     new LithiumInterner<>()

diff --git a/src/main/java/me/jellysquid/mods/lithium/mixin/ai/raid/PickupBannerAsLeaderGoalMixin.java b/src/main/java/me/jellysquid/mods/lithium/mixin/ai/raid/PickupBannerAsLeaderGoalMixin.java
@@ -5,6 +5,7 @@
 import net.minecraft.entity.raid.RaiderEntity;
 import net.minecraft.item.ItemStack;
 import net.minecraft.registry.RegistryEntryLookup;
+import net.minecraft.village.raid.Raid;
 import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Shadow;
@@ -26,6 +27,10 @@ public class PickupBannerAsLeaderGoalMixin<T extends RaiderEntity> {
             )
     )
     private ItemStack getOminousBanner(RegistryEntryLookup<BannerPattern> bannerPatternLookup) {
-        return ((LithiumData) this.actor.getWorld()).lithium$getData().ominousBanner();
+        ItemStack ominousBanner = ((LithiumData) this.actor.getWorld()).lithium$getData().ominousBanner();
+        if (ominousBanner == null) {
+            ominousBanner = Raid.getOminousBanner(bannerPatternLookup);
+        }
+        return ominousBanner;
     }
 }
diff --git a/src/main/java/me/jellysquid/mods/lithium/mixin/ai/raid/RaiderEntityMixin.java b/src/main/java/me/jellysquid/mods/lithium/mixin/ai/raid/RaiderEntityMixin.java
@@ -8,6 +8,8 @@
 import net.minecraft.entity.raid.RaiderEntity;
 import net.minecraft.item.ItemStack;
 import net.minecraft.registry.RegistryEntryLookup;
+import net.minecraft.registry.RegistryKeys;
+import net.minecraft.village.raid.Raid;
 import net.minecraft.world.World;
 import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
@@ -27,8 +29,15 @@ public abstract class RaiderEntityMixin extends Entity {
 
     static {
         // The call to Raid#getOminousBanner() is very expensive, use a cached banner during AI ticking
-        OBTAINABLE_OMINOUS_BANNER_PREDICATE = (itemEntity) -> !itemEntity.cannotPickup() && itemEntity.isAlive() &&
-                ItemStack.areEqual(itemEntity.getStack(), ((LithiumData) itemEntity.getWorld()).lithium$getData().ominousBanner());
+        OBTAINABLE_OMINOUS_BANNER_PREDICATE = (itemEntity) -> {
+            ItemStack ominousBanner = ((LithiumData) itemEntity.getWorld()).lithium$getData().ominousBanner();
+            if (ominousBanner == null) {
+                ominousBanner = Raid.getOminousBanner(itemEntity.getRegistryManager().getWrapperOrThrow(RegistryKeys.BANNER_PATTERN));
+            }
+
+            return !itemEntity.cannotPickup() && itemEntity.isAlive() &&
+                    ItemStack.areEqual(itemEntity.getStack(), ominousBanner);
+        };
     }
 
     public RaiderEntityMixin(EntityType<?> type, World world) {
@@ -40,6 +49,10 @@ public RaiderEntityMixin(EntityType<?> type, World world) {
             at = @At(value = "INVOKE", target = "Lnet/minecraft/village/raid/Raid;getOminousBanner(Lnet/minecraft/registry/RegistryEntryLookup;)Lnet/minecraft/item/ItemStack;")
     )
     private ItemStack getOminousBanner(RegistryEntryLookup<BannerPattern> bannerPatternLookup) {
-        return ((LithiumData) this.getWorld()).lithium$getData().ominousBanner();
+        ItemStack ominousBanner = ((LithiumData) this.getWorld()).lithium$getData().ominousBanner();
+        if (ominousBanner == null) {
+            ominousBanner = Raid.getOminousBanner(bannerPatternLookup);
+        }
+        return ominousBanner;
     }
 }