Move target selection for shields into a separate class

src/main/java/com/minelittlepony/unicopia/ability/magic/spell/effect/AttractiveSpell.java

@@ -1,7 +1,5 @@
 package com.minelittlepony.unicopia.ability.magic.spell.effect;
 
-import java.util.List;
-
 import com.minelittlepony.unicopia.ability.magic.Caster;
 import com.minelittlepony.unicopia.ability.magic.spell.ProjectileSpell;
 import com.minelittlepony.unicopia.ability.magic.spell.trait.SpellTraits;
@@ -9,7 +7,6 @@ import com.minelittlepony.unicopia.ability.magic.spell.trait.Trait;
 import com.minelittlepony.unicopia.entity.player.Pony;
 import com.minelittlepony.unicopia.particle.MagicParticleEffect;
 import com.minelittlepony.unicopia.util.MagicalDamageSource;
-import com.minelittlepony.unicopia.util.VecHelper;
 import com.minelittlepony.unicopia.util.shape.Sphere;
 
 import net.minecraft.entity.Entity;
@@ -39,13 +36,8 @@ public class AttractiveSpell extends ShieldSpell implements ProjectileSpell {
     }
 
     @Override
-    protected List<Entity> getTargets(Caster<?> source, double radius) {
+    protected boolean isValidTarget(Entity entity) {
-
+        return getTraits().get(Trait.FOCUS) > 10 ? entity instanceof ItemEntity : super.isValidTarget(entity);
-        if (getTraits().get(Trait.FOCUS) > 10) {
-            return VecHelper.findInRange(source.getEntity(), source.getWorld(), source.getOriginVector(), radius, i -> i instanceof ItemEntity);
-        }
-
-        return super.getTargets(source, radius);
     }
 
     @Override

src/main/java/com/minelittlepony/unicopia/ability/magic/spell/effect/ShieldSpell.java

@@ -1,21 +1,12 @@
 package com.minelittlepony.unicopia.ability.magic.spell.effect;
 
-import java.util.List;
-import java.util.Map;
-import java.util.TreeMap;
-import java.util.UUID;
-import java.util.stream.Collectors;
-
 import com.minelittlepony.unicopia.Affinity;
-import com.minelittlepony.unicopia.EquinePredicates;
 import com.minelittlepony.unicopia.Unicopia;
 import com.minelittlepony.unicopia.ability.magic.Caster;
-import com.minelittlepony.unicopia.ability.magic.SpellPredicate;
 import com.minelittlepony.unicopia.ability.magic.spell.Situation;
 import com.minelittlepony.unicopia.ability.magic.spell.trait.SpellTraits;
 import com.minelittlepony.unicopia.ability.magic.spell.trait.Trait;
 import com.minelittlepony.unicopia.entity.player.Pony;
-import com.minelittlepony.unicopia.item.FriendshipBraceletItem;
 import com.minelittlepony.unicopia.item.enchantment.UEnchantments;
 import com.minelittlepony.unicopia.particle.MagicParticleEffect;
 import com.minelittlepony.unicopia.particle.ParticleHandle;
@@ -47,7 +38,7 @@ public class ShieldSpell extends AbstractSpell {
 
     protected final ParticleHandle particlEffect = new ParticleHandle();
 
-    private final Map<UUID, Target> targets = new TreeMap<>();
+    private final TargetSelecter targetSelecter = new TargetSelecter(this);
 
     protected ShieldSpell(SpellType<?> type, SpellTraits traits) {
         super(type, traits);
@@ -120,22 +111,6 @@ public class ShieldSpell extends AbstractSpell {
         return (min + (source.getLevel().get() * 2)) / multiplier;
     }
 
-    protected List<Entity> getTargets(Caster<?> source, double radius) {
-
-        Entity owner = source.getMaster();
-
-        boolean ownerIsValid = isFriendlyTogether(source) && (EquinePredicates.PLAYER_UNICORN.test(owner) && owner.isSneaking());
-
-        return source.findAllEntitiesInRange(radius)
-            .filter(entity -> {
-                return !FriendshipBraceletItem.isComrade(source, entity)
-                        && !SpellPredicate.IS_SHIELD_LIKE.isOn(entity)
-                        && isValidTarget(entity)
-                        && !(ownerIsValid && (Pony.equal(entity, owner) || owner.isConnectedThroughVehicle(entity)));
-            })
-            .collect(Collectors.toList());
-    }
-
     protected boolean isValidTarget(Entity entity) {
         return (entity instanceof LivingEntity
                 || entity instanceof TntEntity
@@ -152,21 +127,15 @@ public class ShieldSpell extends AbstractSpell {
 
         Vec3d origin = source.getOriginVector();
 
-        this.targets.values().removeIf(Target::tick);
+        targetSelecter.getEntities(source, radius, this::isValidTarget).forEach(i -> {
-
-        List<Entity> targets = getTargets(source, radius);
-        targets.forEach(i -> {
             try {
-                this.targets.computeIfAbsent(i.getUuid(), Target::new);
+                applyRadialEffect(source, i, i.getPos().distanceTo(origin), radius);
-                double dist = i.getPos().distanceTo(origin);
-
-                applyRadialEffect(source, i, dist, radius);
             } catch (Throwable e) {
                 Unicopia.LOGGER.error("Error updating shield effect", e);
             }
         });
 
-        return this.targets.values().stream().filter(Target::canHurt).count();
+        return targetSelecter.getTotalDamaged();
     }
 
     protected void applyRadialEffect(Caster<?> source, Entity target, double distance, double radius) {
@@ -230,18 +199,4 @@ public class ShieldSpell extends AbstractSpell {
         return force;
     }
 
-    class Target {
-
-        int cooldown = 20;
-
-        Target(UUID id) { }
-
-        boolean tick() {
-            return --cooldown < 0;
-        }
-
-        boolean canHurt() {
-            return cooldown == 20;
-        }
-    }
 }

src/main/java/com/minelittlepony/unicopia/ability/magic/spell/effect/TargetSelecter.java

@@ -0,0 +1,64 @@
+package com.minelittlepony.unicopia.ability.magic.spell.effect;
+
+import java.util.Map;
+import java.util.TreeMap;
+import java.util.UUID;
+import java.util.function.Predicate;
+import java.util.stream.Stream;
+
+import com.minelittlepony.unicopia.EquinePredicates;
+import com.minelittlepony.unicopia.ability.magic.Caster;
+import com.minelittlepony.unicopia.ability.magic.SpellPredicate;
+import com.minelittlepony.unicopia.ability.magic.spell.Spell;
+import com.minelittlepony.unicopia.entity.player.Pony;
+import com.minelittlepony.unicopia.item.FriendshipBraceletItem;
+import net.minecraft.entity.Entity;
+
+public class TargetSelecter {
+
+    private final Map<UUID, Target> targets = new TreeMap<>();
+
+    private final Spell spell;
+
+    public TargetSelecter(Spell spell) {
+        this.spell = spell;
+    }
+
+    public Stream<Entity> getEntities(Caster<?> source, double radius, Predicate<Entity> filter) {
+        targets.values().removeIf(Target::tick);
+
+        Entity owner = source.getMaster();
+
+        boolean ownerIsValid = spell.isFriendlyTogether(source) && (EquinePredicates.PLAYER_UNICORN.test(owner) && owner.isSneaking());
+
+        return source.findAllEntitiesInRange(radius)
+            .filter(entity -> {
+                return !FriendshipBraceletItem.isComrade(source, entity)
+                        && !SpellPredicate.IS_SHIELD_LIKE.isOn(entity)
+                        && !(ownerIsValid && (Pony.equal(entity, owner) || owner.isConnectedThroughVehicle(entity)));
+            })
+            .filter(filter)
+            .map(i -> {
+                targets.computeIfAbsent(i.getUuid(), Target::new);
+                return i;
+            });
+    }
+
+    public long getTotalDamaged() {
+        return targets.values().stream().filter(Target::canHurt).count();
+    }
+
+    static final class Target {
+        private int cooldown = 20;
+
+        Target(UUID id) { }
+
+        boolean tick() {
+            return --cooldown < 0;
+        }
+
+        boolean canHurt() {
+            return cooldown == 20;
+        }
+    }
+}