Fix bugs with spell syncing

src/main/java/com/minelittlepony/unicopia/ability/magic/SpellContainer.java

@@ -1,6 +1,7 @@
 package com.minelittlepony.unicopia.ability.magic;
 
 import java.util.Optional;
+import java.util.function.Function;
 import java.util.function.Predicate;
 
 import org.jetbrains.annotations.Nullable;
@@ -22,6 +23,11 @@ public interface SpellContainer {
 
         @Override
         public void removeIf(Predicate<Spell> effect, boolean update) { }
+
+        @Override
+        public boolean forEach(Function<Spell, Operation> action, boolean update) {
+            return false;
+        }
     };
 
     /**
@@ -51,7 +57,14 @@ public interface SpellContainer {
     /**
      * Removes all matching active effects.
      */
-    void removeIf(Predicate<Spell> effect, boolean update);
+    void removeIf(Predicate<Spell> test, boolean update);
+
+    /**
+     * Iterates active spells and optionally removes matching ones.
+     *
+     * @return True if any matching spells remain active
+     */
+    boolean forEach(Function<Spell, Operation> action, boolean update);
 
     /**
      * Removes all effects currently active in this slot.
@@ -77,9 +90,20 @@ public interface SpellContainer {
             delegate().removeIf(effect, update);
         }
 
+        @Override
+        default boolean forEach(Function<Spell, Operation> action, boolean update) {
+            return delegate().forEach(action, update);
+        }
+
         @Override
         default void clear() {
             delegate().clear();
         }
     }
+
+    public enum Operation {
+        SKIP,
+        KEEP,
+        REMOVE
+    }
 }

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

@@ -141,7 +141,6 @@ public final class SpellType<T extends Spell> implements Affine, SpellPredicate<
     @Nullable
     public T apply(Caster<?> caster, SpellTraits traits) {
         if (isEmpty()) {
-            caster.getSpellSlot().clear();
             return null;
         }
 

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

@@ -7,6 +7,7 @@ import com.minelittlepony.unicopia.Affinity;
 import com.minelittlepony.unicopia.ability.magic.Caster;
 import com.minelittlepony.unicopia.ability.magic.Levelled;
 import com.minelittlepony.unicopia.ability.magic.SpellContainer;
+import com.minelittlepony.unicopia.ability.magic.SpellContainer.Operation;
 import com.minelittlepony.unicopia.ability.magic.spell.Situation;
 import com.minelittlepony.unicopia.ability.magic.spell.Spell;
 import com.minelittlepony.unicopia.network.Channel;
@@ -97,24 +98,21 @@ public class CastSpellEntity extends Entity implements Caster<LivingEntity> {
             if (orphanedTicks-- > 0) {
                 return;
             }
-            remove(RemovalReason.DISCARDED);
+            discard();
             return;
         }
 
         orphanedTicks = 0;
 
-        if (!Caster.of(master).filter(c -> {
-            UUID spellId = dataTracker.get(SPELL).orElse(null);
-
-            if (!c.getSpellSlot().get(true).filter(s -> s.getUuid().equals(spellId) && s.tick(this, Situation.GROUND_ENTITY)).isPresent()) {
-                c.getSpellSlot().clear();
-
-                return false;
-            }
-
-            return true;
+        if (!dataTracker.get(SPELL).filter(spellId -> {
+            return getSpellSlot().forEach(spell -> {
+                if (!spell.getUuid().equals(spellId)) {
+                    return Operation.SKIP;
+                }
+                return spell.tick(this, Situation.GROUND_ENTITY) ? Operation.KEEP: Operation.REMOVE;
+            }, true);
         }).isPresent()) {
-            remove(RemovalReason.DISCARDED);
+            discard();
         }
     }
 

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

@@ -6,6 +6,8 @@ import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
 
 import com.minelittlepony.unicopia.ability.magic.Caster;
+import com.minelittlepony.unicopia.ability.magic.SpellContainer;
+
 import net.minecraft.entity.LivingEntity;
 import net.minecraft.item.Item;
 import net.minecraft.item.ItemStack;
@@ -18,6 +20,6 @@ abstract class MixinMilkBucketItem extends Item {
 
     @Inject(method = "finishUsing", at = @At("HEAD"), cancellable = true)
     private void finishUsing(ItemStack stack, World world, LivingEntity entity, CallbackInfoReturnable<ItemStack> info) {
-        Caster.of(entity).ifPresent(c -> c.getSpellSlot().clear());
+        Caster.of(entity).map(Caster::getSpellSlot).ifPresent(SpellContainer::clear);
     }
 }

src/main/java/com/minelittlepony/unicopia/network/datasync/EffectSync.java

@@ -2,6 +2,7 @@ package com.minelittlepony.unicopia.network.datasync;
 
 import java.util.Optional;
 import java.util.function.Consumer;
+import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.Stream;
 
@@ -61,16 +62,37 @@ public class EffectSync implements SpellContainer {
 
     @Override
     public void removeIf(Predicate<Spell> test, boolean update) {
-        read(effect -> {
+        forEach(effect -> {
             if (test.test(effect)) {
                 spells.removeReference(effect);
             }
         });
     }
 
+    @Override
+    public boolean forEach(Function<Spell, Operation> test, boolean update) {
+        boolean[] matched = new boolean[1];
+
+        forEach(effect -> {
+            Operation op = test.apply(effect);
+            if (op == Operation.REMOVE) {
+                spells.removeReference(effect);
+            } else {
+                matched[0] |= op != Operation.SKIP;
+            }
+        });
+
+        return matched[0];
+    }
+
     @Override
     public void clear() {
-        put(null);
+        if (spells.clear()) {
+            write();
+            if (owner instanceof UpdateCallback) {
+                ((UpdateCallback)owner).onSpellSet(null);
+            }
+        }
     }
 
     private Stream<Spell> read(boolean synchronize, boolean sendUpdate) {
@@ -81,7 +103,7 @@ public class EffectSync implements SpellContainer {
         return spells.getReferences();
     }
 
-    private void read(Consumer<Spell> consumer) {
+    public void forEach(Consumer<Spell> consumer) {
         spells.fromNbt(owner.getEntity().getDataTracker().get(param));
         spells.getReferences().toList().forEach(consumer);
         write();

src/main/java/com/minelittlepony/unicopia/network/datasync/NetworkedReferenceSet.java

@@ -40,9 +40,16 @@ public class NetworkedReferenceSet<T> {
                 .map(Optional::get);
     }
 
+    public boolean clear() {
+        dirty |= !ids.isEmpty() || !values.isEmpty();
+        ids.clear();
+        values.clear();
+        return dirty;
+    }
+
     public void addReference(@Nullable T newValue) {
         if (newValue != null) {
-            addReference(uuidConverter.apply(newValue));
+            addReference(uuidConverter.apply(newValue)).updateReference(newValue);
         }
     }
 
@@ -77,7 +84,7 @@ public class NetworkedReferenceSet<T> {
             incoming.add(UUID.fromString(key.asString()));
         });
 
-        ids.stream().filter(id -> !incoming.contains(id)).forEach(this::removeReference);
+        ids.stream().filter(id -> !incoming.contains(id)).toList().forEach(this::removeReference);
 
         boolean[] send = new boolean[1];
         incoming.forEach(kept -> {
@@ -99,7 +106,7 @@ public class NetworkedReferenceSet<T> {
             ids.add(NbtString.of(sid));
             tag.put(sid, values.get(id).toNbt());
         });
-        tag.put("key", ids);
+        tag.put("keys", ids);
         dirty = false;
         return tag;
     }