Fix reading enchantments from server (#4836)

core/src/main/java/org/geysermc/geyser/inventory/recipe/TrimRecipe.java

@@ -32,9 +32,8 @@ import org.cloudburstmc.protocol.bedrock.data.TrimPattern;
 import org.cloudburstmc.protocol.bedrock.data.inventory.descriptor.ItemDescriptorWithCount;
 import org.cloudburstmc.protocol.bedrock.data.inventory.descriptor.ItemTagDescriptor;
 import org.geysermc.geyser.registry.type.ItemMapping;
-import org.geysermc.geyser.session.GeyserSession;
+import org.geysermc.geyser.session.cache.registry.RegistryEntryContext;
 import org.geysermc.geyser.translator.text.MessageTranslator;
-import org.geysermc.mcprotocollib.protocol.data.game.RegistryEntry;
 
 /**
  * Stores information on trim materials and patterns, including smithing armor hacks for pre-1.20.
@@ -46,18 +45,18 @@ public final class TrimRecipe {
     public static final ItemDescriptorWithCount ADDITION = tagDescriptor("minecraft:trim_materials");
     public static final ItemDescriptorWithCount TEMPLATE = tagDescriptor("minecraft:trim_templates");
 
-    public static TrimMaterial readTrimMaterial(GeyserSession session, RegistryEntry entry) {
+    public static TrimMaterial readTrimMaterial(RegistryEntryContext context) {
-        String key = entry.getId().asMinimalString();
+        String key = context.id().asMinimalString();
 
         // Color is used when hovering over the item
         // Find the nearest legacy color from the RGB Java gives us to work with
         // Also yes this is a COMPLETE hack but it works ok!!!!!
-        String colorTag = entry.getData().getCompound("description").getString("color");
+        String colorTag = context.data().getCompound("description").getString("color");
         TextColor color = TextColor.fromHexString(colorTag);
         String legacy = MessageTranslator.convertMessage(Component.space().color(color));
 
-        String itemIdentifier = entry.getData().getString("ingredient");
+        String itemIdentifier = context.data().getString("ingredient");
-        ItemMapping itemMapping = session.getItemMappings().getMapping(itemIdentifier);
+        ItemMapping itemMapping = context.session().getItemMappings().getMapping(itemIdentifier);
         if (itemMapping == null) {
             // This should never happen so not sure what to do here.
             itemMapping = ItemMapping.AIR;
@@ -66,11 +65,11 @@ public final class TrimRecipe {
         return new TrimMaterial(key, legacy.substring(2).trim(), itemMapping.getBedrockIdentifier());
     }
 
-    public static TrimPattern readTrimPattern(GeyserSession session, RegistryEntry entry) {
+    public static TrimPattern readTrimPattern(RegistryEntryContext context) {
-        String key = entry.getId().asMinimalString();
+        String key = context.id().asMinimalString();
 
-        String itemIdentifier = entry.getData().getString("template_item");
+        String itemIdentifier = context.data().getString("template_item");
-        ItemMapping itemMapping = session.getItemMappings().getMapping(itemIdentifier);
+        ItemMapping itemMapping = context.session().getItemMappings().getMapping(itemIdentifier);
         if (itemMapping == null) {
             // This should never happen so not sure what to do here.
             itemMapping = ItemMapping.AIR;

core/src/main/java/org/geysermc/geyser/inventory/updater/AnvilInventoryUpdater.java

@@ -27,6 +27,7 @@ package org.geysermc.geyser.inventory.updater;
 
 import it.unimi.dsi.fastutil.objects.Object2IntMap;
 import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
+import java.util.stream.IntStream;
 import net.kyori.adventure.text.Component;
 import org.cloudburstmc.nbt.NbtMap;
 import org.cloudburstmc.nbt.NbtMapBuilder;
@@ -41,11 +42,14 @@ import org.geysermc.geyser.inventory.item.BedrockEnchantment;
 import org.geysermc.geyser.item.enchantment.Enchantment;
 import org.geysermc.geyser.item.Items;
 import org.geysermc.geyser.session.GeyserSession;
+import org.geysermc.geyser.session.cache.tags.EnchantmentTag;
+import org.geysermc.geyser.session.cache.tags.ItemTag;
 import org.geysermc.geyser.translator.inventory.InventoryTranslator;
 import org.geysermc.geyser.translator.text.MessageTranslator;
 import org.geysermc.geyser.util.ItemUtils;
 import org.geysermc.mcprotocollib.protocol.data.game.entity.player.GameMode;
 import org.geysermc.mcprotocollib.protocol.data.game.item.component.DataComponentType;
+import org.geysermc.mcprotocollib.protocol.data.game.item.component.HolderSet;
 import org.geysermc.mcprotocollib.protocol.data.game.item.component.ItemEnchantments;
 import org.geysermc.mcprotocollib.protocol.packet.ingame.serverbound.inventory.ServerboundRenameItemPacket;
 
@@ -310,10 +314,12 @@ public class AnvilInventoryUpdater extends InventoryUpdater {
         for (Object2IntMap.Entry<Enchantment> entry : getEnchantments(session, material).object2IntEntrySet()) {
             Enchantment enchantment = entry.getKey();
 
-            boolean canApply = isEnchantedBook(input) || session.getTagCache().is(enchantment.supportedItems(), input);
+            HolderSet supportedItems = enchantment.supportedItems();
-            var exclusiveSet = enchantment.exclusiveSet();
+            int[] supportedItemIds = supportedItems.resolve(tagId -> session.getTagCache().get(ItemTag.ALL_ITEM_TAGS.get(tagId)));
-            if (exclusiveSet != null) {
+            boolean canApply = isEnchantedBook(input) || IntStream.of(supportedItemIds).anyMatch(id -> id == input.getJavaId());
-                int[] incompatibleEnchantments = session.getTagCache().get(exclusiveSet);
+
+            HolderSet exclusiveSet = enchantment.exclusiveSet();
+            int[] incompatibleEnchantments = exclusiveSet.resolve(tagId -> session.getTagCache().get(EnchantmentTag.ALL_ENCHANTMENT_TAGS.get(tagId)));
             for (int i : incompatibleEnchantments) {
                 Enchantment incompatible = session.getRegistryCache().enchantments().byId(i);
                 if (combinedEnchantments.containsKey(incompatible)) {
@@ -323,7 +329,6 @@ public class AnvilInventoryUpdater extends InventoryUpdater {
                     }
                 }
             }
-            }
 
             if (canApply || (!bedrock && session.getGameMode() == GameMode.CREATIVE)) {
                 int currentLevel = combinedEnchantments.getOrDefault(enchantment, 0);

core/src/main/java/org/geysermc/geyser/item/enchantment/Enchantment.java

@@ -25,18 +25,21 @@
 
 package org.geysermc.geyser.item.enchantment;
 
+import java.util.List;
+import java.util.function.Function;
+import net.kyori.adventure.key.Key;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.cloudburstmc.nbt.NbtMap;
 import org.geysermc.geyser.inventory.item.BedrockEnchantment;
-import org.geysermc.geyser.session.cache.tags.EnchantmentTag;
+import org.geysermc.geyser.item.Items;
-import org.geysermc.geyser.session.cache.tags.ItemTag;
+import org.geysermc.geyser.registry.Registries;
+import org.geysermc.geyser.session.cache.registry.RegistryEntryContext;
 import org.geysermc.geyser.translator.text.MessageTranslator;
-import org.geysermc.geyser.util.MinecraftKey;
-import org.geysermc.mcprotocollib.protocol.data.game.RegistryEntry;
 
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
+import org.geysermc.mcprotocollib.protocol.data.game.item.component.HolderSet;
 
 /**
  * @param description only populated if {@link #bedrockEnchantment()} is not null.
@@ -44,28 +47,32 @@ import java.util.Set;
  */
 public record Enchantment(String identifier,
                           Set<EnchantmentComponent> effects,
-                          ItemTag supportedItems,
+                          HolderSet supportedItems,
                           int maxLevel,
                           String description,
                           int anvilCost,
-                          @Nullable EnchantmentTag exclusiveSet,
+                          HolderSet exclusiveSet,
                           @Nullable BedrockEnchantment bedrockEnchantment) {
 
-    // Implementation note: I have a feeling the tags can be a list of items, because in vanilla they're HolderSet classes.
+    public static Enchantment read(RegistryEntryContext context) {
-    // I'm not sure how that's wired over the network, so we'll put it off.
+        NbtMap data = context.data();
-    public static Enchantment read(RegistryEntry entry) {
-        NbtMap data = entry.getData();
         Set<EnchantmentComponent> effects = readEnchantmentComponents(data.getCompound("effects"));
-        String supportedItems = data.getString("supported_items").substring(1); // Remove '#' at beginning that indicates tag
+
+        HolderSet supportedItems = readHolderSet(data.get("supported_items"), itemId -> Registries.JAVA_ITEM_IDENTIFIERS.getOrDefault(itemId.asString(), Items.AIR).javaId());
+
         int maxLevel = data.getInt("max_level");
         int anvilCost = data.getInt("anvil_cost");
-        String exclusiveSet = data.getString("exclusive_set", null);
+
-        EnchantmentTag exclusiveSetTag = exclusiveSet == null ? null : EnchantmentTag.ALL_ENCHANTMENT_TAGS.get(MinecraftKey.key(exclusiveSet.substring(1)));
+        HolderSet exclusiveSet = readHolderSet(data.getOrDefault("exclusive_set", null), context::getNetworkId);
-        BedrockEnchantment bedrockEnchantment = BedrockEnchantment.getByJavaIdentifier(entry.getId().asString());
+
+        BedrockEnchantment bedrockEnchantment = BedrockEnchantment.getByJavaIdentifier(context.id().asString());
+
+        // TODO - description is a component. So if a hardcoded literal string is given, this will display normally on Java,
+        //  but Geyser will attempt to lookup the literal string as translation - and will fail, displaying an empty string as enchantment name.
         String description = bedrockEnchantment == null ? MessageTranslator.deserializeDescription(data) : null;
 
-        return new Enchantment(entry.getId().asString(), effects, ItemTag.ALL_ITEM_TAGS.get(MinecraftKey.key(supportedItems)), maxLevel,
+        return new Enchantment(context.id().asString(), effects, supportedItems, maxLevel,
-                description, anvilCost, exclusiveSetTag, bedrockEnchantment);
+                description, anvilCost, exclusiveSet, bedrockEnchantment);
     }
 
     private static Set<EnchantmentComponent> readEnchantmentComponents(NbtMap effects) {
@@ -77,4 +84,24 @@ public record Enchantment(String identifier,
         }
         return Set.copyOf(components); // Also ensures any empty sets are consolidated
     }
+
+    // TODO holder set util?
+    private static HolderSet readHolderSet(@Nullable Object holderSet, Function<Key, Integer> keyIdMapping) {
+        if (holderSet == null) {
+            return new HolderSet(new int[]{});
+        }
+
+        if (holderSet instanceof String stringTag) {
+            // Tag
+            if (stringTag.startsWith("#")) {
+                return new HolderSet(Key.key(stringTag.substring(1))); // Remove '#' at beginning that indicates tag
+            } else {
+                return new HolderSet(new int[]{keyIdMapping.apply(Key.key(stringTag))});
+            }
+        } else if (holderSet instanceof List<?> list) {
+            // Assume the list is a list of strings
+            return new HolderSet(list.stream().map(o -> (String) o).map(Key::key).map(keyIdMapping).mapToInt(Integer::intValue).toArray());
+        }
+        throw new IllegalArgumentException("Holder set must either be a tag, a string ID or a list of string IDs");
+    }
 }

core/src/main/java/org/geysermc/geyser/level/JavaDimension.java

@@ -26,7 +26,7 @@
 package org.geysermc.geyser.level;
 
 import org.cloudburstmc.nbt.NbtMap;
-import org.geysermc.mcprotocollib.protocol.data.game.RegistryEntry;
+import org.geysermc.geyser.session.cache.registry.RegistryEntryContext;
 
 /**
  * Represents the information we store from the current Java dimension
@@ -35,8 +35,8 @@ import org.geysermc.mcprotocollib.protocol.data.game.RegistryEntry;
  */
 public record JavaDimension(int minY, int maxY, boolean piglinSafe, double worldCoordinateScale) {
 
-    public static JavaDimension read(RegistryEntry entry) {
+    public static JavaDimension read(RegistryEntryContext entry) {
-        NbtMap dimension = entry.getData();
+        NbtMap dimension = entry.data();
         int minY = dimension.getInt("min_y");
         int maxY = dimension.getInt("height");
         // Logical height can be ignored probably - seems to be for artificial limits like the Nether.

core/src/main/java/org/geysermc/geyser/level/JukeboxSong.java

@@ -27,13 +27,13 @@ package org.geysermc.geyser.level;
 
 import org.cloudburstmc.nbt.NbtMap;
 import org.geysermc.geyser.GeyserImpl;
+import org.geysermc.geyser.session.cache.registry.RegistryEntryContext;
 import org.geysermc.geyser.translator.text.MessageTranslator;
-import org.geysermc.mcprotocollib.protocol.data.game.RegistryEntry;
 
 public record JukeboxSong(String soundEvent, String description) {
 
-    public static JukeboxSong read(RegistryEntry entry) {
+    public static JukeboxSong read(RegistryEntryContext context) {
-        NbtMap data = entry.getData();
+        NbtMap data = context.data();
         Object soundEventObject = data.get("sound_event");
         String soundEvent;
         if (soundEventObject instanceof NbtMap map) {
@@ -42,7 +42,7 @@ public record JukeboxSong(String soundEvent, String description) {
             soundEvent = string;
         } else {
             soundEvent = "";
-            GeyserImpl.getInstance().getLogger().debug("Sound event for " + entry.getId() + " was of an unexpected type! Expected string or NBT map, got " + soundEventObject);
+            GeyserImpl.getInstance().getLogger().debug("Sound event for " + context.id() + " was of an unexpected type! Expected string or NBT map, got " + soundEventObject);
         }
         String description = MessageTranslator.deserializeDescription(data);
         return new JukeboxSong(soundEvent, description);

core/src/main/java/org/geysermc/geyser/session/cache/RegistryCache.java

@@ -27,6 +27,8 @@ package org.geysermc.geyser.session.cache;
 
 import it.unimi.dsi.fastutil.ints.Int2IntMap;
 import it.unimi.dsi.fastutil.ints.Int2IntOpenHashMap;
+import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
 import lombok.AccessLevel;
 import lombok.Getter;
 import lombok.experimental.Accessors;
@@ -45,6 +47,7 @@ import org.geysermc.geyser.level.JukeboxSong;
 import org.geysermc.geyser.level.PaintingType;
 import org.geysermc.geyser.session.GeyserSession;
 import org.geysermc.geyser.session.cache.registry.JavaRegistry;
+import org.geysermc.geyser.session.cache.registry.RegistryEntryContext;
 import org.geysermc.geyser.session.cache.registry.SimpleJavaRegistry;
 import org.geysermc.geyser.text.TextDecoration;
 import org.geysermc.geyser.translator.level.BiomeTranslator;
@@ -59,7 +62,6 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.function.BiConsumer;
-import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.ToIntFunction;
 
@@ -76,16 +78,16 @@ public final class RegistryCache {
     private static final Map<Key, BiConsumer<RegistryCache, List<RegistryEntry>>> REGISTRIES = new HashMap<>();
 
     static {
-        register("chat_type", cache -> cache.chatTypes, ($, entry) -> TextDecoration.readChatType(entry));
+        register("chat_type", cache -> cache.chatTypes, TextDecoration::readChatType);
-        register("dimension_type", cache -> cache.dimensions, ($, entry) -> JavaDimension.read(entry));
+        register("dimension_type", cache -> cache.dimensions, JavaDimension::read);
-        register("enchantment", cache -> cache.enchantments, ($, entry) -> Enchantment.read(entry));
+        register("enchantment", cache -> cache.enchantments, Enchantment::read);
-        register("jukebox_song", cache -> cache.jukeboxSongs, ($, entry) -> JukeboxSong.read(entry));
+        register("jukebox_song", cache -> cache.jukeboxSongs, JukeboxSong::read);
-        register("painting_variant", cache -> cache.paintings, ($, entry) -> PaintingType.getByName(entry.getId()));
+        register("painting_variant", cache -> cache.paintings, context -> PaintingType.getByName(context.id()));
         register("trim_material", cache -> cache.trimMaterials, TrimRecipe::readTrimMaterial);
         register("trim_pattern", cache -> cache.trimPatterns, TrimRecipe::readTrimPattern);
         register("worldgen/biome", (cache, array) -> cache.biomeTranslations = array, BiomeTranslator::loadServerBiome);
-        register("banner_pattern", cache -> cache.bannerPatterns, ($, entry) -> BannerPattern.getByJavaIdentifier(entry.getId()));
+        register("banner_pattern", cache -> cache.bannerPatterns, context -> BannerPattern.getByJavaIdentifier(context.id()));
-        register("wolf_variant", cache -> cache.wolfVariants, ($, entry) -> WolfEntity.BuiltInWolfVariant.getByJavaIdentifier(entry.getId().asString()));
+        register("wolf_variant", cache -> cache.wolfVariants, context -> WolfEntity.BuiltInWolfVariant.getByJavaIdentifier(context.id().asString()));
 
         // Load from MCProtocolLib's classloader
         NbtMap tag = MinecraftProtocol.loadNetworkCodec();
@@ -149,25 +151,35 @@ public final class RegistryCache {
      * @param reader converts the RegistryEntry NBT into a class file
      * @param <T> the class that represents these entries.
      */
-    private static <T> void register(String registry, Function<RegistryCache, JavaRegistry<T>> localCacheFunction, BiFunction<GeyserSession, RegistryEntry, T> reader) {
+    private static <T> void register(String registry, Function<RegistryCache, JavaRegistry<T>> localCacheFunction, Function<RegistryEntryContext, T> reader) {
-        Key key = MinecraftKey.key(registry);
+        Key registryKey = MinecraftKey.key(registry);
-        REGISTRIES.put(key, (registryCache, entries) -> {
+        REGISTRIES.put(registryKey, (registryCache, entries) -> {
             Map<Key, NbtMap> localRegistry = null;
             JavaRegistry<T> localCache = localCacheFunction.apply(registryCache);
             // Clear each local cache every time a new registry entry is given to us
             // (e.g. proxy server switches)
+
+            // Store each of the entries resource location IDs and their respective network ID,
+            // used for the key mapper that's currently only used by the Enchantment class
+            Object2IntMap<Key> entryIdMap = new Object2IntOpenHashMap<>();
+            for (int i = 0; i < entries.size(); i++) {
+                entryIdMap.put(entries.get(i).getId(), i);
+            }
+
             List<T> builder = new ArrayList<>(entries.size());
             for (int i = 0; i < entries.size(); i++) {
                 RegistryEntry entry = entries.get(i);
                 // If the data is null, that's the server telling us we need to use our default values.
                 if (entry.getData() == null) {
                     if (localRegistry == null) { // Lazy initialize
-                        localRegistry = DEFAULTS.get(key);
+                        localRegistry = DEFAULTS.get(registryKey);
                     }
                     entry = new RegistryEntry(entry.getId(), localRegistry.get(entry.getId()));
                 }
+
+                RegistryEntryContext context = new RegistryEntryContext(entry, entryIdMap, registryCache.session);
                 // This is what Geyser wants to keep as a value for this registry.
-                T cacheEntry = reader.apply(registryCache.session, entry);
+                T cacheEntry = reader.apply(context);
                 builder.add(i, cacheEntry);
             }
             localCache.reset(builder);

core/src/main/java/org/geysermc/geyser/session/cache/TagCache.java

@@ -130,8 +130,12 @@ public final class TagCache {
         return contains(values, item.javaId());
     }
 
-    public int[] get(EnchantmentTag tag) {
+    public int[] get(ItemTag itemTag) {
-        return this.enchantments[tag.ordinal()];
+        return this.items[itemTag.ordinal()];
+    }
+
+    public int[] get(EnchantmentTag enchantmentTag) {
+        return this.enchantments[enchantmentTag.ordinal()];
     }
 
     private static boolean contains(int[] array, int i) {

core/src/main/java/org/geysermc/geyser/session/cache/registry/RegistryEntryContext.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright (c) 2024 GeyserMC. http://geysermc.org
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ *
+ * @author GeyserMC
+ * @link https://github.com/GeyserMC/Geyser
+ */
+
+package org.geysermc.geyser.session.cache.registry;
+
+import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import java.util.Map;
+import net.kyori.adventure.key.Key;
+import org.cloudburstmc.nbt.NbtMap;
+import org.geysermc.geyser.session.GeyserSession;
+import org.geysermc.mcprotocollib.protocol.data.game.RegistryEntry;
+
+/**
+ * Used to store context around a single registry entry when reading said entry's NBT.
+ *
+ * @param entry the registry entry being read.
+ * @param keyIdMap a map for each of the resource location's in the registry and their respective network IDs.
+ * @param session the Geyser session.
+ */
+public record RegistryEntryContext(RegistryEntry entry, Object2IntMap<Key> keyIdMap, GeyserSession session) {
+
+    public int getNetworkId(Key registryKey) {
+        return keyIdMap.getOrDefault(registryKey, 0);
+    }
+
+    public Key id() {
+        return entry.getId();
+    }
+
+    // Not annotated as nullable because data should never be null here
+    public NbtMap data() {
+        return entry.getData();
+    }
+}

core/src/main/java/org/geysermc/geyser/text/TextDecoration.java

@@ -29,7 +29,7 @@ import net.kyori.adventure.text.format.NamedTextColor;
 import net.kyori.adventure.text.format.Style;
 import org.cloudburstmc.nbt.NbtMap;
 import org.cloudburstmc.nbt.NbtType;
-import org.geysermc.mcprotocollib.protocol.data.game.RegistryEntry;
+import org.geysermc.geyser.session.cache.registry.RegistryEntryContext;
 import org.geysermc.mcprotocollib.protocol.data.game.chat.ChatType;
 import org.geysermc.mcprotocollib.protocol.data.game.chat.ChatTypeDecoration;
 
@@ -43,11 +43,11 @@ public record TextDecoration(String translationKey, List<Parameter> parameters,
         throw new UnsupportedOperationException();
     }
 
-    public static ChatType readChatType(RegistryEntry entry) {
+    public static ChatType readChatType(RegistryEntryContext context) {
         // Note: The ID is NOT ALWAYS THE SAME! ViaVersion as of 1.19 adds two registry entries that do NOT match vanilla.
         // (This note has been passed around through several classes and iterations. It stays as a warning
         // to anyone that dares to try and hardcode registry IDs.)
-        NbtMap tag = entry.getData();
+        NbtMap tag = context.data();
         NbtMap chat = tag.getCompound("chat", null);
         if (chat != null) {
             String translationKey = chat.getString("translation_key");