Fix SearchFilteringRule

2 files changed, 69 insertions, 4 deletions

diff --git a/api/src/main/java/me/shedaniel/rei/api/common/util/CollectionUtils.java b/api/src/main/java/me/shedaniel/rei/api/common/util/CollectionUtils.java
index c80e0dc25..67894fb27 100644
--- a/api/src/main/java/me/shedaniel/rei/api/common/util/CollectionUtils.java
+++ b/api/src/main/java/me/shedaniel/rei/api/common/util/CollectionUtils.java
@@ -23,6 +23,7 @@
 
 package me.shedaniel.rei.api.common.util;
 
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Sets;
 import com.google.common.collect.UnmodifiableIterator;
@@ -361,7 +362,7 @@ public class CollectionUtils {
     }
     
     public static <T> Iterable<List<T>> partition(List<T> list, int size) {
-        return () -> new UnmodifiableIterator<List<T>>() {
+        return () -> new UnmodifiableIterator<>() {
             int i = 0;
             int partitionSize = Mth.ceil(list.size() / (float) size);
             
@@ -374,7 +375,7 @@ public class CollectionUtils {
             public List<T> next() {
                 int cursor = i++ * size;
                 int realSize = Math.min(list.size() - cursor, size);
-                return new AbstractList<T>() {
+                return new AbstractList<>() {
                     @Override
                     public T get(int index) {
                         if (index < 0 || index >= realSize)
@@ -414,6 +415,70 @@ public class CollectionUtils {
         };
     }
     
+    public static <T> Iterable<Iterator<T>> partitionIterator(Iterator<T> iterator, int iteratorSize, int size) {
+        return partitionCollection(new AbstractCollection<>() {
+            
+            @Override
+            public Iterator<T> iterator() {
+                return iterator;
+            }
+            
+            @Override
+            public int size() {
+                return iteratorSize;
+            }
+        }, size);
+    }
+    
+    public static <T> Iterable<Iterator<T>> partitionCollection(Collection<T> collection, int size) {
+        if (collection instanceof List) {
+            return Iterables.transform(partition((List<T>) collection, size), List::iterator);
+        }
+        
+        return () -> new Iterator<>() {
+            int i = 0;
+            int partitionSize = Mth.ceil(collection.size() / (float) size);
+            int advanced = 0;
+            Iterator<T> iterator = collection.iterator();
+            
+            @Override
+            public boolean hasNext() {
+                return i < partitionSize;
+            }
+            
+            @Override
+            public Iterator<T> next() {
+                int cursor = i++ * size;
+                int realSize = Math.min(collection.size() - cursor, size);
+                
+                if (advanced < cursor) {
+                    for (int j = 0; j < cursor - advanced; j++) {
+                        if (iterator.hasNext()) {
+                            iterator.next();
+                        } else {
+                            advanced = cursor;
+                            return Collections.emptyIterator();
+                        }
+                    }
+                    advanced = cursor;
+                }
+                
+                return new Iterator<>() {
+                    @Override
+                    public boolean hasNext() {
+                        return iterator.hasNext() && advanced < cursor + realSize;
+                    }
+                    
+                    @Override
+                    public T next() {
+                        advanced++;
+                        return iterator.next();
+                    }
+                };
+            }
+        };
+    }
+    
     public static Ingredient toIngredient(ItemStack stack) {
         return Ingredient.of(Stream.of(stack));
     }
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/entry/filtering/rules/SearchFilteringRule.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/entry/filtering/rules/SearchFilteringRule.java
index e3bf9c790..e53330953 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/entry/filtering/rules/SearchFilteringRule.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/entry/filtering/rules/SearchFilteringRule.java
@@ -24,12 +24,12 @@
 package me.shedaniel.rei.impl.client.entry.filtering.rules;
 
 import com.google.common.base.Suppliers;
+import com.google.common.collect.Iterators;
 import com.google.common.collect.Lists;
 import me.shedaniel.rei.api.client.entry.filtering.*;
 import me.shedaniel.rei.api.client.search.SearchFilter;
 import me.shedaniel.rei.api.client.search.SearchProvider;
 import me.shedaniel.rei.api.common.entry.EntryStack;
-import me.shedaniel.rei.api.common.util.CollectionUtils;
 import net.fabricmc.api.EnvType;
 import net.fabricmc.api.Environment;
 import net.minecraft.util.StringUtil;
@@ -93,7 +93,7 @@ public class SearchFilteringRule implements FilteringRule<Unit> {
     }
     
     private void processList(Collection<EntryStack<?>> stacks, List<CompletableFuture<List<EntryStack<?>>>> completableFutures) {
-        for (Iterable<EntryStack<?>> partitionStacks : CollectionUtils.partition((List<EntryStack<?>>) stacks, 100)) {
+        for (Iterable<EntryStack<?>> partitionStacks : (Iterable<List<EntryStack<?>>>) () -> Iterators.partition(stacks.iterator(), 100)) {
             completableFutures.add(CompletableFuture.supplyAsync(() -> {
                 List<EntryStack<?>> output = Lists.newArrayList();
                 for (EntryStack<?> stack : partitionStacks) {