[1.20.6] Use stack hash strategy for mutable map in creative tab event

src/main/java/net/neoforged/neoforge/common/util/MutableHashedLinkedMap.java

@@ -12,6 +12,7 @@
 import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Objects;
+import java.util.function.BiPredicate;
 import org.jetbrains.annotations.Nullable;
 
 /**
@@ -34,6 +35,7 @@ public class MutableHashedLinkedMap<K, V> implements Iterable<Map.Entry<K, V>> {
     private final Strategy<? super K> strategy;
     private final Map<K, Entry> entries;
     private final MergeFunction<K, V> merge;
+    private final BiPredicate<K, V> insertTest;
     private Entry head = null;
     private Entry last = null;
     /*
@@ -59,15 +61,37 @@ public MutableHashedLinkedMap(Strategy<? super K> strategy) {
     }
 
     /**
-     * Creates a mutable linked map with a custom merge function.
+     * Creates a mutable linked map with a default new-value-selecting merge function.
      *
      * @param strategy the hashing strategy
      * @param merge    the function used when merging an existing value and a new value
      */
     public MutableHashedLinkedMap(Strategy<? super K> strategy, MergeFunction<K, V> merge) {
+        this(strategy, merge, (k, v) -> true);
+    }
+
+    /**
+     * Creates a mutable linked map with a default new-value-selecting merge function.
+     *
+     * @param strategy   the hashing strategy
+     * @param insertTest the test to apply before inserting a key and value
+     */
+    public MutableHashedLinkedMap(Strategy<? super K> strategy, BiPredicate<K, V> insertTest) {
+        this(strategy, (k, v1, v2) -> v2, insertTest);
+    }
+
+    /**
+     * Creates a mutable linked map with a custom merge function.
+     *
+     * @param strategy   the hashing strategy
+     * @param merge      the function used when merging an existing value and a new value
+     * @param insertTest the test to apply before inserting a key and value
+     */
+    public MutableHashedLinkedMap(Strategy<? super K> strategy, MergeFunction<K, V> merge, BiPredicate<K, V> insertTest) {
         this.strategy = strategy;
         this.entries = new Object2ObjectOpenCustomHashMap<>(strategy);
         this.merge = merge;
+        this.insertTest = insertTest;
     }
 
     /**
@@ -84,6 +108,9 @@ public MutableHashedLinkedMap(Strategy<? super K> strategy, MergeFunction<K, V>
     @Nullable
     public V put(K key, V value) {
         var old = entries.get(key);
+        if (!insertTest.test(key, value)) {
+            return old == null ? null : old.value;
+        }
         if (old != null) {
             V ret = old.value;
             old.value = merge.apply(key, ret, value);
@@ -215,6 +242,9 @@ public V putAfter(K after, K key, V value) {
 
         V ret = null;
         var entry = entries.get(key);
+        if (!insertTest.test(key, value)) {
+            return entry == null ? null : entry.value;
+        }
         if (entry != null) {
             ret = entry.value;
             entry.value = merge.apply(key, ret, value);
@@ -260,6 +290,9 @@ public V putBefore(K before, K key, V value) {
 
         V ret = null;
         var entry = entries.get(key);
+        if (!insertTest.test(key, value)) {
+            return entry == null ? null : entry.value;
+        }
         if (entry != null) {
             ret = entry.value;
             entry.value = merge.apply(key, ret, value);

src/main/java/net/neoforged/neoforge/event/EventHooks.java

@@ -8,6 +8,7 @@
 import com.mojang.authlib.GameProfile;
 import com.mojang.brigadier.CommandDispatcher;
 import com.mojang.datafixers.util.Either;
+import it.unimi.dsi.fastutil.objects.ObjectLinkedOpenCustomHashSet;
 import java.io.File;
 import java.util.EnumSet;
 import java.util.List;
@@ -69,6 +70,7 @@
 import net.minecraft.world.item.CreativeModeTab;
 import net.minecraft.world.item.Item;
 import net.minecraft.world.item.ItemStack;
+import net.minecraft.world.item.ItemStackLinkedSet;
 import net.minecraft.world.item.TooltipFlag;
 import net.minecraft.world.item.context.UseOnContext;
 import net.minecraft.world.item.crafting.RecipeType;
@@ -1035,14 +1037,30 @@ public static ItemEnchantments getEnchantmentLevel(ItemEnchantments enchantments
      */
     @ApiStatus.Internal
     public static void onCreativeModeTabBuildContents(CreativeModeTab tab, ResourceKey<CreativeModeTab> tabKey, CreativeModeTab.DisplayItemsGenerator originalGenerator, CreativeModeTab.ItemDisplayParameters params, CreativeModeTab.Output output) {
-        final var entries = new MutableHashedLinkedMap<ItemStack, CreativeModeTab.TabVisibility>();
+        final var searchDupes = new ObjectLinkedOpenCustomHashSet<ItemStack>(ItemStackLinkedSet.TYPE_AND_TAG);
+        // The ItemStackLinkedSet.TYPE_AND_TAG strategy cannot be used for the MutableHashedLinkedMap due to vanilla
+        // adding multiple identical ItemStacks with different TabVisibility values. The values also cannot be merged
+        // because it does not abide by the intended order. For example, vanilla adds all max enchanted books to the
+        // "ingredient" tab with "parent only" visibility, then also adds all enchanted books again in increasing order
+        // to their max values but with the "search only" visibility. Because the parent-only is added first and then
+        // the search-only entries are added after, the max enchantments would show up first and then the enchantments
+        // in increasing order up to max-1.
+        final var entries = new MutableHashedLinkedMap<ItemStack, CreativeModeTab.TabVisibility>(MutableHashedLinkedMap.BASIC, (stack, tabVisibility) -> {
+            if (!searchDupes.add(stack) && tabVisibility != CreativeModeTab.TabVisibility.SEARCH_TAB_ONLY) {
+                throw new IllegalStateException(
+                        "Accidentally adding the same item stack twice "
+                                + stack.getDisplayName().getString()
+                                + " to a Creative Mode Tab: "
+                                + tab.getDisplayName().getString());
+            }
+            return true;
+        });
 
         originalGenerator.accept(params, (stack, vis) -> {
             if (stack.getCount() != 1)
                 throw new IllegalArgumentException("The stack count must be 1");
             entries.put(stack, vis);
         });
-
         ModLoader.postEvent(new BuildCreativeModeTabContentsEvent(tab, tabKey, params, entries));
 
         for (var entry : entries)

tests/src/junit/java/net/neoforged/neoforge/unittest/CreativeTabOrderTest.java

@@ -0,0 +1,151 @@
+/*
+ * Copyright (c) NeoForged and contributors
+ * SPDX-License-Identifier: LGPL-2.1-only
+ */
+
+package net.neoforged.neoforge.unittest;
+
+import it.unimi.dsi.fastutil.objects.ObjectOpenCustomHashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.IntStream;
+import net.minecraft.core.Holder;
+import net.minecraft.core.HolderLookup;
+import net.minecraft.core.component.DataComponents;
+import net.minecraft.core.registries.Registries;
+import net.minecraft.server.MinecraftServer;
+import net.minecraft.tags.ItemTags;
+import net.minecraft.tags.TagKey;
+import net.minecraft.world.flag.FeatureFlagSet;
+import net.minecraft.world.flag.FeatureFlags;
+import net.minecraft.world.item.CreativeModeTab;
+import net.minecraft.world.item.CreativeModeTabs;
+import net.minecraft.world.item.EnchantedBookItem;
+import net.minecraft.world.item.Item;
+import net.minecraft.world.item.ItemStack;
+import net.minecraft.world.item.ItemStackLinkedSet;
+import net.minecraft.world.item.Items;
+import net.minecraft.world.item.enchantment.Enchantment;
+import net.minecraft.world.item.enchantment.EnchantmentInstance;
+import net.neoforged.bus.api.IEventBus;
+import net.neoforged.fml.common.Mod;
+import net.neoforged.neoforge.event.BuildCreativeModeTabContentsEvent;
+import net.neoforged.testframework.junit.EphemeralTestServerProvider;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.MethodOrderer;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestMethodOrder;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+@ExtendWith(EphemeralTestServerProvider.class)
+@TestMethodOrder(MethodOrderer.MethodName.class)
+public class CreativeTabOrderTest {
+    public static final String MOD_ID = "creative_tab_order_test";
+    private static final Set<TagKey<Item>> ENCHANTABLES = Set.of(
+            ItemTags.FOOT_ARMOR_ENCHANTABLE,
+            ItemTags.LEG_ARMOR_ENCHANTABLE,
+            ItemTags.CHEST_ARMOR_ENCHANTABLE,
+            ItemTags.HEAD_ARMOR_ENCHANTABLE,
+            ItemTags.ARMOR_ENCHANTABLE,
+            ItemTags.SWORD_ENCHANTABLE,
+            ItemTags.SHARP_WEAPON_ENCHANTABLE,
+            ItemTags.MACE_ENCHANTABLE,
+            ItemTags.FIRE_ASPECT_ENCHANTABLE,
+            ItemTags.WEAPON_ENCHANTABLE,
+            ItemTags.MINING_ENCHANTABLE,
+            ItemTags.MINING_LOOT_ENCHANTABLE,
+            ItemTags.FISHING_ENCHANTABLE,
+            ItemTags.TRIDENT_ENCHANTABLE,
+            ItemTags.DURABILITY_ENCHANTABLE,
+            ItemTags.BOW_ENCHANTABLE,
+            ItemTags.EQUIPPABLE_ENCHANTABLE,
+            ItemTags.CROSSBOW_ENCHANTABLE,
+            ItemTags.VANISHING_ENCHANTABLE);
+    public static Iterable<Map.Entry<ItemStack, CreativeModeTab.TabVisibility>> ingredientsTab;
+    public static Iterable<Map.Entry<ItemStack, CreativeModeTab.TabVisibility>> searchTab;
+
+    @BeforeAll
+    static void testSetupTabs(MinecraftServer server) {
+        CreativeModeTabs.tryRebuildTabContents(FeatureFlags.DEFAULT_FLAGS, true, server.registryAccess());
+    }
+
+    /**
+     * The local tabEnchantments variable comes from {@link CreativeModeTabs#generateEnchantmentBookTypesOnlyMaxLevel(CreativeModeTab.Output, HolderLookup, Set, CreativeModeTab.TabVisibility, FeatureFlagSet)}
+     * 
+     * @param server Ephemeral server from extension
+     */
+    @Test
+    void testIngredientsEnchantmentExistence(MinecraftServer server) {
+        final Set<ItemStack> tabEnchantments = server.registryAccess().lookupOrThrow(Registries.ENCHANTMENT).listElements()
+                .map(Holder::value)
+                .filter(enchantment -> enchantment.isEnabled(FeatureFlags.DEFAULT_FLAGS))
+                .filter(enchantment -> enchantment.allowedInCreativeTab(Items.ENCHANTED_BOOK, ENCHANTABLES))
+                .map(enchantment -> EnchantedBookItem.createForEnchantment(new EnchantmentInstance(enchantment, enchantment.getMaxLevel())))
+                .collect(() -> new ObjectOpenCustomHashSet<>(ItemStackLinkedSet.TYPE_AND_TAG), ObjectOpenCustomHashSet::add, ObjectOpenCustomHashSet::addAll);
+        for (Map.Entry<ItemStack, CreativeModeTab.TabVisibility> entry : ingredientsTab) {
+            if (entry.getValue() == CreativeModeTab.TabVisibility.SEARCH_TAB_ONLY) {
+                continue;
+            }
+            if (entry.getKey().getItem() == Items.ENCHANTED_BOOK) {
+                Assertions.assertTrue(tabEnchantments.remove(entry.getKey()), "Enchanted book present that does not exist in the default set?");
+            }
+        }
+
+        Assertions.assertTrue(tabEnchantments.isEmpty(), "Missing enchantments in Ingredient tab.");
+    }
+
+    /**
+     * The local tabEnchantments variable comes from {@link CreativeModeTabs#generateEnchantmentBookTypesAllLevels(CreativeModeTab.Output, HolderLookup, Set, CreativeModeTab.TabVisibility, FeatureFlagSet)}
+     * 
+     * @param server Ephemeral server from extension
+     */
+    @Test
+    void testSearchEnchantmentOrder(MinecraftServer server) {
+        final var tabEnchantments = server.registryAccess().lookupOrThrow(Registries.ENCHANTMENT).listElements()
+                .map(Holder::value)
+                .filter(enchantment -> enchantment.isEnabled(FeatureFlags.DEFAULT_FLAGS))
+                .filter(enchantment -> enchantment.allowedInCreativeTab(Items.ENCHANTED_BOOK, ENCHANTABLES))
+                .flatMap(
+                        enchantment -> IntStream.rangeClosed(enchantment.getMinLevel(), enchantment.getMaxLevel())
+                                .mapToObj(p_270006_ -> EnchantedBookItem.createForEnchantment(new EnchantmentInstance(enchantment, p_270006_))))
+                .collect(() -> new ObjectOpenCustomHashSet<>(ItemStackLinkedSet.TYPE_AND_TAG), ObjectOpenCustomHashSet::add, ObjectOpenCustomHashSet::addAll);
+
+        Enchantment enchantment = null;
+        int level = 0;
+        for (Map.Entry<ItemStack, CreativeModeTab.TabVisibility> entry : searchTab) {
+            if (entry.getKey().getItem() != Items.ENCHANTED_BOOK) {
+                continue;
+            }
+            final var enchantmentEntry = entry.getKey().get(DataComponents.STORED_ENCHANTMENTS).entrySet().iterator().next();
+            final var entryEnchantment = enchantmentEntry.getKey().value();
+            final var entryEnchantmentLevel = enchantmentEntry.getIntValue();
+            if (enchantment == null || enchantment != entryEnchantment) {
+                enchantment = entryEnchantment;
+                Assertions.assertFalse(entryEnchantmentLevel > enchantment.getMinLevel(), "Enchantment does not start at the minimum level");
+            } else {
+                Assertions.assertTrue(entryEnchantmentLevel > level);
+            }
+            Assertions.assertTrue(tabEnchantments.remove(entry.getKey()), "Enchanted book present that does not exist in the default set?");
+            level = entryEnchantmentLevel;
+        }
+
+        Assertions.assertTrue(tabEnchantments.isEmpty(), "Missing enchantments in Search tab.");
+    }
+
+    @Mod(MOD_ID)
+    public static class CreativeTabOrderTestMod {
+        public CreativeTabOrderTestMod(IEventBus modBus) {
+            modBus.addListener(this::buildCreativeTab);
+        }
+
+        private void buildCreativeTab(final BuildCreativeModeTabContentsEvent event) {
+            if (event.getTabKey() == CreativeModeTabs.INGREDIENTS) {
+                ingredientsTab = event.getEntries();
+            }
+            if (event.getTabKey() == CreativeModeTabs.SEARCH) {
+                searchTab = event.getEntries();
+            }
+        }
+    }
+}

tests/src/main/java/net/neoforged/neoforge/oldtest/CreativeModeTabTest.java

@@ -6,6 +6,8 @@
 package net.neoforged.neoforge.oldtest;
 
 import java.util.List;
+import net.minecraft.core.component.DataComponentPatch;
+import net.minecraft.core.component.DataComponents;
 import net.minecraft.core.registries.Registries;
 import net.minecraft.network.chat.Component;
 import net.minecraft.resources.ResourceKey;
@@ -16,6 +18,7 @@
 import net.minecraft.world.item.DyeColor;
 import net.minecraft.world.item.DyeItem;
 import net.minecraft.world.item.ItemStack;
+import net.minecraft.world.item.Items;
 import net.minecraft.world.level.ItemLike;
 import net.minecraft.world.level.block.Block;
 import net.minecraft.world.level.block.Blocks;
@@ -31,6 +34,7 @@ public class CreativeModeTabTest {
 
     private static final ResourceKey<CreativeModeTab> LOGS = ResourceKey.create(Registries.CREATIVE_MODE_TAB, new ResourceLocation(MOD_ID, "logs"));
     private static final ResourceKey<CreativeModeTab> STONE = ResourceKey.create(Registries.CREATIVE_MODE_TAB, new ResourceLocation(MOD_ID, "stone"));
+    private static final ResourceKey<CreativeModeTab> DAMAGED_SWORDS = ResourceKey.create(Registries.CREATIVE_MODE_TAB, new ResourceLocation(MOD_ID, "damaged_swords"));
 
     public CreativeModeTabTest(IEventBus modEventBus) {
         if (!ENABLED)
@@ -76,6 +80,18 @@ private static void onCreativeModeTabRegister(RegisterEvent event) {
                     .withTabsBefore(CreativeModeTabs.COLORED_BLOCKS)
                     .build());
 
+            helper.register(DAMAGED_SWORDS, CreativeModeTab.builder().title(Component.literal("Damaged Wooden Swords"))
+                    .displayItems((params, output) -> {
+                        output.accept(new ItemStack(Items.WOODEN_SWORD));
+                        output.accept(new ItemStack(Items.WOODEN_SWORD), TabVisibility.SEARCH_TAB_ONLY); // Should still be added
+                        for (int i = 1; i <= 59; i++) {
+                            // Each should be added since they have different data component values
+                            output.accept(new ItemStack(Items.WOODEN_SWORD.builtInRegistryHolder(), 1, DataComponentPatch.builder().set(DataComponents.DAMAGE, i).build()));
+                        }
+                    })
+                    .icon(() -> new ItemStack(Items.WOODEN_SWORD))
+                    .build());
+
             List<Block> blocks = List.of(Blocks.GRANITE, Blocks.DIORITE, Blocks.ANDESITE, Blocks.COBBLESTONE);
             for (int i = 0; i < blocks.size(); i++) {
                 Block block = blocks.get(i);
@@ -118,6 +134,11 @@ private static void onCreativeModeTabBuildContents(BuildCreativeModeTabContentsE
             entries.putBefore(i(Blocks.DIORITE), i(Blocks.POLISHED_DIORITE), vis);
             entries.putBefore(i(Blocks.ANDESITE), i(Blocks.POLISHED_ANDESITE), vis);
         }
+
+        // Adding this causes a crash (as it should) when opening the creative inventory
+//        if (event.getTabKey() == DAMAGED_SWORDS) {
+//            entries.putBefore(i(Items.WOODEN_SWORD), i(Items.WOODEN_SWORD), vis);
+//        }
     }
 
     private static class CreativeModeColorTab extends CreativeModeTab {

tests/src/main/resources/META-INF/neoforge.mods.toml

@@ -37,6 +37,9 @@ license="LGPL v2.1"
 [[mods]]
     modId="multiple_entrypoints_test"
 
+[[mods]]
+    modId="creative_tab_order_test"
+
 [[mods]]
     modId="ordered_test_1"
 [[mods]]