Fix more issues with registering vanilla tags, works better now

core/src/main/java/org/geysermc/geyser/GeyserImpl.java

@@ -79,6 +79,7 @@ import org.geysermc.geyser.session.GeyserSession;
 import org.geysermc.geyser.session.PendingMicrosoftAuthentication;
 import org.geysermc.geyser.session.SessionManager;
 import org.geysermc.geyser.session.cache.RegistryCache;
+import org.geysermc.geyser.session.cache.tags.TagRegistry;
 import org.geysermc.geyser.skin.FloodgateSkinUploader;
 import org.geysermc.geyser.skin.ProvidedSkins;
 import org.geysermc.geyser.skin.SkinProvider;
@@ -219,6 +220,7 @@ public class GeyserImpl implements GeyserApi {
         /* Initialize registries */
         Registries.init();
         BlockRegistries.init();
+        TagRegistry.init();
 
         RegistryCache.init();
 

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

@@ -66,7 +66,7 @@ public final class TagCache {
 
         int i = 0;
         for (Key registryKey : allTags.keySet()) {
-            TagRegistry registry = TagRegistry.valueOf(registryKey);
+            TagRegistry registry = TagRegistry.fromKey(registryKey);
             if (registry == null) {
                 GeyserImpl.getInstance().getLogger().debug("Not loading tags for registry " + registryKey + " (registry is not defined in TagRegistry enum)");
                 continue;
@@ -91,19 +91,18 @@ public final class TagCache {
                 }
             }
 
-            int[][] registryTagsArray = new int[0][];
+            Object2IntMap<Key> tagIndexMap = new Object2IntOpenHashMap<>();
-            this.tagIndexMaps.add(loadTags(registryTags, registryTagsArray, registry));
+            this.tags[i] = loadTags(registryTags, tagIndexMap, registry);
-            this.tags[i] = registryTagsArray;
+            this.tagIndexMaps.add(tagIndexMap);
             i++;
         }
     }
 
-    private Object2IntMap<Key> loadTags(Map<Key, int[]> packetTags, int[][] tags, TagRegistry registry) {
+    private int[][] loadTags(Map<Key, int[]> packetTags, Object2IntMap<Key> tagIndexMap, TagRegistry registry) {
         List<Key> vanillaTagKeys = List.copyOf(registry.getVanillaTags().keySet());
         List<Key> nonVanillaTagKeys = packetTags.keySet().stream().filter(tag -> !vanillaTagKeys.contains(tag)).toList();
 
         List<int[]> tagsBuilder = new ArrayList<>(vanillaTagKeys.size() + nonVanillaTagKeys.size());
-        Object2IntMap<Key> tagIndexMap = new Object2IntOpenHashMap<>();
 
         for (Key vanillaTagKey : registry.getVanillaTags().keySet()) {
             tagsBuilder.add(packetTags.getOrDefault(vanillaTagKey, new int[0]));
@@ -114,8 +113,7 @@ public final class TagCache {
             tagsBuilder.add(packetTags.get(nonVanillaTagKey));
         }
 
-        tagsBuilder.toArray(tags);
+        return tagsBuilder.toArray(new int[0][]);
-        return tagIndexMap;
     }
 
     /**

core/src/main/java/org/geysermc/geyser/session/cache/tags/BlockTag.java

@@ -36,8 +36,6 @@ import java.util.Map;
  */
 @SuppressWarnings("unused")
 public final class BlockTag {
-    public static final Map<Key, Tag> ALL_BLOCK_TAGS = new HashMap<>();
-
     public static final Tag WOOL = register("wool");
     public static final Tag PLANKS = register("planks");
     public static final Tag STONE_BRICKS = register("stone_bricks");
@@ -226,10 +224,10 @@ public final class BlockTag {
     private BlockTag() {}
 
     private static Tag register(String name) {
-        Key identifier = MinecraftKey.key(name);
+        return TagRegistry.BLOCK.registerVanillaTag(MinecraftKey.key(name));
-        int geyserId = ALL_BLOCK_TAGS.size();
+    }
-        Tag tag = new VanillaTag(TagRegistry.BLOCK, identifier, geyserId);
+
-        ALL_BLOCK_TAGS.put(identifier, tag);
+    public static void init() {
-        return tag;
+        // no-op
     }
 }

core/src/main/java/org/geysermc/geyser/session/cache/tags/EnchantmentTag.java

@@ -28,16 +28,11 @@ package org.geysermc.geyser.session.cache.tags;
 import net.kyori.adventure.key.Key;
 import org.geysermc.geyser.util.MinecraftKey;
 
-import java.util.HashMap;
-import java.util.Map;
-
 /**
  * Lists vanilla enchantment tags.
  */
 @SuppressWarnings("unused")
 public final class EnchantmentTag {
-    public static final Map<Key, Tag> ALL_ENCHANTMENT_TAGS = new HashMap<>();
-
     public static final Tag TOOLTIP_ORDER = register("tooltip_order");
     public static final Tag EXCLUSIVE_SET_ARMOR = register("exclusive_set/armor");
     public static final Tag EXCLUSIVE_SET_BOOTS = register("exclusive_set/boots");
@@ -78,10 +73,10 @@ public final class EnchantmentTag {
     private EnchantmentTag() {}
 
     private static Tag register(String name) {
-        Key identifier = MinecraftKey.key(name);
+        return TagRegistry.ENCHANTMENT.registerVanillaTag(MinecraftKey.key(name));
-        int geyserId = ALL_ENCHANTMENT_TAGS.size();
+    }
-        Tag tag = new VanillaTag(TagRegistry.ENCHANTMENT, identifier, geyserId);
+
-        ALL_ENCHANTMENT_TAGS.put(MinecraftKey.key(name), tag);
+    public static void init() {
-        return tag;
+        // no-op
     }
 }

core/src/main/java/org/geysermc/geyser/session/cache/tags/ItemTag.java

@@ -36,8 +36,6 @@ import java.util.Map;
  */
 @SuppressWarnings("unused")
 public final class ItemTag {
-    public static final Map<Key, Tag> ALL_ITEM_TAGS = new HashMap<>();
-
     public static final Tag WOOL = register("wool");
     public static final Tag PLANKS = register("planks");
     public static final Tag STONE_BRICKS = register("stone_bricks");
@@ -189,10 +187,10 @@ public final class ItemTag {
     private ItemTag() {}
 
     private static Tag register(String name) {
-        Key identifier = MinecraftKey.key(name);
+        return TagRegistry.ITEM.registerVanillaTag(MinecraftKey.key(name));
-        int geyserId = ALL_ITEM_TAGS.size();
+    }
-        Tag tag = new VanillaTag(TagRegistry.ITEM, identifier, geyserId);
+
-        ALL_ITEM_TAGS.put(MinecraftKey.key(name), tag);
+    public static void init() {
-        return tag;
+        // no-op
     }
 }

core/src/main/java/org/geysermc/geyser/session/cache/tags/TagRegistry.java

@@ -25,8 +25,8 @@
 
 package org.geysermc.geyser.session.cache.tags;
 
+import java.util.HashMap;
 import java.util.Map;
-import lombok.Getter;
 import net.kyori.adventure.key.Key;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.geysermc.geyser.util.MinecraftKey;
@@ -34,29 +34,43 @@ import org.geysermc.geyser.util.MinecraftKey;
 /**
  * Lists registries that Geyser stores tags for.
  *
- * When wanting to store tags from a new registry, add the registry here, and add a map for all the vanilla tags in the registry. These vanilla tags
+ * When wanting to store tags from a new registry, add the registry here, and register all vanilla tags for it using {@link TagRegistry#registerVanillaTag}. These vanilla tags
- * can be stored in a vanilla tag class, like {@link BlockTag} and {@link ItemTag}.
+ * can be stored in a vanilla tag class, like {@link BlockTag} and {@link ItemTag}. This class can then have an init method that's called in {@link TagRegistry#init)}, to ensure
+ * that all vanilla tags are registered before any connection is made.
  */
-@Getter
 public enum TagRegistry {
-    BLOCK("block", BlockTag.ALL_BLOCK_TAGS),
+    BLOCK("block"),
-    ITEM("item", ItemTag.ALL_ITEM_TAGS),
+    ITEM("item"),
-    ENCHANTMENT("enchantment", EnchantmentTag.ALL_ENCHANTMENT_TAGS);
+    ENCHANTMENT("enchantment");
 
     private final Key registryKey;
     private final Map<Key, Tag> vanillaTags;
 
-    TagRegistry(String registry, Map<Key, Tag> vanillaTags) {
+    TagRegistry(String registry) {
         this.registryKey = MinecraftKey.key(registry);
-        this.vanillaTags = vanillaTags;
+        this.vanillaTags = new HashMap<>();
+    }
+
+    public Tag registerVanillaTag(Key identifier) {
+        if (vanillaTags.containsKey(identifier)) {
+            throw new IllegalArgumentException("Vanilla tag " + identifier + " was already registered!");
+        }
+
+        Tag tag = new VanillaTag(this, identifier, vanillaTags.size());
+        vanillaTags.put(identifier, tag);
+        return tag;
+    }
+
+    public Map<Key, Tag> getVanillaTags() {
+        return Map.copyOf(vanillaTags);
     }
 
     public static boolean shouldLoad(Key registryKey) {
-        return valueOf(registryKey) != null;
+        return fromKey(registryKey) != null;
     }
 
     @Nullable
-    public static TagRegistry valueOf(Key registryKey) {
+    public static TagRegistry fromKey(Key registryKey) {
         for (TagRegistry registry : TagRegistry.values()) {
             if (registry.registryKey.equals(registryKey)) {
                 return registry;
@@ -64,4 +78,10 @@ public enum TagRegistry {
         }
         return null;
     }
+
+    public static void init() {
+        BlockTag.init();
+        ItemTag.init();
+        EnchantmentTag.init();
+    }
 }