From 0910ad72cad26e91d134b7b01eb0eaeb8556b3ed Mon Sep 17 00:00:00 2001
From: Sollace <sollacea@gmail.com>
Date: Fri, 26 Jan 2024 15:58:22 +0000
Subject: [PATCH] Fixed dismissing spells from the gui killing them without
 going through the normal removal process

---
 .../unicopia/ability/magic/SpellContainer.java         |  7 +++++++
 .../client/render/spell/PortalFrameBuffer.java         |  6 +-----
 .../unicopia/network/MsgRemoveSpell.java               |  4 ++--
 .../unicopia/network/datasync/EffectSync.java          | 10 ++++++++++
 .../network/datasync/NetworkedReferenceSet.java        |  8 +++++++-
 5 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/src/main/java/com/minelittlepony/unicopia/ability/magic/SpellContainer.java b/src/main/java/com/minelittlepony/unicopia/ability/magic/SpellContainer.java
index bbe1ad70..e703911d 100644
--- a/src/main/java/com/minelittlepony/unicopia/ability/magic/SpellContainer.java
+++ b/src/main/java/com/minelittlepony/unicopia/ability/magic/SpellContainer.java
@@ -39,6 +39,13 @@ public interface SpellContainer {
      */
     void put(@Nullable Spell effect);
 
+    /**
+     * Cleanly removes a spell from this spell container.
+     *
+     * @param spellid ID of the spell to remove.
+     */
+    void remove(UUID spellid);
+
     /**
      * Removes all active effects that match or contain a matching effect.
      *
diff --git a/src/main/java/com/minelittlepony/unicopia/client/render/spell/PortalFrameBuffer.java b/src/main/java/com/minelittlepony/unicopia/client/render/spell/PortalFrameBuffer.java
index dffa297b..81a99039 100644
--- a/src/main/java/com/minelittlepony/unicopia/client/render/spell/PortalFrameBuffer.java
+++ b/src/main/java/com/minelittlepony/unicopia/client/render/spell/PortalFrameBuffer.java
@@ -76,14 +76,10 @@ class PortalFrameBuffer implements AutoCloseable {
 
     private boolean pendingDraw;
 
-    private final UUID id;
-
     @Nullable
     private Frustum frustum;
 
-    PortalFrameBuffer(UUID id) {
-        this.id = id;
-    }
+    PortalFrameBuffer(UUID id) { }
 
     public void draw(MatrixStack matrices, VertexConsumerProvider vertices) {
         matrices.translate(0, -0.001, 0);
diff --git a/src/main/java/com/minelittlepony/unicopia/network/MsgRemoveSpell.java b/src/main/java/com/minelittlepony/unicopia/network/MsgRemoveSpell.java
index 37b1e0a4..771b2c73 100644
--- a/src/main/java/com/minelittlepony/unicopia/network/MsgRemoveSpell.java
+++ b/src/main/java/com/minelittlepony/unicopia/network/MsgRemoveSpell.java
@@ -10,7 +10,7 @@ import net.minecraft.network.PacketByteBuf;
 import net.minecraft.server.network.ServerPlayerEntity;
 
 /**
- * Sent to the server when a player activates an ability.
+ * Sent to the server when a player dismisses a spell from their dismiss spell screen
  */
 public record MsgRemoveSpell (UUID id) implements HandledPacket<ServerPlayerEntity> {
     MsgRemoveSpell(PacketByteBuf buffer) {
@@ -30,7 +30,7 @@ public record MsgRemoveSpell (UUID id) implements HandledPacket<ServerPlayerEnti
     public void handle(ServerPlayerEntity sender) {
         Pony player = Pony.of(sender);
         if (player != null) {
-            player.getSpellSlot().removeIf(spell -> spell.getUuid().equals(id), true);
+            player.getSpellSlot().remove(id);
         }
     }
 }
diff --git a/src/main/java/com/minelittlepony/unicopia/network/datasync/EffectSync.java b/src/main/java/com/minelittlepony/unicopia/network/datasync/EffectSync.java
index b2cad93c..83507e6b 100644
--- a/src/main/java/com/minelittlepony/unicopia/network/datasync/EffectSync.java
+++ b/src/main/java/com/minelittlepony/unicopia/network/datasync/EffectSync.java
@@ -93,6 +93,16 @@ public class EffectSync implements SpellContainer, NbtSerialisable {
         }
     }
 
+    @Override
+    public void remove(UUID id) {
+        Spell spell = spells.getReference(id);
+        spell.setDead();
+        spell.tickDying(owner);
+        if (spell.isDead()) {
+            spells.removeReference(id);
+        }
+    }
+
     @Override
     public boolean removeWhere(Predicate<Spell> test, boolean update) {
         return reduce(update, (initial, effect) -> {
diff --git a/src/main/java/com/minelittlepony/unicopia/network/datasync/NetworkedReferenceSet.java b/src/main/java/com/minelittlepony/unicopia/network/datasync/NetworkedReferenceSet.java
index df54150f..90582e52 100644
--- a/src/main/java/com/minelittlepony/unicopia/network/datasync/NetworkedReferenceSet.java
+++ b/src/main/java/com/minelittlepony/unicopia/network/datasync/NetworkedReferenceSet.java
@@ -83,7 +83,13 @@ public class NetworkedReferenceSet<T> {
         }
     }
 
-    private synchronized void removeReference(UUID id) {
+    @Nullable
+    synchronized T getReference(UUID id) {
+        NetworkedReference<T> i = values.get(id);
+        return i == null ? null : i.getReference().orElse(null);
+    }
+
+    synchronized void removeReference(UUID id) {
         dirty |= ids.remove(id);
         NetworkedReference<T> i = values.remove(id);
         if (i != null) {