Fix #662

11 files changed, 156 insertions, 5 deletions

diff --git a/api/src/main/java/me/shedaniel/rei/api/client/registry/category/CategoryRegistry.java b/api/src/main/java/me/shedaniel/rei/api/client/registry/category/CategoryRegistry.java
index 6e83dfac2..ebe4f1000 100644
--- a/api/src/main/java/me/shedaniel/rei/api/client/registry/category/CategoryRegistry.java
+++ b/api/src/main/java/me/shedaniel/rei/api/client/registry/category/CategoryRegistry.java
@@ -24,8 +24,10 @@
 package me.shedaniel.rei.api.client.registry.category;
 
 import me.shedaniel.rei.api.client.plugins.REIClientPlugin;
+import me.shedaniel.rei.api.client.registry.category.extension.CategoryExtensionProvider;
 import me.shedaniel.rei.api.client.registry.category.visibility.CategoryVisibilityPredicate;
 import me.shedaniel.rei.api.client.registry.display.DisplayCategory;
+import me.shedaniel.rei.api.client.registry.display.DisplayCategoryView;
 import me.shedaniel.rei.api.common.category.CategoryIdentifier;
 import me.shedaniel.rei.api.common.display.Display;
 import me.shedaniel.rei.api.common.entry.EntryIngredient;
@@ -37,6 +39,7 @@ import me.shedaniel.rei.api.common.util.Identifiable;
 import net.fabricmc.api.EnvType;
 import net.fabricmc.api.Environment;
 import net.minecraft.resources.ResourceLocation;
+import org.jetbrains.annotations.ApiStatus;
 
 import java.util.List;
 import java.util.Optional;
@@ -215,6 +218,12 @@ public interface CategoryRegistry extends Reloadable<REIClientPlugin>, Iterable<
          */
         CategoryIdentifier<?> getCategoryIdentifier();
         
+        @ApiStatus.Experimental
+        void registerExtension(CategoryExtensionProvider<T> provider);
+        
+        @ApiStatus.Experimental
+        DisplayCategoryView<T> getView(T display);
+        
         @Override
         default ResourceLocation getIdentifier() {
             return getCategoryIdentifier().getIdentifier();
diff --git a/api/src/main/java/me/shedaniel/rei/api/client/registry/category/extension/CategoryExtensionProvider.java b/api/src/main/java/me/shedaniel/rei/api/client/registry/category/extension/CategoryExtensionProvider.java
new file mode 100644
index 000000000..58fcbb8de
--- /dev/null
+++ b/api/src/main/java/me/shedaniel/rei/api/client/registry/category/extension/CategoryExtensionProvider.java
@@ -0,0 +1,49 @@
+/*
+ * This file is licensed under the MIT License, part of Roughly Enough Items.
+ * Copyright (c) 2018, 2019, 2020, 2021 shedaniel
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package me.shedaniel.rei.api.client.registry.category.extension;
+
+import me.shedaniel.rei.api.client.registry.display.DisplayCategory;
+import me.shedaniel.rei.api.client.registry.display.DisplayCategoryView;
+import me.shedaniel.rei.api.common.display.Display;
+import net.fabricmc.api.EnvType;
+import net.fabricmc.api.Environment;
+import org.jetbrains.annotations.ApiStatus;
+
+@FunctionalInterface
+@ApiStatus.Experimental
+@Environment(EnvType.CLIENT)
+public interface CategoryExtensionProvider<T extends Display> {
+    /**
+     * Returns a new {@link DisplayCategoryView} for a specific {@link Display},
+     * a previous {@link DisplayCategoryView} will be provided, this may be the original category.
+     * {@code null} is not an accepted value, return the previous view if this provider
+     * does not modify the view.
+     *
+     * @param display  the display to display for, do not store or cache this
+     * @param category the category of the display, do not store or cache this
+     * @param lastView the previous category view
+     * @return the new category view, {@code null} is not accepted here
+     */
+    DisplayCategoryView<T> provide(T display, DisplayCategory<T> category, DisplayCategoryView<T> lastView);
+}
diff --git a/api/src/main/java/me/shedaniel/rei/api/client/registry/display/DisplayCategory.java b/api/src/main/java/me/shedaniel/rei/api/client/registry/display/DisplayCategory.java
index 9202a18ca..0b672e9a6 100644
--- a/api/src/main/java/me/shedaniel/rei/api/client/registry/display/DisplayCategory.java
+++ b/api/src/main/java/me/shedaniel/rei/api/client/registry/display/DisplayCategory.java
@@ -44,7 +44,7 @@ import java.util.Collections;
 import java.util.List;
 
 @Environment(EnvType.CLIENT)
-public interface DisplayCategory<T extends Display> extends Identifiable {
+public interface DisplayCategory<T extends Display> extends DisplayCategoryView<T>, Identifiable {
     /**
      * Returns the renderer of the icon.
      *
@@ -66,6 +66,7 @@ public interface DisplayCategory<T extends Display> extends Identifiable {
      * @return the display renderer
      */
     @ApiStatus.OverrideOnly
+    @Override
     default DisplayRenderer getDisplayRenderer(T display) {
         return SimpleDisplayRenderer.from(display.getInputEntries(), display.getOutputEntries());
     }
@@ -78,6 +79,7 @@ public interface DisplayCategory<T extends Display> extends Identifiable {
      * @return the list of widgets
      */
     @ApiStatus.OverrideOnly
+    @Override
     default List<Widget> setupDisplay(T display, Rectangle bounds) {
         return Collections.singletonList(Widgets.createRecipeBase(bounds));
     }
diff --git a/api/src/main/java/me/shedaniel/rei/api/client/registry/display/DisplayCategoryView.java b/api/src/main/java/me/shedaniel/rei/api/client/registry/display/DisplayCategoryView.java
new file mode 100644
index 000000000..57fb3b3cc
--- /dev/null
+++ b/api/src/main/java/me/shedaniel/rei/api/client/registry/display/DisplayCategoryView.java
@@ -0,0 +1,54 @@
+/*
+ * This file is licensed under the MIT License, part of Roughly Enough Items.
+ * Copyright (c) 2018, 2019, 2020, 2021 shedaniel
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package me.shedaniel.rei.api.client.registry.display;
+
+import me.shedaniel.math.Rectangle;
+import me.shedaniel.rei.api.client.gui.DisplayRenderer;
+import me.shedaniel.rei.api.client.gui.widgets.Widget;
+import me.shedaniel.rei.api.common.display.Display;
+import org.jetbrains.annotations.ApiStatus;
+
+import java.util.List;
+
+@ApiStatus.Experimental
+public interface DisplayCategoryView<T extends Display> {
+    /**
+     * Gets the recipe renderer for the category, used in {@link me.shedaniel.rei.impl.client.gui.CompositeRecipeViewingScreen} for rendering simple recipes
+     *
+     * @param display the display to render
+     * @return the display renderer
+     */
+    @ApiStatus.OverrideOnly
+    DisplayRenderer getDisplayRenderer(T display);
+    
+    /**
+     * Setup the widgets for displaying the recipe
+     *
+     * @param display the recipe
+     * @param bounds  the bounds of the display, configurable with overriding the width, height methods.
+     * @return the list of widgets
+     */
+    @ApiStatus.OverrideOnly
+    List<Widget> setupDisplay(T display, Rectangle bounds);
+}
diff --git a/api/src/main/java/me/shedaniel/rei/api/client/registry/display/DisplayRegistry.java b/api/src/main/java/me/shedaniel/rei/api/client/registry/display/DisplayRegistry.java
index ad3a2aa18..9f895feeb 100644
--- a/api/src/main/java/me/shedaniel/rei/api/client/registry/display/DisplayRegistry.java
+++ b/api/src/main/java/me/shedaniel/rei/api/client/registry/display/DisplayRegistry.java
@@ -331,6 +331,12 @@ public interface DisplayRegistry extends RecipeManagerContext<REIClientPlugin> {
     @ApiStatus.Experimental
     <T> Collection<Display> tryFillDisplay(T value, DisplayAdditionReason... reasons);
     
+    /**
+     * Returns the origin of the display, this may be the recipe which the display was created from.
+     *
+     * @param display the display
+     * @return the origin
+     */
     @Nullable
     Object getDisplayOrigin(Display display);
 }
diff --git a/forge/build.gradle b/forge/build.gradle
index ed93685d4..aa99b6d02 100644
--- a/forge/build.gradle
+++ b/forge/build.gradle
@@ -101,6 +101,11 @@ dependencies {
     // modRuntime("curse.maven:cc-tweaked-282001:3434680")
     // modRuntime("curse.maven:ender-storage-245174:3361748")
     // modRuntime("curse.maven:codechicken-242818:3528603")
+    // modRuntime("curse.maven:recipe-stages-280554:3492760")
+    // modRuntime("curse.maven:game-stages-268655:3437962")
+    // modRuntime("curse.maven:bookshelf-228525:3474499")
+    // modRuntime("curse.maven:kubejs-238086:3533307")
+    // modRuntime("curse.maven:rhino-416294:3525704")
 }
 
 shadowJar {
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/screen/AbstractDisplayViewingScreen.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/screen/AbstractDisplayViewingScreen.java
index 96d290a9a..680527b8e 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/screen/AbstractDisplayViewingScreen.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/screen/AbstractDisplayViewingScreen.java
@@ -30,7 +30,9 @@ import me.shedaniel.rei.api.client.gui.screen.DisplayScreen;
 import me.shedaniel.rei.api.client.gui.widgets.Slot;
 import me.shedaniel.rei.api.client.gui.widgets.Widget;
 import me.shedaniel.rei.api.client.gui.widgets.Widgets;
+import me.shedaniel.rei.api.client.registry.category.CategoryRegistry;
 import me.shedaniel.rei.api.client.registry.display.DisplayCategory;
+import me.shedaniel.rei.api.client.registry.display.DisplayCategoryView;
 import me.shedaniel.rei.api.common.category.CategoryIdentifier;
 import me.shedaniel.rei.api.common.display.Display;
 import me.shedaniel.rei.api.common.entry.EntryStack;
@@ -120,6 +122,11 @@ public abstract class AbstractDisplayViewingScreen extends Screen implements Dis
         return (DisplayCategory<Display>) categories.get(selectedCategoryIndex);
     }
     
+    protected DisplayCategoryView<Display> getCurrentCategoryView(Display display) {
+        return CategoryRegistry.getInstance().get(categories.get(selectedCategoryIndex).getCategoryIdentifier().cast())
+                .getView(display);
+    }
+    
     @Override
     public void previousCategory() {
         int currentCategoryIndex = selectedCategoryIndex;
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/screen/CompositeDisplayViewingScreen.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/screen/CompositeDisplayViewingScreen.java
index c207a44ab..27b6c3cb0 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/screen/CompositeDisplayViewingScreen.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/screen/CompositeDisplayViewingScreen.java
@@ -118,7 +118,7 @@ public class CompositeDisplayViewingScreen extends AbstractDisplayViewingScreen
         this.tabs.clear();
         int largestWidth = width - 100;
         int largestHeight = height - 40;
-        DisplayCategory<Display> category = (DisplayCategory<Display>) categories.get(selectedCategoryIndex);
+        DisplayCategory<Display> category = getCurrentCategory();
         DisplaySpec display = categoryMap.get(category).get(selectedRecipeIndex);
         int guiWidth = Mth.clamp(category.getDisplayWidth(display.provideInternalDisplay()) + 30, 0, largestWidth) + 100;
         int guiHeight = Mth.clamp(category.getDisplayHeight() + 40, 166, largestHeight);
@@ -157,7 +157,7 @@ public class CompositeDisplayViewingScreen extends AbstractDisplayViewingScreen
         Rectangle recipeBounds = new Rectangle(bounds.x + 100 + (guiWidth - 100) / 2 - category.getDisplayWidth(display.provideInternalDisplay()) / 2, bounds.y + bounds.height / 2 - category.getDisplayHeight() / 2, category.getDisplayWidth(display.provideInternalDisplay()), category.getDisplayHeight());
         List<Widget> setupDisplay;
         try {
-            setupDisplay = category.setupDisplay(display.provideInternalDisplay(), recipeBounds);
+            setupDisplay = getCurrentCategoryView(display.provideInternalDisplay()).setupDisplay(display.provideInternalDisplay(), recipeBounds);
         } catch (Throwable throwable) {
             throwable.printStackTrace();
             setupDisplay = new ArrayList<>();
@@ -180,7 +180,7 @@ public class CompositeDisplayViewingScreen extends AbstractDisplayViewingScreen
         for (DisplaySpec recipeDisplay : categoryMap.get(category)) {
             int finalIndex = index;
             DisplayRenderer displayRenderer;
-            displayRenderers.add(displayRenderer = category.getDisplayRenderer(recipeDisplay.provideInternalDisplay()));
+            displayRenderers.add(displayRenderer = getCurrentCategoryView(recipeDisplay.provideInternalDisplay()).getDisplayRenderer(recipeDisplay.provideInternalDisplay()));
             buttonList.add(Widgets.createButton(new Rectangle(bounds.x + 5, 0, displayRenderer.getWidth(), displayRenderer.getHeight()), NarratorChatListener.NO_TITLE)
                     .onClick(button -> {
                         selectedRecipeIndex = finalIndex;
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/screen/DefaultDisplayViewingScreen.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/screen/DefaultDisplayViewingScreen.java
index 37e329a95..5ee3ece56 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/screen/DefaultDisplayViewingScreen.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/screen/DefaultDisplayViewingScreen.java
@@ -256,7 +256,7 @@ public class DefaultDisplayViewingScreen extends AbstractDisplayViewingScreen {
             final Rectangle displayBounds = new Rectangle(getBounds().getCenterX() - displayWidth / 2, getBounds().getCenterY() + 16 - displayHeight * (getRecipesPerPage() + 1) / 2 - 2 * (getRecipesPerPage() + 1) + displayHeight * i + 4 * i, displayWidth, displayHeight);
             List<Widget> setupDisplay;
             try {
-                setupDisplay = getCurrentCategory().setupDisplay(display.provideInternalDisplay(), displayBounds);
+                setupDisplay = getCurrentCategoryView(display.provideInternalDisplay()).setupDisplay(display.provideInternalDisplay(), displayBounds);
             } catch (Throwable throwable) {
                 throwable.printStackTrace();
                 setupDisplay = new ArrayList<>();
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/InternalWidgets.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/InternalWidgets.java
index 5a9712575..47a1edde0 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/InternalWidgets.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/InternalWidgets.java
@@ -107,6 +107,7 @@ public final class InternalWidgets {
                 });
         return new DelegateWidget(autoCraftingButton) {
             boolean didJustRender = false;
+            
             @Override
             public void render(PoseStack poses, int mouseX, int mouseY, float delta) {
                 didJustRender = false;
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/category/CategoryRegistryImpl.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/category/CategoryRegistryImpl.java
index ad1aaa431..8deeecdc5 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/category/CategoryRegistryImpl.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/category/CategoryRegistryImpl.java
@@ -28,8 +28,10 @@ import me.shedaniel.rei.RoughlyEnoughItemsCore;
 import me.shedaniel.rei.api.client.plugins.REIClientPlugin;
 import me.shedaniel.rei.api.client.registry.category.ButtonArea;
 import me.shedaniel.rei.api.client.registry.category.CategoryRegistry;
+import me.shedaniel.rei.api.client.registry.category.extension.CategoryExtensionProvider;
 import me.shedaniel.rei.api.client.registry.category.visibility.CategoryVisibilityPredicate;
 import me.shedaniel.rei.api.client.registry.display.DisplayCategory;
+import me.shedaniel.rei.api.client.registry.display.DisplayCategoryView;
 import me.shedaniel.rei.api.common.category.CategoryIdentifier;
 import me.shedaniel.rei.api.common.display.Display;
 import me.shedaniel.rei.api.common.entry.EntryIngredient;
@@ -129,6 +131,7 @@ public class CategoryRegistryImpl implements CategoryRegistry {
     private static class Configuration<T extends Display> implements CategoryConfiguration<T> {
         private final DisplayCategory<T> category;
         private final List<EntryIngredient> workstations = Collections.synchronizedList(new ArrayList<>());
+        private final List<CategoryExtensionProvider<T>> extensionProviders = Collections.synchronizedList(new ArrayList<>());
         
         private Optional<ButtonArea> plusButtonArea = Optional.of(ButtonArea.defaultArea());
         
@@ -165,5 +168,20 @@ public class CategoryRegistryImpl implements CategoryRegistry {
         public CategoryIdentifier<?> getCategoryIdentifier() {
             return this.category.getCategoryIdentifier();
         }
+        
+        @Override
+        public void registerExtension(CategoryExtensionProvider<T> provider) {
+            this.extensionProviders.add(provider);
+        }
+        
+        @Override
+        public DisplayCategoryView<T> getView(T display) {
+            DisplayCategory<T> category = getCategory();
+            DisplayCategoryView<T> view = getCategory();
+            for (CategoryExtensionProvider<T> provider : this.extensionProviders) {
+                view = Objects.requireNonNull(provider.provide(display, category, view));
+            }
+            return view;
+        }
     }
 }