Support heterogeneous lists in GeyserHolderSet

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

@@ -138,7 +138,7 @@ public <T> boolean is(@NonNull GeyserHolderSet<T> holderSet, @Nullable T object)
         if (object == null) {
             return false;
         }
-        return holderSet.resolveRaw(this).contains(holderSet.getRegistry().networkId(session, object));
+        return holderSet.contains(session, object);
     }
 
     /**

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

@@ -39,6 +39,7 @@
 import org.geysermc.geyser.util.MinecraftKey;
 import org.geysermc.mcprotocollib.protocol.data.game.item.component.HolderSet;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
 import java.util.function.Function;
@@ -64,28 +65,46 @@ public final class GeyserHolderSet<T> {
     private final @Nullable Tag<T> tag;
     private final @Nullable IntList holders;
     private final @Nullable List<T> inline;
+    private final @Nullable List<HolderEntry<T>> entries;
+
+    public sealed interface HolderEntry<T> permits HolderEntry.Direct, HolderEntry.Id {
+        record Direct<T>(T value) implements HolderEntry<T> {}
+        record Id<T>(int id) implements HolderEntry<T> {}
+    }
 
     private GeyserHolderSet(JavaRegistryKey<T> registry) {
         this(registry, IntLists.emptyList());
     }
 
     public GeyserHolderSet(JavaRegistryKey<T> registry, @Nullable IntList holders) {
-        this(registry, null, holders, null);
+        this(registry, null, holders, null, null);
     }
 
     public GeyserHolderSet(JavaRegistryKey<T> registry, @NonNull Tag<T> tagId) {
-        this(registry, tagId, null, null);
+        this(registry, tagId, null, null, null);
     }
 
     public GeyserHolderSet(JavaRegistryKey<T> registry, @NonNull List<T> inline) {
-        this(registry, null, null, inline);
+        this(registry, null, null, inline, null);
     }
 
     private GeyserHolderSet(JavaRegistryKey<T> registry, @Nullable Tag<T> tag, @Nullable IntList holders, @Nullable List<T> inline) {
+        this(registry, tag, holders, inline, null);
+    }
+
+    private GeyserHolderSet(JavaRegistryKey<T> registry, @Nullable Tag<T> tag, @Nullable IntList holders, @Nullable List<T> inline, @Nullable List<HolderEntry<T>> entries) {
         this.registry = registry;
         this.tag = tag;
         this.holders = holders;
         this.inline = inline;
+        this.entries = entries;
+    }
+
+    /**
+     * Constructs a {@link GeyserHolderSet} from an ordered list of {@link HolderEntry} elements.
+     */
+    public static <T> GeyserHolderSet<T> ofEntries(JavaRegistryKey<T> registry, @NonNull List<HolderEntry<T>> entries) {
+        return new GeyserHolderSet<>(registry, null, null, null, entries);
     }
 
     /**
@@ -101,39 +120,72 @@ public static <T> GeyserHolderSet<T> empty(JavaRegistryKey<T> registry) {
     public static <T> GeyserHolderSet<T> fromHolderSet(JavaRegistryKey<T> registry, @NonNull HolderSet holderSet) {
         // MCPL HolderSets don't have to support inline elements... for now (TODO CHECK ME)
         Tag<T> tag = holderSet.getLocation() == null ? null : new Tag<>(registry, holderSet.getLocation());
-        return new GeyserHolderSet<>(registry, tag, holderSet.getHolders(), null);
+        return new GeyserHolderSet<>(registry, tag, holderSet.getHolders(), null, null);
     }
 
     public boolean contains(@NonNull GeyserSession session, @Nullable T object) {
         if (object == null) {
             return false;
         }
+        if (inline != null) {
+            return inline.contains(object);
+        }
+        if (entries != null) {
+            int id = registry.networkId(session, object);
+            for (HolderEntry<T> entry : entries) {
+                if (entry instanceof HolderEntry.Direct<T> direct) {
+                    if (Objects.equals(direct.value(), object)) {
+                        return true;
+                    }
+                } else if (entry instanceof HolderEntry.Id<T> idEntry) {
+                    if (id != -1 && idEntry.id() == id) {
+                        return true;
+                    }
+                }
+            }
+            return false;
+        }
         return session.getTagCache().is(this, object);
     }
 
     /**
      * Resolves the HolderSet, and automatically maps the network IDs to their respective object types.
      * If the HolderSet is a list of IDs, this will be returned. If it is a tag, the tag will be resolved from the tag cache. If it is an inline HolderSet, the list of inline elements will be returned.
+     * If it is a heterogeneous HolderSet, the inline elements and resolved registered objects are returned in order.
      *
      * @return the HolderSet turned into a list of objects.
      */
     public List<T> resolve(GeyserSession session) {
         if (inline != null) {
             return inline;
         }
+        if (entries != null) {
+            List<T> result = new ArrayList<>(entries.size());
+            for (HolderEntry<T> entry : entries) {
+                if (entry instanceof HolderEntry.Direct<T> direct) {
+                    result.add(direct.value());
+                } else if (entry instanceof HolderEntry.Id<T> idEntry) {
+                    T value = registry.value(session, idEntry.id());
+                    if (value != null) {
+                        result.add(value);
+                    }
+                }
+            }
+            return result;
+        }
         return TagCache.mapRawArray(session, resolveRaw(session.getTagCache()), registry);
     }
 
     /**
      * Resolves the HolderSet into a list of network IDs. If the HolderSet is a list of IDs, this will be returned. If it is a tag, the tag will be resolved from the tag cache.
      *
-     * <p>If the HolderSet is a list of inline elements, this method will throw! Inline elements are not registered and as such do not have a network ID.</p>
+     * <p>If the HolderSet has inline elements, this method will throw! Inline elements are not registered and as such do not have a network ID.</p>
      *
      * @return the HolderSet turned into a list of network IDs.
-     * @throws IllegalStateException when the HolderSet is a list of inline elements.
+     * @throws IllegalStateException when the HolderSet contains inline elements.
      */
     public IntList resolveRaw(TagCache tagCache) {
-        if (inline != null) {
+        if (inline != null || entries != null) {
             throw new IllegalStateException("Tried to resolve network IDs of a GeyserHolderSet(registry=" + registry  + ") with inline elements!");
         } else if (holders != null) {
             return holders;
@@ -185,7 +237,8 @@ public static <T> GeyserHolderSet<T> readHolderSet(JavaRegistryKey<T> registry,
         // so it works. If this ever changes, we'll have to accommodate for that here
         if (holderSet instanceof NbtMap singleInlineElement && reader != null) {
             return new GeyserHolderSet<>(registry, List.of(reader.apply(singleInlineElement)));
-        } if (holderSet instanceof String elementOrTag) {
+        }
+        if (holderSet instanceof String elementOrTag) {
             if (elementOrTag.startsWith("#")) {
                 // Tag
                 return new GeyserHolderSet<>(registry, new Tag<>(registry, MinecraftKey.key(elementOrTag.substring(1)))); // Remove '#' at beginning that indicates a tag
@@ -196,19 +249,45 @@ public static <T> GeyserHolderSet<T> readHolderSet(JavaRegistryKey<T> registry,
         } else if (holderSet instanceof List<?> list) {
             if (list.isEmpty()) {
                 return new GeyserHolderSet<>(registry);
-            } else if (list.getFirst() instanceof NbtMap) {
+            }
+
+            boolean hasInline = false;
+            boolean hasReference = false;
+            for (Object element : list) {
+                if (element instanceof NbtMap) {
+                    hasInline = true;
+                } else if (element instanceof String) {
+                    hasReference = true;
+                }
+            }
+
+            if (hasInline && !hasReference) {
                 if (reader != null) {
                     return new GeyserHolderSet<>(registry, list.stream().map(o -> (NbtMap) o).map(reader).toList());
                 }
-            } else {
+            } else if (hasReference && !hasInline) {
                 // Assume the list is a list of strings (resource locations)
-                return new GeyserHolderSet<>(registry, IntList.of(list.stream().map(o -> (String) o).map(Key::key).mapToInt(idMapper).toArray()));
+                return new GeyserHolderSet<>(registry, IntList.of(list.stream().map(o -> (String) o).map(MinecraftKey::key).mapToInt(idMapper).toArray()));
+            } else if (hasInline && hasReference) {
+                if (reader != null) {
+                    List<HolderEntry<T>> entries = new ArrayList<>(list.size());
+                    for (Object element : list) {
+                        if (element instanceof NbtMap map) {
+                            entries.add(new HolderEntry.Direct<>(reader.apply(map)));
+                        } else if (element instanceof String str) {
+                            entries.add(new HolderEntry.Id<>(idMapper.applyAsInt(MinecraftKey.key(str))));
+                        } else {
+                            GeyserImpl.getInstance().getLogger().warning("Unexpected element type in mixed HolderSet for registry " + registry + ": " + element);
+                        }
+                    }
+                    return ofEntries(registry, entries);
+                }
             }
         }
 
         String expected = reader == null ? "either a tag, a string ID, or a list of string IDs"
-            : "either a tag, a string ID, an inline registry element, a list of string IDs, or a list of inline registry elements";
-        GeyserImpl.getInstance().getLogger().warning("Failed parsing HolderSet for registry + " + registry + "! Expected " + expected + ", found " + holderSet);
+            : "either a tag, a string ID, an inline registry element, a list of string IDs, a list of inline registry elements, or a heterogeneous list of both";
+        GeyserImpl.getInstance().getLogger().warning("Failed parsing HolderSet for registry " + registry + "! Expected " + expected + ", found " + holderSet);
         return new GeyserHolderSet<>(registry);
     }
 }

core/src/test/java/org/geysermc/geyser/session/cache/tags/GeyserHolderSetTest.java

@@ -0,0 +1,191 @@
+package org.geysermc.geyser.session.cache.tags;
+
+import it.unimi.dsi.fastutil.ints.IntList;
+import net.kyori.adventure.key.Key;
+import org.cloudburstmc.nbt.NbtMap;
+import org.geysermc.geyser.scoreboard.network.util.GeyserMockContext;
+import org.geysermc.geyser.session.GeyserSession;
+import org.geysermc.geyser.session.cache.RegistryCache;
+import org.geysermc.geyser.session.cache.TagCache;
+import org.geysermc.geyser.session.cache.registry.JavaRegistryKey;
+import org.geysermc.geyser.session.cache.registry.JavaRegistryProvider;
+import org.geysermc.geyser.session.cache.registry.RegistryEntryData;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+import java.util.Optional;
+import java.util.function.Function;
+import java.util.function.ToIntFunction;
+
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.Mockito.when;
+
+class GeyserHolderSetTest {
+
+    private void runWithContext(TestRunner runner) {
+        GeyserMockContext.mockContext(context -> {
+            GeyserSession session = context.mock(GeyserSession.class);
+            RegistryCache registryCache = context.mock(RegistryCache.class);
+            TagCache tagCache = context.mock(TagCache.class);
+            when(session.getRegistryCache()).thenReturn(registryCache);
+            when(session.getTagCache()).thenReturn(tagCache);
+
+            JavaRegistryKey.RegistryLookup<String> lookup = new JavaRegistryKey.RegistryLookup<>() {
+                @Override
+                public Optional<RegistryEntryData<String>> entry(JavaRegistryProvider registries, JavaRegistryKey<String> registry, int networkId) {
+                    if (networkId == 100) {
+                        return Optional.of(new RegistryEntryData<>(100, Key.key("minecraft:dialog_ref"), "ResolvedDialog"));
+                    }
+                    return Optional.empty();
+                }
+
+                @Override
+                public Optional<RegistryEntryData<String>> entry(JavaRegistryProvider registries, JavaRegistryKey<String> registry, Key key) {
+                    if (key.equals(Key.key("minecraft:dialog_ref"))) {
+                        return Optional.of(new RegistryEntryData<>(100, Key.key("minecraft:dialog_ref"), "ResolvedDialog"));
+                    }
+                    return Optional.empty();
+                }
+
+                @Override
+                public Optional<RegistryEntryData<String>> entry(JavaRegistryProvider registries, JavaRegistryKey<String> registry, String object) {
+                    if ("ResolvedDialog".equals(object)) {
+                        return Optional.of(new RegistryEntryData<>(100, Key.key("minecraft:dialog_ref"), "ResolvedDialog"));
+                    }
+                    return Optional.empty();
+                }
+            };
+
+            JavaRegistryKey<String> registryKey = new JavaRegistryKey<>(Key.key("minecraft:dialog"), lookup);
+            runner.run(session, registryKey, tagCache);
+        });
+    }
+
+    @FunctionalInterface
+    interface TestRunner {
+        void run(GeyserSession session, JavaRegistryKey<String> registryKey, TagCache tagCache);
+    }
+
+    @Test
+    void testEmptyHolderSet() {
+        runWithContext((session, registryKey, tagCache) -> {
+            GeyserHolderSet<String> set = GeyserHolderSet.readHolderSet(registryKey, null, key -> -1, null);
+            assertNotNull(set);
+            assertNull(set.getTag());
+            assertEquals(0, set.getHolders().size());
+            assertNull(set.getInline());
+            assertNull(set.getEntries());
+            assertTrue(set.resolve(session).isEmpty());
+        });
+    }
+
+    @Test
+    void testTagHolderSet() {
+        runWithContext((session, registryKey, tagCache) -> {
+            GeyserHolderSet<String> set = GeyserHolderSet.readHolderSet(registryKey, "#minecraft:all_dialogs", key -> -1, null);
+            assertNotNull(set.getTag());
+            assertEquals("minecraft", set.getTag().tag().namespace());
+            assertEquals("all_dialogs", set.getTag().tag().value());
+            assertNull(set.getHolders());
+            assertNull(set.getInline());
+            assertNull(set.getEntries());
+        });
+    }
+
+    @Test
+    void testSingleStringHolderSet() {
+        runWithContext((session, registryKey, tagCache) -> {
+            ToIntFunction<Key> idMapper = key -> key.asString().equals("minecraft:dialog_ref") ? 100 : -1;
+            GeyserHolderSet<String> set = GeyserHolderSet.readHolderSet(registryKey, "minecraft:dialog_ref", idMapper, null);
+            assertNull(set.getTag());
+            assertNotNull(set.getHolders());
+            assertEquals(IntList.of(100), set.getHolders());
+            assertNull(set.getInline());
+            assertNull(set.getEntries());
+            assertEquals(List.of("ResolvedDialog"), set.resolve(session));
+        });
+    }
+
+    @Test
+    void testSingleInlineHolderSet() {
+        runWithContext((session, registryKey, tagCache) -> {
+            NbtMap inlineNbt = NbtMap.builder().putString("title", "Inline 1").build();
+            Function<NbtMap, String> reader = map -> map.getString("title");
+
+            GeyserHolderSet<String> set = GeyserHolderSet.readHolderSet(registryKey, inlineNbt, key -> -1, reader);
+            assertNull(set.getTag());
+            assertNull(set.getHolders());
+            assertNotNull(set.getInline());
+            assertEquals(List.of("Inline 1"), set.getInline());
+            assertNull(set.getEntries());
+            assertEquals(List.of("Inline 1"), set.resolve(session));
+            assertThrows(IllegalStateException.class, () -> set.resolveRaw(tagCache));
+        });
+    }
+
+    @Test
+    void testHomogeneousStringListHolderSet() {
+        runWithContext((session, registryKey, tagCache) -> {
+            ToIntFunction<Key> idMapper = key -> 100;
+            GeyserHolderSet<String> set = GeyserHolderSet.readHolderSet(registryKey, List.of("minecraft:dialog_ref"), idMapper, null);
+            assertNotNull(set.getHolders());
+            assertEquals(IntList.of(100), set.getHolders());
+            assertNull(set.getInline());
+            assertNull(set.getEntries());
+            assertEquals(List.of("ResolvedDialog"), set.resolve(session));
+        });
+    }
+
+    @Test
+    void testHomogeneousInlineListHolderSet() {
+        runWithContext((session, registryKey, tagCache) -> {
+            NbtMap nbt1 = NbtMap.builder().putString("title", "Inline 1").build();
+            NbtMap nbt2 = NbtMap.builder().putString("title", "Inline 2").build();
+            Function<NbtMap, String> reader = map -> map.getString("title");
+
+            GeyserHolderSet<String> set = GeyserHolderSet.readHolderSet(registryKey, List.of(nbt1, nbt2), key -> -1, reader);
+            assertNull(set.getHolders());
+            assertNotNull(set.getInline());
+            assertEquals(List.of("Inline 1", "Inline 2"), set.getInline());
+            assertNull(set.getEntries());
+            assertEquals(List.of("Inline 1", "Inline 2"), set.resolve(session));
+        });
+    }
+
+    @Test
+    void testHeterogeneousListHolderSet() {
+        runWithContext((session, registryKey, tagCache) -> {
+            NbtMap inline1 = NbtMap.builder().putString("title", "Inline First").build();
+            String ref = "minecraft:dialog_ref";
+            NbtMap inline2 = NbtMap.builder().putString("title", "Inline Last").build();
+
+            ToIntFunction<Key> idMapper = key -> key.asString().equals("minecraft:dialog_ref") ? 100 : -1;
+            Function<NbtMap, String> reader = map -> map.getString("title");
+
+            List<Object> mixedList = List.of(inline1, ref, inline2);
+            GeyserHolderSet<String> set = GeyserHolderSet.readHolderSet(registryKey, mixedList, idMapper, reader);
+
+            assertNull(set.getHolders());
+            assertNull(set.getInline());
+            assertNotNull(set.getEntries());
+            assertEquals(3, set.getEntries().size());
+
+            assertInstanceOf(GeyserHolderSet.HolderEntry.Direct.class, set.getEntries().get(0));
+            assertInstanceOf(GeyserHolderSet.HolderEntry.Id.class, set.getEntries().get(1));
+            assertInstanceOf(GeyserHolderSet.HolderEntry.Direct.class, set.getEntries().get(2));
+
+            // Test resolution preserves exact sequence
+            List<String> resolved = set.resolve(session);
+            assertEquals(List.of("Inline First", "ResolvedDialog", "Inline Last"), resolved);
+
+            // Test resolveRaw throws
+            assertThrows(IllegalStateException.class, () -> set.resolveRaw(tagCache));
+
+            // Test contains
+            assertTrue(set.contains(session, "Inline First"));
+            assertTrue(set.contains(session, "ResolvedDialog"));
+            assertTrue(set.contains(session, "Inline Last"));
+            assertFalse(set.contains(session, "NonExistentDialog"));
+        });
+    }
+}