Rewrite entity minion/discorded ai changes to be non-destructive

src/main/java/com/minelittlepony/unicopia/entity/Creature.java

@@ -1,6 +1,7 @@
 package com.minelittlepony.unicopia.entity;
 
 import java.util.Optional;
+import java.util.function.BooleanSupplier;
 import java.util.function.Predicate;
 
 import org.jetbrains.annotations.NotNull;
@@ -16,6 +17,7 @@ import com.minelittlepony.unicopia.entity.ai.BreakHeartGoal;
 import com.minelittlepony.unicopia.entity.ai.DynamicTargetGoal;
 import com.minelittlepony.unicopia.entity.ai.EatMuffinGoal;
 import com.minelittlepony.unicopia.entity.ai.FleeExplosionGoal;
+import com.minelittlepony.unicopia.entity.ai.PredicatedGoal;
 import com.minelittlepony.unicopia.entity.ai.PrioritizedActiveTargetGoal;
 import com.minelittlepony.unicopia.entity.ai.TargettingUtil;
 import com.minelittlepony.unicopia.entity.ai.WantItTakeItGoal;
@@ -34,6 +36,7 @@ import net.minecraft.entity.passive.*;
 import net.minecraft.entity.player.PlayerEntity;
 import net.minecraft.nbt.NbtCompound;
 import net.minecraft.nbt.NbtElement;
+import net.minecraft.particle.ParticleTypes;
 import net.minecraft.registry.RegistryWrapper.WrapperLookup;
 import net.minecraft.util.math.MathHelper;
 
@@ -45,13 +48,11 @@ public class Creature extends Living<LivingEntity> implements WeaklyOwned.Mutabl
     private final EntityReference<LivingEntity> owner = new EntityReference<>();
 
     private Optional<GoalSelector> goals = Optional.empty();
-    private Optional<GoalSelector> targets = Optional.empty();
 
     private int eatTimer;
     @Nullable
     private EatMuffinGoal eatMuffinGoal;
 
-    private boolean discordedChanged = true;
     private int smittenTicks;
 
     private final Predicate<LivingEntity> targetPredicate = TargetSelecter.<LivingEntity>validTarget(() -> getOriginatingCaster().getAffinity(), this).and(e -> {
@@ -80,9 +81,6 @@ public class Creature extends Living<LivingEntity> implements WeaklyOwned.Mutabl
     @Override
     public void setMaster(LivingEntity owner) {
         this.owner.set(owner);
-        if (owner != null) {
-            targets.ifPresent(this::initMinionAi);
-        }
     }
 
     public boolean isMinion() {
@@ -104,7 +102,6 @@ public class Creature extends Living<LivingEntity> implements WeaklyOwned.Mutabl
 
     public void setDiscorded(boolean discorded) {
         this.discorded.set(discorded);
-        discordedChanged = true;
     }
 
     @Override
@@ -118,17 +115,12 @@ public class Creature extends Living<LivingEntity> implements WeaklyOwned.Mutabl
         return owner;
     }
 
-    public Optional<GoalSelector> getTargets() {
-        return targets;
-    }
-
     public Optional<GoalSelector> getGoals() {
         return goals;
     }
 
     public void initAi(GoalSelector goals, GoalSelector targets) {
         this.goals = Optional.of(goals);
-        this.targets = Optional.of(targets);
 
         if (entity instanceof MagicImmune) {
             return;
@@ -153,48 +145,31 @@ public class Creature extends Living<LivingEntity> implements WeaklyOwned.Mutabl
             goals.add(3, new FleeExplosionGoal(tameable, 6, 1, 1.2));
         }
 
-        initMinionAi(targets);
-        initDiscordedAi();
-
         if (entity instanceof CreeperEntity mob) {
             goals.add(1, new FleeEntityGoal<>(mob, LivingEntity.class, 10, 1.5, 1.9, AmuletSelectors.ALICORN_AMULET));
         }
         if (entity instanceof PassiveEntity mob) {
             goals.add(1, new FleeEntityGoal<>(mob, LivingEntity.class, 10, 1.1, 1.7, AmuletSelectors.ALICORN_AMULET_AFTER_1_DAYS));
         }
-    }
 
-    private void initMinionAi(GoalSelector targets) {
-        if (entity instanceof MagicImmune || !getMasterReference().isSet()) {
-            return;
-        }
+        final BooleanSupplier isMinion = () -> getMasterReference().isSet();
+        final BooleanSupplier isNotMinion = () -> !isMinion.getAsBoolean();
 
-        clearGoals(targets);
-        targets.add(2, new ActiveEnemyGoal<>(PlayerEntity.class));
-        targets.add(2, new ActiveEnemyGoal<>(HostileEntity.class));
-        targets.add(2, new ActiveEnemyGoal<>(SlimeEntity.class));
-    }
+        PredicatedGoal.applyToAll(targets, isNotMinion);
+        targets.add(2, new PredicatedGoal(new ActiveEnemyGoal<>(PlayerEntity.class), isMinion));
+        targets.add(2, new PredicatedGoal(new ActiveEnemyGoal<>(HostileEntity.class), isMinion));
+        targets.add(2, new PredicatedGoal(new ActiveEnemyGoal<>(SlimeEntity.class), isMinion));
 
-    private void initDiscordedAi() {
-        if (entity instanceof MagicImmune || getMasterReference().isSet() || !isDiscorded()) {
-            return;
-        }
-        targets.ifPresent(this::clearGoals);
-        // the brain drain
-        entity.getBrain().clear();
         if (entity instanceof MobEntity mob) {
-            mob.setTarget(null);
-            goals.ifPresent(goalSelector -> {
-                clearGoals(goalSelector);
-                goalSelector.add(1, new SwimGoal(mob));
+            final BooleanSupplier isDiscorded = () -> isNotMinion.getAsBoolean() && isDiscorded();
+            PredicatedGoal.applyToAll(goals, () -> !isDiscorded.getAsBoolean());
+
+            goals.add(1, new PredicatedGoal(new SwimGoal(mob), isDiscorded));
             if (mob instanceof PathAwareEntity pae) {
-                    goalSelector.add(5, new WanderAroundFarGoal(pae, 0.8));
+                goals.add(5, new PredicatedGoal(new WanderAroundFarGoal(pae, 0.8), isDiscorded));
             }
-                goalSelector.add(6, new LookAtEntityGoal(mob, PlayerEntity.class, 8.0f));
-                goalSelector.add(6, new LookAroundGoal(mob));
-            });
-        } else {
-            goals.ifPresent(this::clearGoals);
+            goals.add(6, new PredicatedGoal(new LookAtEntityGoal(mob, PlayerEntity.class, 8.0f), isDiscorded));
+            goals.add(6, new PredicatedGoal(new LookAroundGoal(mob), isDiscorded));
         }
     }
 
@@ -206,9 +181,8 @@ public class Creature extends Living<LivingEntity> implements WeaklyOwned.Mutabl
 
     @Override
     public boolean beforeUpdate() {
-        if (discordedChanged) {
-            discordedChanged = false;
-            initDiscordedAi();
+        if (isDiscorded()) {
+            spawnParticles(ParticleTypes.ELECTRIC_SPARK, 1);
         }
 
         if (!isClient() && smittenTicks > 0) {
@@ -220,10 +194,6 @@ public class Creature extends Living<LivingEntity> implements WeaklyOwned.Mutabl
         return super.beforeUpdate();
     }
 
-    private void clearGoals(GoalSelector t) {
-        t.clear(g -> true);
-    }
-
     private void updateConsumption() {
         if (isClient()) {
             eatTimer = eating.get();
@@ -312,9 +282,6 @@ public class Creature extends Living<LivingEntity> implements WeaklyOwned.Mutabl
         super.fromNBT(compound, lookup);
         if (compound.contains("master", NbtElement.COMPOUND_TYPE)) {
             getMasterReference().fromNBT(compound.getCompound("master"), lookup);
-            if (getMasterReference().isSet()) {
-                targets.ifPresent(this::initMinionAi);
-            }
         }
         physics.fromNBT(compound, lookup);
         setDiscorded(compound.getBoolean("discorded"));

src/main/java/com/minelittlepony/unicopia/entity/ai/PredicatedGoal.java

@@ -0,0 +1,93 @@
+package com.minelittlepony.unicopia.entity.ai;
+
+import java.util.EnumSet;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.function.BooleanSupplier;
+
+import org.jetbrains.annotations.Nullable;
+
+import net.minecraft.entity.ai.goal.Goal;
+import net.minecraft.entity.ai.goal.GoalSelector;
+import net.minecraft.entity.ai.goal.PrioritizedGoal;
+
+public class PredicatedGoal extends Goal {
+    private final Goal goal;
+    private final BooleanSupplier predicate;
+
+    public static void applyToAll(GoalSelector goals, BooleanSupplier predicate) {
+        Set<PrioritizedGoal> existingTasks = new HashSet<>(goals.getGoals());
+        goals.clear(g -> !(g instanceof PredicatedGoal));
+        existingTasks.forEach(goal -> {
+            if (!(goal.getGoal() instanceof PredicatedGoal)) {
+                goals.add(goal.getPriority(), new PredicatedGoal(goal.getGoal(), predicate));
+            }
+        });
+    }
+
+    public PredicatedGoal(Goal goal, BooleanSupplier predicate) {
+        this.goal = goal;
+        this.predicate = predicate;
+    }
+
+    @Override
+    public boolean canStart() {
+        return predicate.getAsBoolean() && goal.canStart();
+    }
+
+    @Override
+    public boolean shouldContinue() {
+        return predicate.getAsBoolean() && goal.shouldContinue();
+    }
+
+    @Override
+    public boolean canStop() {
+        return goal.canStop();
+    }
+
+    @Override
+    public void start() {
+        if (predicate.getAsBoolean()) {
+            goal.start();
+        }
+    }
+
+    @Override
+    public void stop() {
+        goal.stop();
+    }
+
+    @Override
+    public boolean shouldRunEveryTick() {
+        return goal.shouldRunEveryTick();
+    }
+
+    @Override
+    public void tick() {
+        goal.tick();
+    }
+
+    @Override
+    public void setControls(EnumSet<Goal.Control> controls) {
+        goal.setControls(controls);
+    }
+
+    @Override
+    public EnumSet<Goal.Control> getControls() {
+        return goal.getControls();
+    }
+
+    @Override
+    public boolean equals(@Nullable Object o) {
+        if (this == o) {
+            return true;
+        } else {
+            return o != null && getClass() == o.getClass() ? goal.equals(((PredicatedGoal)o).goal) : false;
+        }
+    }
+
+    @Override
+    public int hashCode() {
+        return goal.hashCode();
+    }
+}

src/main/java/com/minelittlepony/unicopia/item/BellItem.java

@@ -120,6 +120,11 @@ public class BellItem extends Item {
         user.setTarget(null);
         user.playSound(USounds.ITEM_GROGAR_BELL_STOP_USING, 0.2F, 0.3F);
         if (target instanceof Creature creature && (completed || target.asEntity().getHealth() < (target.asEntity().getMaxHealth() * 0.5F) + 1)) {
+            ItemStack handStack = creature.asEntity().getStackInHand(Hand.MAIN_HAND);
+            if (!handStack.isEmpty()) {
+                creature.asEntity().setStackInHand(Hand.MAIN_HAND, ItemStack.EMPTY);
+                creature.asEntity().dropStack(handStack);
+            }
             creature.setDiscorded(true);
         }
     }

src/main/java/com/minelittlepony/unicopia/mixin/MixinBrain.java

@@ -0,0 +1,21 @@
+package com.minelittlepony.unicopia.mixin;
+
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.*;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+
+import com.minelittlepony.unicopia.entity.Creature;
+import com.minelittlepony.unicopia.entity.Living;
+import net.minecraft.entity.LivingEntity;
+import net.minecraft.entity.ai.brain.Brain;
+import net.minecraft.server.world.ServerWorld;
+
+@Mixin(Brain.class)
+abstract class MixinBrain<E extends LivingEntity> {
+    @Inject(method = "tick", at = @At("HEAD"), cancellable = true)
+    private void onTick(ServerWorld world, E entity, CallbackInfo info) {
+        if (Living.living(entity) instanceof Creature c && c.isDiscorded()) {
+            info.cancel();
+        }
+    }
+}

src/main/resources/unicopia.mixin.json

@@ -14,6 +14,7 @@
     "MixinBlockEntityType",
     "MixinBlockItem",
     "MixinBoatEntity",
+    "MixinBrain",
     "MixinChunkBlockLightProvider",
     "MixinComponentHolder",
     "MutableBlockLightStorage",