A bit of cleanup

src/main/java/com/minelittlepony/unicopia/entity/behaviour/BeeBehaviour.java

@@ -1,7 +1,8 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
 import com.minelittlepony.unicopia.ability.magic.Caster;
-import com.minelittlepony.unicopia.ability.magic.Spell;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
+
 import net.minecraft.client.MinecraftClient;
 import net.minecraft.client.sound.AggressiveBeeSoundInstance;
 import net.minecraft.client.sound.PassiveBeeSoundInstance;
@@ -20,7 +21,7 @@ public class BeeBehaviour extends EntityBehaviour<BeeEntity> {
     }
 
     @Override
-    public void update(Caster<?> source, BeeEntity entity, Spell spell) {
+    public void update(Caster<?> source, BeeEntity entity, DisguiseSpell spell) {
 
         if (source.getOwner().isSneaking()) {
             entity.setAngerTime(10);

src/main/java/com/minelittlepony/unicopia/entity/behaviour/ChickenBehaviour.java

@@ -1,7 +1,7 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
 import com.minelittlepony.unicopia.ability.magic.Caster;
-import com.minelittlepony.unicopia.ability.magic.Spell;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
 import com.minelittlepony.unicopia.entity.player.Pony;
 
 import net.minecraft.entity.Entity;
@@ -21,7 +21,7 @@ public class ChickenBehaviour extends EntityBehaviour<ChickenEntity> {
     }
 
     @Override
-    public void update(Caster<?> source, ChickenEntity entity, Spell spell) {
+    public void update(Caster<?> source, ChickenEntity entity, DisguiseSpell spell) {
         entity.eggLayTime = Integer.MAX_VALUE;
 
         if (source instanceof Pony) {

src/main/java/com/minelittlepony/unicopia/entity/behaviour/CreeperBehaviour.java

@@ -1,13 +1,13 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
 import com.minelittlepony.unicopia.ability.magic.Caster;
-import com.minelittlepony.unicopia.ability.magic.Spell;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
 
 import net.minecraft.entity.mob.CreeperEntity;
 
 public class CreeperBehaviour extends EntityBehaviour<CreeperEntity> {
     @Override
-    public void update(Caster<?> source, CreeperEntity entity, Spell spell) {
+    public void update(Caster<?> source, CreeperEntity entity, DisguiseSpell spell) {
         if (isSneakingOnGround(source)) {
             entity.setFuseSpeed(1);
         } else {

src/main/java/com/minelittlepony/unicopia/entity/behaviour/EndermanBehaviour.java

@@ -1,7 +1,7 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
 import com.minelittlepony.unicopia.ability.magic.Caster;
-import com.minelittlepony.unicopia.ability.magic.Spell;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
 
 import net.minecraft.entity.mob.EndermanEntity;
 import net.minecraft.item.BlockItem;
@@ -10,7 +10,7 @@ import net.minecraft.util.Hand;
 
 public class EndermanBehaviour extends EntityBehaviour<EndermanEntity> {
     @Override
-    public void update(Caster<?> source, EndermanEntity entity, Spell spell) {
+    public void update(Caster<?> source, EndermanEntity entity, DisguiseSpell spell) {
         if (source.getOwner().isSneaking() || source.getOwner().isSprinting()) {
             entity.setTarget(entity);
         } else {

src/main/java/com/minelittlepony/unicopia/entity/behaviour/EntityBehaviour.java

@@ -6,8 +6,9 @@ import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
 import com.minelittlepony.unicopia.ability.magic.Caster;
-import com.minelittlepony.unicopia.ability.magic.Spell;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
 import com.minelittlepony.unicopia.entity.ItemWielder;
+import com.minelittlepony.unicopia.entity.player.Pony;
 import com.minelittlepony.unicopia.util.Registries;
 
 import net.minecraft.entity.Entity;
@@ -35,7 +36,13 @@ public class EntityBehaviour<T extends Entity> {
     private static final EntityBehaviour<Entity> DEFAULT = new EntityBehaviour<>();
     private static final Registry<EntityBehaviour<?>> REGISTRY = Registries.createSimple(new Identifier("unicopia", "entity_behaviour"));
 
-    public void update(Caster<?> source, T entity, Spell spell) {
+    public void update(Caster<?> source, T entity, DisguiseSpell spell) {
+        if (source instanceof Pony) {
+            update((Pony)source, entity, spell);
+        }
+    }
+
+    protected void update(Pony pony, T entity, DisguiseSpell spell) {
 
     }
 
@@ -224,7 +231,7 @@ public class EntityBehaviour<T extends Entity> {
         register(FallingBlockBehaviour::new, EntityType.FALLING_BLOCK);
         register(MobBehaviour::new, EntityType.RAVAGER, EntityType.IRON_GOLEM);
         register(RabbitBehaviour::new, EntityType.RABBIT);
-        register(VillagerBehaviour::new, EntityType.VILLAGER, EntityType.WANDERING_TRADER);
+        register(TraderBehaviour::new, EntityType.VILLAGER, EntityType.WANDERING_TRADER);
         register(SheepBehaviour::new, EntityType.SHEEP);
         register(BeeBehaviour::new, EntityType.BEE);
         register(EndermanBehaviour::new, EntityType.ENDERMAN);

src/main/java/com/minelittlepony/unicopia/entity/behaviour/FallingBlockBehaviour.java

@@ -4,7 +4,6 @@ import java.util.List;
 import java.util.Optional;
 
 import com.minelittlepony.unicopia.ability.magic.Caster;
-import com.minelittlepony.unicopia.ability.magic.Spell;
 import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
 import com.minelittlepony.unicopia.entity.player.Pony;
 import com.minelittlepony.unicopia.mixin.MixinBlockEntity;
@@ -68,7 +67,7 @@ public class FallingBlockBehaviour extends EntityBehaviour<FallingBlockEntity> {
     }
 
     @Override
-    public void update(Caster<?> source, FallingBlockEntity entity, Spell spell) {
+    public void update(Caster<?> source, FallingBlockEntity entity, DisguiseSpell spell) {
 
         BlockState state = entity.getBlockState();
         if (state.contains(Properties.WATERLOGGED)) {
@@ -76,12 +75,12 @@ public class FallingBlockBehaviour extends EntityBehaviour<FallingBlockEntity> {
 
             if (state.get(Properties.WATERLOGGED) != logged) {
                 entity = new FallingBlockEntity(entity.world, entity.getX(), entity.getY(), entity.getZ(), state.with(Properties.WATERLOGGED, logged));
-                ((DisguiseSpell)spell).getDisguise().setAppearance(entity);
+                spell.getDisguise().setAppearance(entity);
                 return;
             }
         }
 
-        Disguise disguise = ((DisguiseSpell)spell).getDisguise();
+        Disguise disguise = spell.getDisguise();
         List<Entity> attachments = disguise.getAttachments();
         if (attachments.size() > 0) {
             copyBaseAttributes(source.getOwner(), attachments.get(0), UP);

src/main/java/com/minelittlepony/unicopia/entity/behaviour/MinecartBehaviour.java

@@ -1,7 +1,7 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
 import com.minelittlepony.unicopia.ability.magic.Caster;
-import com.minelittlepony.unicopia.ability.magic.Spell;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
 
 import net.minecraft.client.MinecraftClient;
 import net.minecraft.client.sound.MovingMinecartSoundInstance;
@@ -20,7 +20,7 @@ public class MinecartBehaviour extends EntityBehaviour<AbstractMinecartEntity> {
     }
 
     @Override
-    public void update(Caster<?> source, AbstractMinecartEntity entity, Spell spell) {
+    public void update(Caster<?> source, AbstractMinecartEntity entity, DisguiseSpell spell) {
         entity.yaw -= 90;
         entity.prevYaw -= 90;
 

src/main/java/com/minelittlepony/unicopia/entity/behaviour/MobBehaviour.java

@@ -1,7 +1,6 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
-import com.minelittlepony.unicopia.ability.magic.Caster;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
-import com.minelittlepony.unicopia.ability.magic.Spell;
 import com.minelittlepony.unicopia.entity.player.Pony;
 import com.minelittlepony.unicopia.util.RayTraceHelper;
 
@@ -10,14 +9,10 @@ import net.minecraft.entity.mob.MobEntity;
 
 public class MobBehaviour extends EntityBehaviour<MobEntity> {
     @Override
-    public void update(Caster<?> source, MobEntity entity, Spell spell) {
+    public void update(Pony player, MobEntity entity, DisguiseSpell spell) {
 
-        if (source instanceof Pony) {
+        if (player.sneakingChanged() && isSneakingOnGround(player)) {
-            Pony player = (Pony)source;
+            entity.tryAttack(RayTraceHelper.findEntity(player.getEntity(), 6, 1, e -> e instanceof LivingEntity).orElse(entity));
-
-            if (player.sneakingChanged() && isSneakingOnGround(source)) {
-                entity.tryAttack(RayTraceHelper.findEntity(source.getEntity(), 6, 1, e -> e instanceof LivingEntity).orElse(entity));
-            }
         }
     }
 }

src/main/java/com/minelittlepony/unicopia/entity/behaviour/RabbitBehaviour.java

@@ -1,7 +1,6 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
-import com.minelittlepony.unicopia.ability.magic.Caster;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
-import com.minelittlepony.unicopia.ability.magic.Spell;
 import com.minelittlepony.unicopia.entity.player.Pony;
 
 import net.minecraft.entity.Entity;
@@ -9,10 +8,7 @@ import net.minecraft.entity.passive.RabbitEntity;
 
 public class RabbitBehaviour extends EntityBehaviour<RabbitEntity> {
     @Override
-    public void update(Caster<?> source, RabbitEntity entity, Spell spell) {
+    public void update(Pony player, RabbitEntity entity, DisguiseSpell spell) {
-
-        if (source instanceof Pony) {
-            Pony player = (Pony)source;
 
         if (player.getEntity().isOnGround()) {
             if (Entity.squaredHorizontalLength(player.getEntity().getVelocity()) > 0.01) {
@@ -24,4 +20,3 @@ public class RabbitBehaviour extends EntityBehaviour<RabbitEntity> {
         }
     }
 }
-}

src/main/java/com/minelittlepony/unicopia/entity/behaviour/RangedAttackBehaviour.java

@@ -1,8 +1,7 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
 import java.util.function.BiFunction;
-import com.minelittlepony.unicopia.ability.magic.Caster;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
-import com.minelittlepony.unicopia.ability.magic.Spell;
 import com.minelittlepony.unicopia.entity.player.Pony;
 
 import net.minecraft.entity.Entity;
@@ -24,19 +23,16 @@ public class RangedAttackBehaviour<T extends Entity & RangedAttackMob> extends E
     }
 
     @Override
-    public void update(Caster<?> source, T entity, Spell spell) {
+    public void update(Pony player, T entity, DisguiseSpell spell) {
-
-        if (source instanceof Pony) {
-            Pony player = (Pony)source;
 
         if (player.sneakingChanged() && isSneakingOnGround(player)) {
 
             ProjectileEntity spit = projectileSupplier.apply(entity.world, entity);
 
-                Vec3d rot = source.getEntity().getRotationVec(1);
+            Vec3d rot = player.getEntity().getRotationVec(1);
 
             spit.setVelocity(rot.getX(), rot.getY(), rot.getZ(), 1.5F, 3);
-                spit.setOwner(source.getOwner());
+            spit.setOwner(player.getOwner());
 
             if (!entity.isSilent()) {
                 entity.world.playSound(null, entity.getX(), entity.getY(), entity.getZ(),
@@ -48,4 +44,3 @@ public class RangedAttackBehaviour<T extends Entity & RangedAttackMob> extends E
         }
     }
 }
-}

src/main/java/com/minelittlepony/unicopia/entity/behaviour/SheepBehaviour.java

@@ -2,8 +2,7 @@ package com.minelittlepony.unicopia.entity.behaviour;
 
 import java.util.Random;
 
-import com.minelittlepony.unicopia.ability.magic.Caster;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
-import com.minelittlepony.unicopia.ability.magic.Spell;
 import com.minelittlepony.unicopia.entity.player.Pony;
 import com.minelittlepony.unicopia.mixin.MixinSheepEntity;
 import com.minelittlepony.unicopia.util.WorldEvent;
@@ -20,9 +19,7 @@ import net.minecraft.util.math.BlockPos;
 
 public class SheepBehaviour extends EntityBehaviour<SheepEntity> {
     @Override
-    public void update(Caster<?> source, SheepEntity entity, Spell spell) {
+    public void update(Pony player, SheepEntity entity, DisguiseSpell spell) {
-        if (source instanceof Pony) {
-            Pony player = (Pony)source;
 
         if (player.sneakingChanged()) {
 
@@ -74,4 +71,3 @@ public class SheepBehaviour extends EntityBehaviour<SheepEntity> {
         }
     }
 }
-}

src/main/java/com/minelittlepony/unicopia/entity/behaviour/ShulkerBehaviour.java

@@ -1,7 +1,7 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
 import com.minelittlepony.unicopia.ability.magic.Caster;
-import com.minelittlepony.unicopia.ability.magic.Spell;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
 import com.minelittlepony.unicopia.entity.player.Pony;
 import com.minelittlepony.unicopia.mixin.MixinShulkerEntity;
 
@@ -16,38 +16,14 @@ import net.minecraft.util.math.Vec3d;
 
 public class ShulkerBehaviour extends EntityBehaviour<ShulkerEntity> {
     @Override
-    public void update(Caster<?> source, ShulkerEntity shulker, Spell spell) {
+    public void update(Caster<?> source, ShulkerEntity shulker, DisguiseSpell spell) {
         shulker.yaw = 0;
         shulker.prevBodyYaw = 0;
         shulker.bodyYaw = 0;
 
         shulker.setAttachedBlock(null);
 
-        if (source instanceof Pony) {
+        super.update(source, shulker, spell);
-            Pony player = (Pony)source;
-
-            float peekAmount = 30;
-
-            double speed = !source.getEntity().isSneaking() ? 0.29 : 0;
-            speed += Math.sqrt(Entity.squaredHorizontalLength(source.getEntity().getVelocity())) * 2;
-
-            peekAmount = (float)MathHelper.clamp(speed, 0, 1);
-            peekAmount = ((Pony)source).getInterpolator().interpolate("peek", peekAmount, 5);
-
-            MixinShulkerEntity mx = (MixinShulkerEntity)shulker;
-
-            mx.setPrevOpenProgress(mx.getOpenProgress());
-            mx.setOpenProgress(peekAmount);
-
-            if (player.sneakingChanged()) {
-                shulker.setPeekAmount((int)(peekAmount / 0.01F));
-            } else if (peekAmount > 0.2 && shulker.getPeekAmount() == 0) {
-                if (shulker.isAlive() && shulker.world.random.nextInt(1000) < shulker.ambientSoundChance++) {
-                    shulker.ambientSoundChance = -shulker.getMinAmbientSoundDelay();
-                    shulker.playSound(SoundEvents.ENTITY_SHULKER_AMBIENT, 1, 1);
-                 }
-            }
-        }
 
         Direction attachmentFace = shulker.getAttachedFace();
         BlockPos pos = shulker.getBlockPos().offset(attachmentFace);
@@ -63,4 +39,29 @@ public class ShulkerBehaviour extends EntityBehaviour<ShulkerEntity> {
             }
         }
     }
+
+    @Override
+    protected void update(Pony player, ShulkerEntity shulker, DisguiseSpell spell) {
+        float peekAmount = 30;
+
+        double speed = !player.getEntity().isSneaking() ? 0.29 : 0;
+        speed += Math.sqrt(Entity.squaredHorizontalLength(player.getEntity().getVelocity())) * 2;
+
+        peekAmount = (float)MathHelper.clamp(speed, 0, 1);
+        peekAmount = player.getInterpolator().interpolate("peek", peekAmount, 5);
+
+        MixinShulkerEntity mx = (MixinShulkerEntity)shulker;
+
+        mx.setPrevOpenProgress(mx.getOpenProgress());
+        mx.setOpenProgress(peekAmount);
+
+        if (player.sneakingChanged()) {
+            shulker.setPeekAmount((int)(peekAmount / 0.01F));
+        } else if (peekAmount > 0.2 && shulker.getPeekAmount() == 0) {
+            if (shulker.isAlive() && shulker.world.random.nextInt(1000) < shulker.ambientSoundChance++) {
+                shulker.ambientSoundChance = -shulker.getMinAmbientSoundDelay();
+                shulker.playSound(SoundEvents.ENTITY_SHULKER_AMBIENT, 1, 1);
+             }
+        }
+    }
 }

src/main/java/com/minelittlepony/unicopia/entity/behaviour/SilverfishBehaviour.java

@@ -1,7 +1,6 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
-import com.minelittlepony.unicopia.ability.magic.Caster;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
-import com.minelittlepony.unicopia.ability.magic.Spell;
 import com.minelittlepony.unicopia.block.state.StateMaps;
 import com.minelittlepony.unicopia.entity.player.Pony;
 import com.minelittlepony.unicopia.util.WorldEvent;
@@ -12,11 +11,8 @@ import net.minecraft.util.math.BlockPos;
 
 public class SilverfishBehaviour extends EntityBehaviour<SilverfishEntity> {
     @Override
-    public void update(Caster<?> source, SilverfishEntity entity, Spell spell) {
+    public void update(Pony player, SilverfishEntity entity, DisguiseSpell spell) {
-        if (source instanceof Pony && !source.isClient()) {
+        if (!player.isClient() && player.sneakingChanged() && player.getOwner().isSneaking()) {
-            Pony player = (Pony)source;
-
-            if (player.sneakingChanged() && player.getOwner().isSneaking()) {
             BlockPos pos = entity.getBlockPos().down();
             BlockState state = entity.world.getBlockState(pos);
 
@@ -28,4 +24,3 @@ public class SilverfishBehaviour extends EntityBehaviour<SilverfishEntity> {
         }
     }
 }
-}

src/main/java/com/minelittlepony/unicopia/entity/behaviour/SpellcastingIllagerBehaviour.java

@@ -1,18 +1,13 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
-import com.minelittlepony.unicopia.ability.magic.Caster;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
-import com.minelittlepony.unicopia.ability.magic.Spell;
 import com.minelittlepony.unicopia.entity.player.Pony;
 
 import net.minecraft.entity.mob.SpellcastingIllagerEntity;
 
 public class SpellcastingIllagerBehaviour extends EntityBehaviour<SpellcastingIllagerEntity> {
     @Override
-    public void update(Caster<?> source, SpellcastingIllagerEntity entity, Spell s) {
+    public void update(Pony player, SpellcastingIllagerEntity entity, DisguiseSpell s) {
-
-        if (source instanceof Pony) {
-            Pony player = (Pony)source;
-
         if (player.sneakingChanged()) {
             if (player.getOwner().isSneaking()) {
                 SpellcastingIllagerEntity.Spell[] spells = SpellcastingIllagerEntity.Spell.values();
@@ -27,4 +22,3 @@ public class SpellcastingIllagerBehaviour extends EntityBehaviour<SpellcastingIl
         }
     }
 }
-}

src/main/java/com/minelittlepony/unicopia/entity/behaviour/TraderBehaviour.java

@@ -0,0 +1,20 @@
+package com.minelittlepony.unicopia.entity.behaviour;
+
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
+import com.minelittlepony.unicopia.entity.player.Pony;
+
+import net.minecraft.entity.passive.AbstractTraderEntity;
+import net.minecraft.sound.SoundEvents;
+
+public class TraderBehaviour extends EntityBehaviour<AbstractTraderEntity> {
+    @Override
+    public void update(Pony pony, AbstractTraderEntity entity, DisguiseSpell spell) {
+        if (pony.sneakingChanged() && pony.getOwner().isSneaking()) {
+            entity.setHeadRollingTimeLeft(40);
+
+            if (!entity.world.isClient()) {
+               entity.playSound(SoundEvents.ENTITY_VILLAGER_NO, 1, 1);
+            }
+        }
+    }
+}

src/main/java/com/minelittlepony/unicopia/entity/behaviour/VillagerBehaviour.java

@@ -1,26 +0,0 @@
-package com.minelittlepony.unicopia.entity.behaviour;
-
-import com.minelittlepony.unicopia.ability.magic.Caster;
-import com.minelittlepony.unicopia.ability.magic.Spell;
-import com.minelittlepony.unicopia.entity.player.Pony;
-
-import net.minecraft.entity.passive.AbstractTraderEntity;
-import net.minecraft.sound.SoundEvents;
-
-public class VillagerBehaviour extends EntityBehaviour<AbstractTraderEntity> {
-    @Override
-    public void update(Caster<?> source, AbstractTraderEntity entity, Spell spell) {
-
-        if (source instanceof Pony) {
-            Pony player = (Pony)source;
-
-            if (player.sneakingChanged() && player.getOwner().isSneaking()) {
-                entity.setHeadRollingTimeLeft(40);
-
-                if (!entity.world.isClient()) {
-                   entity.playSound(SoundEvents.ENTITY_VILLAGER_NO, 1, 1);
-                }
-            }
-        }
-    }
-}