1 files changed, 27 insertions, 14 deletions

diff --git a/src/main/java/me/shedaniel/rei/impl/ClientHelperImpl.java b/src/main/java/me/shedaniel/rei/impl/ClientHelperImpl.java
index 10e91151e..6c83db607 100644
--- a/src/main/java/me/shedaniel/rei/impl/ClientHelperImpl.java
+++ b/src/main/java/me/shedaniel/rei/impl/ClientHelperImpl.java
@@ -32,8 +32,8 @@ import me.shedaniel.rei.RoughlyEnoughItemsCore;
 import me.shedaniel.rei.RoughlyEnoughItemsNetwork;
 import me.shedaniel.rei.api.*;
 import me.shedaniel.rei.gui.PreRecipeViewingScreen;
+import me.shedaniel.rei.gui.RecipeScreen;
 import me.shedaniel.rei.gui.RecipeViewingScreen;
-import me.shedaniel.rei.gui.StackToNoticeScreen;
 import me.shedaniel.rei.gui.VillagerRecipeViewingScreen;
 import me.shedaniel.rei.gui.config.RecipeScreenType;
 import me.shedaniel.rei.utils.CollectionUtils;
@@ -68,7 +68,7 @@ import java.util.UUID;
 @ApiStatus.Internal
 public class ClientHelperImpl implements ClientHelper, ClientModInitializer {
     
-    public static ClientHelperImpl instance;
+    private static ClientHelperImpl instance;
     @ApiStatus.Internal public final Lazy<Boolean> isYog = new Lazy<>(() -> {
         try {
             if (MinecraftClient.getInstance().getSession().getProfile().getId().equals(UUID.fromString("f9546389-9415-4358-9c29-2c26b25bff5b")))
@@ -89,6 +89,11 @@ public class ClientHelperImpl implements ClientHelper, ClientModInitializer {
     });
     private final Map<String, String> modNameCache = Maps.newHashMap();
     
+    @ApiStatus.Internal
+    public static ClientHelperImpl getInstance() {
+        return instance;
+    }
+    
     @Override
     public String getFormattedModFromItem(Item item) {
         String mod = getModFromItem(item);
@@ -176,7 +181,7 @@ public class ClientHelperImpl implements ClientHelper, ClientModInitializer {
     public boolean executeRecipeKeyBind(EntryStack stack) {
         Map<RecipeCategory<?>, List<RecipeDisplay>> map = RecipeHelper.getInstance().getRecipesFor(stack);
         if (map.keySet().size() > 0)
-            openRecipeViewingScreen(map, null, stack);
+            openRecipeViewingScreen(map, null, null, stack);
         return map.keySet().size() > 0;
     }
     
@@ -184,7 +189,7 @@ public class ClientHelperImpl implements ClientHelper, ClientModInitializer {
     public boolean executeUsageKeyBind(EntryStack stack) {
         Map<RecipeCategory<?>, List<RecipeDisplay>> map = RecipeHelper.getInstance().getUsagesFor(stack);
         if (map.keySet().size() > 0)
-            openRecipeViewingScreen(map, stack, null);
+            openRecipeViewingScreen(map, null, stack, null);
         return map.keySet().size() > 0;
     }
     
@@ -200,7 +205,7 @@ public class ClientHelperImpl implements ClientHelper, ClientModInitializer {
     public boolean executeViewAllRecipesKeyBind() {
         Map<RecipeCategory<?>, List<RecipeDisplay>> map = RecipeHelper.getInstance().getAllRecipes();
         if (map.keySet().size() > 0)
-            openRecipeViewingScreen(map);
+            openRecipeViewingScreen(map, null, null, null);
         return map.keySet().size() > 0;
     }
     
@@ -232,30 +237,38 @@ public class ClientHelperImpl implements ClientHelper, ClientModInitializer {
     
     @Override
     public void openRecipeViewingScreen(Map<RecipeCategory<?>, List<RecipeDisplay>> map) {
-        openRecipeViewingScreen(map, null, null);
+        openRecipeViewingScreen(map, null, null, null);
     }
     
     @ApiStatus.Internal
-    public void openRecipeViewingScreen(Map<RecipeCategory<?>, List<RecipeDisplay>> map, @Nullable EntryStack ingredientNotice, @Nullable EntryStack resultNotice) {
+    public void openRecipeViewingScreen(Map<RecipeCategory<?>, List<RecipeDisplay>> map, @Nullable Identifier category, @Nullable EntryStack ingredientNotice, @Nullable EntryStack resultNotice) {
+        if (category == null) {
+            Screen currentScreen = MinecraftClient.getInstance().currentScreen;
+            if (currentScreen instanceof RecipeScreen) {
+                category = ((RecipeScreen) currentScreen).getCurrentCategory();
+            }
+        }
         Screen screen;
         if (ConfigObject.getInstance().getRecipeScreenType() == RecipeScreenType.VILLAGER) {
-            screen = new VillagerRecipeViewingScreen(map);
+            screen = new VillagerRecipeViewingScreen(map, category);
         } else if (ConfigObject.getInstance().getRecipeScreenType() == RecipeScreenType.UNSET) {
+            @Nullable Identifier finalCategory = category;
             screen = new PreRecipeViewingScreen(ScreenHelper.getLastContainerScreen(), RecipeScreenType.UNSET, true, original -> {
                 ConfigObject.getInstance().setRecipeScreenType(original ? RecipeScreenType.ORIGINAL : RecipeScreenType.VILLAGER);
                 ConfigManager.getInstance().saveConfig();
-                openRecipeViewingScreen(map, ingredientNotice, resultNotice);
+                openRecipeViewingScreen(map, finalCategory, ingredientNotice, resultNotice);
             });
         } else {
-            screen = new RecipeViewingScreen(map);
+            screen = new RecipeViewingScreen(map, category);
         }
-        if (screen instanceof StackToNoticeScreen) {
+        if (screen instanceof RecipeScreen) {
             if (ingredientNotice != null)
-                ((StackToNoticeScreen) screen).addIngredientStackToNotice(ingredientNotice);
+                ((RecipeScreen) screen).addIngredientStackToNotice(ingredientNotice);
             if (resultNotice != null)
-                ((StackToNoticeScreen) screen).addResultStackToNotice(resultNotice);
+                ((RecipeScreen) screen).addResultStackToNotice(resultNotice);
         }
-        ScreenHelper.storeRecipeScreen(MinecraftClient.getInstance().currentScreen);
+        if (MinecraftClient.getInstance().currentScreen instanceof RecipeScreen)
+            ScreenHelper.storeRecipeScreen((RecipeScreen) MinecraftClient.getInstance().currentScreen);
         MinecraftClient.getInstance().openScreen(screen);
     }