Fix #662

5 files changed, 30 insertions, 4 deletions

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;
+        }
     }
 }