From 0a93682662998b11c15b42f4321b18a6915bc07e Mon Sep 17 00:00:00 2001
From: Sollace <sollacea@gmail.com>
Date: Tue, 11 Jan 2022 12:23:20 +0200
Subject: [PATCH] Ensure previously-discovered traits properly update
 server-side when the registered values change for an item

---
 .../ability/magic/spell/trait/Trait.java      | 10 +++++++
 .../magic/spell/trait/TraitDiscovery.java     | 27 ++++++++++---------
 2 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/src/main/java/com/minelittlepony/unicopia/ability/magic/spell/trait/Trait.java b/src/main/java/com/minelittlepony/unicopia/ability/magic/spell/trait/Trait.java
index 44f5d7ce..6d6b25d1 100644
--- a/src/main/java/com/minelittlepony/unicopia/ability/magic/spell/trait/Trait.java
+++ b/src/main/java/com/minelittlepony/unicopia/ability/magic/spell/trait/Trait.java
@@ -6,10 +6,13 @@ import java.util.Map;
 import java.util.Optional;
 import java.util.function.Function;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import com.minelittlepony.common.client.gui.Tooltip;
 
 import net.minecraft.item.ItemStack;
+import net.minecraft.nbt.NbtElement;
+import net.minecraft.nbt.NbtList;
 import net.minecraft.text.LiteralText;
 import net.minecraft.text.TranslatableText;
 import net.minecraft.util.Formatting;
@@ -102,6 +105,13 @@ public enum Trait {
         return IDS.values();
     }
 
+    public static Stream<Trait> fromNbt(NbtList nbt) {
+        return nbt.stream()
+                .map(NbtElement::asString)
+                .map(Trait::fromId)
+                .flatMap(Optional::stream);
+    }
+
     public static Optional<Trait> fromId(Identifier id) {
         return Optional.ofNullable(IDS.getOrDefault(id, null));
     }
diff --git a/src/main/java/com/minelittlepony/unicopia/ability/magic/spell/trait/TraitDiscovery.java b/src/main/java/com/minelittlepony/unicopia/ability/magic/spell/trait/TraitDiscovery.java
index 7285c95a..7bd33d95 100644
--- a/src/main/java/com/minelittlepony/unicopia/ability/magic/spell/trait/TraitDiscovery.java
+++ b/src/main/java/com/minelittlepony/unicopia/ability/magic/spell/trait/TraitDiscovery.java
@@ -122,23 +122,24 @@ public class TraitDiscovery implements NbtSerialisable {
         NbtCompound disco = compound.getCompound("items");
         disco.getKeys().forEach(key -> {
             Optional.ofNullable(Identifier.tryParse(key)).ifPresent(id -> {
-                SpellTraits.fromNbt(disco.getCompound(key)).ifPresent(val -> {
+                loadTraits(id, disco.getCompound(key)).filter(SpellTraits::isPresent).ifPresent(val -> {
                     items.put(id, val);
                 });
             });
         });
-        compound.getList("traits", NbtElement.STRING_TYPE).stream()
-            .map(NbtElement::asString)
-            .map(Trait::fromId)
-            .filter(Optional::isPresent)
-            .map(Optional::get)
-            .forEach(this.traits::add);
-        compound.getList("unreadTraits", NbtElement.STRING_TYPE).stream()
-            .map(NbtElement::asString)
-            .map(Trait::fromId)
-            .filter(Optional::isPresent)
-            .map(Optional::get)
-            .forEach(this.unreadTraits::add);
+        Trait.fromNbt(compound.getList("traits", NbtElement.STRING_TYPE)).forEach(traits::add);
+        Trait.fromNbt(compound.getList("unreadTraits", NbtElement.STRING_TYPE)).forEach(unreadTraits::add);
+    }
+
+    private Optional<SpellTraits> loadTraits(Identifier itemId, NbtCompound nbt) {
+        if (!pony.isClient()) {
+            return Registry.ITEM.getOrEmpty(itemId)
+                    .flatMap(item -> Optional.of(SpellTraits.of(item)))
+                    .filter(SpellTraits::isPresent)
+                    .or(() -> SpellTraits.fromNbt(nbt));
+        }
+
+        return SpellTraits.fromNbt(nbt);
     }
 
     public void copyFrom(TraitDiscovery old) {