Link the spell to the entity it spawns

src/main/java/com/minelittlepony/unicopia/ability/magic/spell/AbstractPlacedSpell.java

@@ -6,9 +6,11 @@ import com.minelittlepony.unicopia.UEntities;
 import com.minelittlepony.unicopia.ability.magic.Attached;
 import com.minelittlepony.unicopia.ability.magic.Caster;
 import com.minelittlepony.unicopia.entity.CastSpellEntity;
+import com.minelittlepony.unicopia.entity.EntityReference;
 import com.minelittlepony.unicopia.particle.OrientedBillboardParticleEffect;
 import com.minelittlepony.unicopia.particle.ParticleHandle;
 import com.minelittlepony.unicopia.particle.UParticles;
+
 import net.minecraft.nbt.CompoundTag;
 import net.minecraft.util.Identifier;
 import net.minecraft.util.math.Vec3d;
@@ -20,6 +22,8 @@ public abstract class AbstractPlacedSpell extends AbstractSpell implements Attac
 
     private final ParticleHandle particlEffect = new ParticleHandle();
 
+    private final EntityReference<CastSpellEntity> castEntity = new EntityReference<>();
+
     protected AbstractPlacedSpell(SpellType<?> type) {
         super(type);
     }
@@ -39,13 +43,15 @@ public abstract class AbstractPlacedSpell extends AbstractSpell implements Attac
                 dimension = source.getWorld().getRegistryKey().getValue();
                 setDirty(true);
 
-                if (!source.isClient()) {
+                if (!source.isClient() && !castEntity.isPresent(source.getWorld())) {
                     CastSpellEntity entity = UEntities.CAST_SPELL.create(source.getWorld());
                     Vec3d pos = source.getOriginVector();
                     entity.updatePositionAndAngles(pos.x, pos.y, pos.z, 0, 0);
                     entity.setSpell(this);
                     entity.setMaster(source.getMaster());
                     entity.world.spawnEntity(entity);
+
+                    castEntity.set(entity);
                 }
             }
 
@@ -75,6 +81,7 @@ public abstract class AbstractPlacedSpell extends AbstractSpell implements Attac
         if (dimension != null) {
             compound.putString("dimension", dimension.toString());
         }
+        compound.put("owner", castEntity.toNBT());
     }
 
     @Override
@@ -84,5 +91,8 @@ public abstract class AbstractPlacedSpell extends AbstractSpell implements Attac
         if (compound.contains("dimension")) {
             dimension = new Identifier(compound.getString("dimension"));
         }
+        if (compound.contains("castEntity")) {
+            castEntity.fromNBT(compound.getCompound("castEntity"));
+        }
     }
 }

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

@@ -21,7 +21,6 @@ import net.minecraft.entity.data.TrackedData;
 import net.minecraft.entity.data.TrackedDataHandlerRegistry;
 import net.minecraft.nbt.CompoundTag;
 import net.minecraft.network.Packet;
-import net.minecraft.server.world.ServerWorld;
 import net.minecraft.world.World;
 
 public class CastSpellEntity extends Entity implements Caster<LivingEntity> {
@@ -36,8 +35,7 @@ public class CastSpellEntity extends Entity implements Caster<LivingEntity> {
 
     private final EntityPhysics<CastSpellEntity> physics = new EntityPhysics<>(this, GRAVITY);
 
-    private UUID ownerUuid;
+    private final EntityReference<LivingEntity> owner = new EntityReference<>();
-    private int ownerEntityId;
 
     public CastSpellEntity(EntityType<?> type, World world) {
         super(type, world);
@@ -80,10 +78,7 @@ public class CastSpellEntity extends Entity implements Caster<LivingEntity> {
 
     @Override
     public void setMaster(LivingEntity owner) {
-        if (owner != null) {
+        this.owner.set(owner);
-            ownerUuid = owner.getUuid();
-            ownerEntityId = owner.getEntityId();
-         }
     }
 
     @Override
@@ -98,15 +93,7 @@ public class CastSpellEntity extends Entity implements Caster<LivingEntity> {
 
     @Override
     public LivingEntity getMaster() {
-        if (ownerUuid != null && world instanceof ServerWorld) {
+        return owner.get(world);
-            return (LivingEntity)((ServerWorld)world).getEntity(ownerUuid);
-        }
-
-        if (ownerEntityId != 0) {
-            return (LivingEntity)world.getEntityById(ownerEntityId);
-        }
-
-        return null;
     }
 
     @Override
@@ -131,18 +118,14 @@ public class CastSpellEntity extends Entity implements Caster<LivingEntity> {
 
     @Override
     protected void writeCustomDataToTag(CompoundTag tag) {
-        if (ownerUuid != null) {
+        tag.put("owner", owner.toNBT());
-            tag.putUuid("Owner", ownerUuid);
-        }
-        tag.putInt("OwnerId", ownerEntityId);
     }
 
     @Override
     protected void readCustomDataFromTag(CompoundTag tag) {
-        if (tag.containsUuid("Owner")) {
+        if (tag.contains("owner")) {
-            ownerUuid = tag.getUuid("Owner");
+            owner.fromNBT(tag.getCompound("owner"));
         }
-        ownerEntityId = tag.getInt("OwnerId");
     }
 
 

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

@@ -0,0 +1,60 @@
+package com.minelittlepony.unicopia.entity;
+
+import java.util.UUID;
+
+import javax.annotation.Nullable;
+
+import com.minelittlepony.unicopia.util.NbtSerialisable;
+
+import net.minecraft.entity.Entity;
+import net.minecraft.nbt.CompoundTag;
+import net.minecraft.server.world.ServerWorld;
+import net.minecraft.world.World;
+
+public class EntityReference<T extends Entity> implements NbtSerialisable {
+
+    private UUID uuid;
+    private int clientId;
+
+    public void set(@Nullable T entity) {
+        if (entity != null) {
+            uuid = entity.getUuid();
+            clientId = entity.getEntityId();
+        }
+    }
+
+    public boolean isPresent(World world) {
+        T entity = get(world);
+        return entity != null && !entity.removed;
+    }
+
+    @SuppressWarnings("unchecked")
+    @Nullable
+    public T get(World world) {
+        if (uuid != null && world instanceof ServerWorld) {
+            return (T)((ServerWorld)world).getEntity(uuid);
+        }
+
+        if (clientId != 0) {
+            return (T)world.getEntityById(clientId);
+        }
+
+        return null;
+    }
+
+    @Override
+    public void toNBT(CompoundTag tag) {
+        if (uuid != null) {
+            tag.putUuid("uuid", uuid);
+        }
+        tag.putInt("clientId", clientId);
+    }
+
+    @Override
+    public void fromNBT(CompoundTag tag) {
+        if (tag.containsUuid("uuid")) {
+            uuid = tag.getUuid("uuid");
+        }
+        clientId = tag.getInt("clientId");
+    }
+}

src/main/java/com/minelittlepony/unicopia/particle/ParticleHandle.java

@@ -1,11 +1,11 @@
 package com.minelittlepony.unicopia.particle;
 
 import java.util.Optional;
+import java.util.UUID;
 import java.util.function.Consumer;
 
 import com.minelittlepony.unicopia.ability.magic.Caster;
 import com.minelittlepony.unicopia.ability.magic.Spell;
-import com.minelittlepony.unicopia.ability.magic.spell.SpellType;
 
 import net.fabricmc.api.EnvType;
 import net.fabricmc.api.Environment;
@@ -60,13 +60,13 @@ public class ParticleHandle {
     public static final class Link {
 
         private Optional<Caster<?>> caster = Optional.empty();
-        private SpellType<?> effect;
+        private UUID effect;
         private boolean linked;
 
         public void attach(Caster<?> caster) {
             this.linked = true;
             this.caster = Optional.of(caster);
-            this.effect = caster.getSpellSlot().get(false).map(Spell::getType).orElse(null);
+            this.effect = caster.getSpellSlot().get(false).map(Spell::getUuid).orElse(null);
         }
 
         public void detach() {
@@ -83,7 +83,7 @@ public class ParticleHandle {
 
                 return Caster.of(e).orElse(null) == c
                         && c.getSpellSlot().get(false)
-                            .filter(s -> s.getType() == effect)
+                            .filter(s -> s.getUuid().equals(effect))
                             .isPresent()
                         && e != null
                         && c.getWorld().getEntityById(e.getEntityId()) != null;