Added workstation usage info when looking for usage of item/block

1 files changed, 14 insertions, 10 deletions

diff --git a/src/main/java/me/shedaniel/rei/impl/RecipeHelperImpl.java b/src/main/java/me/shedaniel/rei/impl/RecipeHelperImpl.java
index 7e5846bc3..21ea623a0 100644
--- a/src/main/java/me/shedaniel/rei/impl/RecipeHelperImpl.java
+++ b/src/main/java/me/shedaniel/rei/impl/RecipeHelperImpl.java
@@ -170,19 +170,23 @@ public class RecipeHelperImpl implements RecipeHelper {
         Map<Identifier, List<RecipeDisplay>> categoriesMap = new HashMap<>();
         categories.forEach(f -> categoriesMap.put(f.getIdentifier(), Lists.newArrayList()));
         for (Map.Entry<Identifier, List<RecipeDisplay>> entry : recipeCategoryListMap.entrySet()) {
-            RecipeCategory<?> category = getCategory(entry.getKey());
+            boolean isWorkstationCategory = getWorkingStations(entry.getKey()).stream().anyMatch(ws -> ws.contains(stack));
             for (RecipeDisplay recipeDisplay : entry.getValue()) {
-                boolean found = false;
-                for (List<EntryStack> input : recipeDisplay.getInputEntries()) {
-                    for (EntryStack otherEntry : input) {
-                        if (otherEntry.equals(stack)) {
-                            categoriesMap.get(recipeDisplay.getRecipeCategory()).add(recipeDisplay);
-                            found = true;
-                            break;
+                if (isWorkstationCategory) {
+                    categoriesMap.get(recipeDisplay.getRecipeCategory()).add(recipeDisplay);
+                } else {
+                    boolean found = false;
+                    for (List<EntryStack> input : recipeDisplay.getInputEntries()) {
+                        for (EntryStack otherEntry : input) {
+                            if (otherEntry.equals(stack)) {
+                                categoriesMap.get(recipeDisplay.getRecipeCategory()).add(recipeDisplay);
+                                found = true;
+                                break;
+                            }
                         }
+                        if (found)
+                            break;
                     }
-                    if (found)
-                        break;
                 }
             }
         }