1 files changed, 65 insertions, 13 deletions

diff --git a/src/main/java/me/shedaniel/rei/impl/EntryRegistryImpl.java b/src/main/java/me/shedaniel/rei/impl/EntryRegistryImpl.java
index b9053e547..7fd4f8d1b 100644
--- a/src/main/java/me/shedaniel/rei/impl/EntryRegistryImpl.java
+++ b/src/main/java/me/shedaniel/rei/impl/EntryRegistryImpl.java
@@ -13,16 +13,39 @@ import net.minecraft.item.Item;
 import net.minecraft.item.ItemStack;
 import net.minecraft.util.DefaultedList;
 
+import it.unimi.dsi.fastutil.Hash;
+import it.unimi.dsi.fastutil.objects.ObjectOpenCustomHashSet;
+import it.unimi.dsi.fastutil.objects.ObjectSets;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
 import java.util.List;
-import java.util.TreeSet;
+import java.util.Set;
 import java.util.concurrent.CopyOnWriteArrayList;
-import java.util.stream.Collectors;
 
 @Deprecated
 @Internal
 public class EntryRegistryImpl implements EntryRegistry {
 
     private final CopyOnWriteArrayList<EntryStack> entries = Lists.newCopyOnWriteArrayList();
+    private final Set<EntryStack> entrySet = ObjectSets.synchronize(new ObjectOpenCustomHashSet<>(new Hash.Strategy<EntryStack>() {
+        @Override
+        public int hashCode(EntryStack entry) {
+            return entry == null ? 0 : entry.hashOfAll();
+        }
+
+        @Override
+        public boolean equals(EntryStack a, EntryStack b) {
+            if (a == null || b == null) {
+                return a == b;
+            } else {
+                boolean result = a.equalsAll(b);
+                assert result == b.equalsAll(a) : "a.equalsAll(b) != b.equalsAll(a); (a = " + a + ", b = " + b + ")";
+                return result;
+            }
+        }
+    }));
 
     @Override
     public List<EntryStack> getStacksList() {
@@ -34,23 +57,52 @@ public class EntryRegistryImpl implements EntryRegistry {
         DefaultedList<ItemStack> list = DefaultedList.of();
         list.add(item.getStackForRender());
         item.appendStacks(item.getGroup(), list);
-        TreeSet<ItemStack> stackSet = list.stream().collect(Collectors.toCollection(() -> new TreeSet<ItemStack>((p1, p2) -> ItemStack.areEqualIgnoreDamage(p1, p2) ? 0 : 1)));
-        return Lists.newArrayList(stackSet).toArray(new ItemStack[0]);
+        ItemStack[] array = list.toArray(new ItemStack[0]);
+        Arrays.sort(array, (a, b) -> ItemStack.areEqualIgnoreDamage(a, b) ? 0 : 1);
+        return array;
     }
 
     @Override
     @Deprecated
     public void registerEntryAfter(EntryStack afterEntry, EntryStack stack, boolean checkAlreadyContains) {
-        if (!stack.isEmpty() && (!checkAlreadyContains || !alreadyContain(stack)))
-            if (afterEntry == null || afterEntry.isEmpty())
-                entries.add(stack);
-            else {
-                int last = entries.size();
-                for (int i = 0; i < entries.size(); i++)
-                    if (entries.get(i).equalsAll(afterEntry))
-                        last = i + 1;
-                entries.add(last, stack);
+        if (stack.isEmpty()) return;
+        boolean isNew = entrySet.add(stack);
+        if (checkAlreadyContains && !isNew) {
+            return;
+        }
+        if (afterEntry == null) {
+            entries.add(stack);
+        } else {
+            int last = entries.size();
+            for (int i = 0; i < entries.size(); i++)
+                if (entries.get(i).equalsAll(afterEntry))
+                    last = i + 1;
+            entries.add(last, stack);
+        }
+    }
+
+    @Override
+    public void registerEntriesAfter(EntryStack afterStack, Collection<? extends EntryStack> stacks) {
+        List<EntryStack> nonDuplicates = new ArrayList<>();
+        for (EntryStack stack : stacks) {
+            if (entrySet.add(stack)) {
+                nonDuplicates.add(stack);
+            }
+        }
+        int index = entries.size();
+        if (afterStack != null) {
+            for (int i = index - 1; i >= 0; i--) {
+                if (entries.get(i).equalsAll(afterStack)) {
+                    index = i + 1;
+                    break;
+                }
             }
+        }
+        entries.addAll(index, nonDuplicates);
     }
 
+    @Override
+    public void registerEntriesAfter(EntryStack afterStack, EntryStack... stacks) {
+        registerEntriesAfter(afterStack, Arrays.asList(stacks));
+    }
 }