Close #1421

5 files changed, 285 insertions, 48 deletions

diff --git a/api/src/main/java/me/shedaniel/rei/api/client/registry/screen/OverlayDecider.java b/api/src/main/java/me/shedaniel/rei/api/client/registry/screen/OverlayDecider.java
index 49b63e1f5..70db769e3 100644
--- a/api/src/main/java/me/shedaniel/rei/api/client/registry/screen/OverlayDecider.java
+++ b/api/src/main/java/me/shedaniel/rei/api/client/registry/screen/OverlayDecider.java
@@ -89,6 +89,19 @@ public interface OverlayDecider extends Comparable<OverlayDecider> {
     }
     
     /**
+     * Returns the provider for rendering the overlay.
+     * Return {@code null} to pass the rendering to the next decider.
+     *
+     * @return the provider for rendering the overlay
+     * @see OverlayRendererProvider for more information
+     * @since 8.3
+     */
+    @ApiStatus.Experimental
+    default OverlayRendererProvider getRendererProvider() {
+        return null;
+    }
+    
+    /**
      * {@inheritDoc}
      */
     @Override
diff --git a/api/src/main/java/me/shedaniel/rei/api/client/registry/screen/OverlayRendererProvider.java b/api/src/main/java/me/shedaniel/rei/api/client/registry/screen/OverlayRendererProvider.java
new file mode 100644
index 000000000..e690218b7
--- /dev/null
+++ b/api/src/main/java/me/shedaniel/rei/api/client/registry/screen/OverlayRendererProvider.java
@@ -0,0 +1,79 @@
+/*
+ * This file is licensed under the MIT License, part of Roughly Enough Items.
+ * Copyright (c) 2018, 2019, 2020, 2021, 2022, 2023 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.screen;
+
+import com.mojang.blaze3d.vertex.PoseStack;
+import me.shedaniel.rei.api.client.overlay.ScreenOverlay;
+import org.jetbrains.annotations.ApiStatus;
+
+/**
+ * The provider for a renderer provider.
+ * <p>
+ * {@link #onApplied(Sink)} is called when this provider is selected for a screen,
+ * and {@link #onRemoved()} is called when this provider is removed from a screen,
+ * for example if the screen is now no longer applicable for this provider, or
+ * the screen is closed.
+ * <p>
+ * {@link Sink} is given to the provider to allow the provider to render the overlay,
+ * whatever it deems necessary.
+ */
+@ApiStatus.Experimental
+public interface OverlayRendererProvider {
+    default void onApplied(Sink sink) {
+    }
+    
+    default void onRemoved() {
+    }
+    
+    @ApiStatus.NonExtendable
+    interface Sink {
+        /**
+         * Renders the overlay.
+         *
+         * @param matrices the matrices
+         * @param mouseX   the mouse x
+         * @param mouseY   the mouse y
+         * @param delta    the delta
+         */
+        void render(PoseStack matrices, int mouseX, int mouseY, float delta);
+        
+        /**
+         * Renders the overlay components that are supposed to be rendered last,
+         * for example, menu entries, or tooltips.
+         *
+         * @param matrices the matrices
+         * @param mouseX   the mouse x
+         * @param mouseY   the mouse y
+         * @param delta    the delta
+         */
+        void lateRender(PoseStack matrices, int mouseX, int mouseY, float delta);
+        
+        /**
+         * Returns the overlay.
+         *
+         * @return the overlay
+         */
+        ScreenOverlay getOverlay();
+    }
+}
diff --git a/runtime/src/main/java/me/shedaniel/rei/RoughlyEnoughItemsCoreClient.java b/runtime/src/main/java/me/shedaniel/rei/RoughlyEnoughItemsCoreClient.java
index 46f6589f2..5d5d53d8f 100644
--- a/runtime/src/main/java/me/shedaniel/rei/RoughlyEnoughItemsCoreClient.java
+++ b/runtime/src/main/java/me/shedaniel/rei/RoughlyEnoughItemsCoreClient.java
@@ -353,6 +353,7 @@ public class RoughlyEnoughItemsCoreClient {
                 }
                 
                 REIRuntimeImpl.getInstance().setPreviousScreen(screen);
+                ((ScreenRegistryImpl) ScreenRegistry.getInstance()).getLastRendererProvider(screen);
             }
             if (ConfigObject.getInstance().doesDisableRecipeBook() && screen instanceof AbstractContainerScreen) {
                 access.getRenderables().removeIf(widget -> widget instanceof ImageButton button && button.resourceLocation.equals(recipeButtonTex));
@@ -410,53 +411,6 @@ public class RoughlyEnoughItemsCoreClient {
                 return EventResult.interruptFalse();
             return EventResult.pass();
         });
-        int[] rendered = {0};
-        ClientGuiEvent.RENDER_PRE.register((screen, matrices, mouseX, mouseY, delta) -> {
-            if (shouldReturn(screen))
-                return EventResult.pass();
-            rendered[0] = 0;
-            return EventResult.pass();
-        });
-        ClientGuiEvent.RENDER_CONTAINER_BACKGROUND.register((screen, matrices, mouseX, mouseY, delta) -> {
-            if (shouldReturn(screen))
-                return;
-            rendered[0] = 1;
-            resetFocused(screen);
-            if (!(screen instanceof DisplayScreen)) {
-                getOverlay().render(matrices, mouseX, mouseY, delta);
-            }
-            resetFocused(screen);
-        });
-        ClientGuiEvent.RENDER_CONTAINER_FOREGROUND.register((screen, matrices, mouseX, mouseY, delta) -> {
-            if (shouldReturn(screen))
-                return;
-            rendered[0] = 2;
-            resetFocused(screen);
-            PoseStack poseStack = RenderSystem.getModelViewStack();
-            poseStack.pushPose();
-            poseStack.translate(-screen.leftPos, -screen.topPos, 0.0);
-            RenderSystem.applyModelViewMatrix();
-            ((ScreenOverlayImpl) getOverlay()).lateRender(matrices, mouseX, mouseY, delta);
-            poseStack.popPose();
-            RenderSystem.applyModelViewMatrix();
-            resetFocused(screen);
-        });
-        ClientGuiEvent.RENDER_POST.register((screen, matrices, mouseX, mouseY, delta) -> {
-            if (shouldReturn(screen) || rendered[0] == 2)
-                return;
-            if (screen instanceof AbstractContainerScreen) {
-                InternalLogger.getInstance().warn("Screen " + screen.getClass().getName() + " did not render background and foreground! This might cause rendering issues!");
-            }
-            resetFocused(screen);
-            if (rendered[0] == 0 && !(screen instanceof DisplayScreen)) {
-                getOverlay().render(matrices, mouseX, mouseY, delta);
-            }
-            rendered[0] = 1;
-            if (rendered[0] == 1) {
-                ((ScreenOverlayImpl) getOverlay()).lateRender(matrices, mouseX, mouseY, delta);
-            }
-            resetFocused(screen);
-        });
         ClientScreenInputEvent.MOUSE_DRAGGED_PRE.register((minecraftClient, screen, mouseX1, mouseY1, button, mouseX2, mouseY2) -> {
             if (shouldReturn(screen) || screen instanceof DisplayScreen)
                 return EventResult.pass();
@@ -506,7 +460,7 @@ public class RoughlyEnoughItemsCoreClient {
         });
     }
     
-    private boolean resetFocused(Screen screen) {
+    public static boolean resetFocused(Screen screen) {
         if (screen.getFocused() instanceof ScreenOverlay || screen.getFocused() == screen) {
             screen.setFocused(null);
         }
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/screen/DefaultScreenOverlayRenderer.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/screen/DefaultScreenOverlayRenderer.java
new file mode 100644
index 000000000..520b89bb7
--- /dev/null
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/screen/DefaultScreenOverlayRenderer.java
@@ -0,0 +1,112 @@
+/*
+ * This file is licensed under the MIT License, part of Roughly Enough Items.
+ * Copyright (c) 2018, 2019, 2020, 2021, 2022, 2023 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.impl.client.registry.screen;
+
+import com.mojang.blaze3d.systems.RenderSystem;
+import com.mojang.blaze3d.vertex.PoseStack;
+import dev.architectury.event.EventResult;
+import dev.architectury.event.events.client.ClientGuiEvent;
+import me.shedaniel.rei.api.client.gui.screen.DisplayScreen;
+import me.shedaniel.rei.api.client.registry.screen.OverlayRendererProvider;
+import me.shedaniel.rei.impl.common.InternalLogger;
+import net.minecraft.client.gui.screens.inventory.AbstractContainerScreen;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import static me.shedaniel.rei.RoughlyEnoughItemsCoreClient.resetFocused;
+import static me.shedaniel.rei.RoughlyEnoughItemsCoreClient.shouldReturn;
+
+public enum DefaultScreenOverlayRenderer implements OverlayRendererProvider {
+    INSTANCE;
+    
+    private final List<Runnable> onRemoved = new ArrayList<>();
+    
+    @Override
+    public void onApplied(Sink sink) {
+        int[] rendered = {0};
+        ClientGuiEvent.ScreenRenderPre renderPre;
+        ClientGuiEvent.ContainerScreenRenderBackground renderContainerBg;
+        ClientGuiEvent.ContainerScreenRenderForeground renderContainerFg;
+        ClientGuiEvent.ScreenRenderPost renderPost;
+        ClientGuiEvent.RENDER_PRE.register(renderPre = (screen, matrices, mouseX, mouseY, delta) -> {
+            if (shouldReturn(screen))
+                return EventResult.pass();
+            rendered[0] = 0;
+            return EventResult.pass();
+        });
+        ClientGuiEvent.RENDER_CONTAINER_BACKGROUND.register(renderContainerBg = (screen, matrices, mouseX, mouseY, delta) -> {
+            if (shouldReturn(screen))
+                return;
+            rendered[0] = 1;
+            resetFocused(screen);
+            if (!(screen instanceof DisplayScreen)) {
+                sink.render(matrices, mouseX, mouseY, delta);
+            }
+            resetFocused(screen);
+        });
+        ClientGuiEvent.RENDER_CONTAINER_FOREGROUND.register(renderContainerFg = (screen, matrices, mouseX, mouseY, delta) -> {
+            if (shouldReturn(screen))
+                return;
+            rendered[0] = 2;
+            resetFocused(screen);
+            PoseStack poseStack = RenderSystem.getModelViewStack();
+            poseStack.pushPose();
+            poseStack.translate(-screen.leftPos, -screen.topPos, 0.0);
+            RenderSystem.applyModelViewMatrix();
+            sink.lateRender(matrices, mouseX, mouseY, delta);
+            poseStack.popPose();
+            RenderSystem.applyModelViewMatrix();
+            resetFocused(screen);
+        });
+        ClientGuiEvent.RENDER_POST.register(renderPost = (screen, matrices, mouseX, mouseY, delta) -> {
+            if (shouldReturn(screen) || rendered[0] == 2)
+                return;
+            if (screen instanceof AbstractContainerScreen) {
+                InternalLogger.getInstance().warn("Screen " + screen.getClass().getName() + " did not render background and foreground! This might cause rendering issues!");
+            }
+            resetFocused(screen);
+            if (rendered[0] == 0 && !(screen instanceof DisplayScreen)) {
+                sink.render(matrices, mouseX, mouseY, delta);
+            }
+            rendered[0] = 1;
+            if (rendered[0] == 1) {
+                sink.lateRender(matrices, mouseX, mouseY, delta);
+            }
+            resetFocused(screen);
+        });
+        this.onRemoved.add(() -> {
+            ClientGuiEvent.RENDER_PRE.unregister(renderPre);
+            ClientGuiEvent.RENDER_CONTAINER_BACKGROUND.unregister(renderContainerBg);
+            ClientGuiEvent.RENDER_CONTAINER_FOREGROUND.unregister(renderContainerFg);
+            ClientGuiEvent.RENDER_POST.unregister(renderPost);
+        });
+    }
+    
+    @Override
+    public void onRemoved() {
+        this.onRemoved.forEach(Runnable::run);
+        this.onRemoved.clear();
+    }
+}
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/screen/ScreenRegistryImpl.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/screen/ScreenRegistryImpl.java
index 85053fb73..248f76cfc 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/screen/ScreenRegistryImpl.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/screen/ScreenRegistryImpl.java
@@ -26,6 +26,7 @@ package me.shedaniel.rei.impl.client.registry.screen;
 import com.google.common.collect.HashMultimap;
 import com.google.common.collect.Multimap;
 import com.mojang.blaze3d.platform.Window;
+import com.mojang.blaze3d.vertex.PoseStack;
 import dev.architectury.event.CompoundEventResult;
 import me.shedaniel.math.Point;
 import me.shedaniel.math.Rectangle;
@@ -37,6 +38,7 @@ import me.shedaniel.rei.api.client.gui.drag.component.DraggableComponentProvider
 import me.shedaniel.rei.api.client.gui.drag.component.DraggableComponentVisitor;
 import me.shedaniel.rei.api.client.gui.drag.component.DraggableComponentVisitorWidget;
 import me.shedaniel.rei.api.client.gui.widgets.Widgets;
+import me.shedaniel.rei.api.client.overlay.ScreenOverlay;
 import me.shedaniel.rei.api.client.plugins.REIClientPlugin;
 import me.shedaniel.rei.api.client.registry.screen.*;
 import me.shedaniel.rei.api.common.category.CategoryIdentifier;
@@ -44,6 +46,7 @@ import me.shedaniel.rei.api.common.entry.EntryStack;
 import me.shedaniel.rei.api.common.registry.ReloadStage;
 import me.shedaniel.rei.api.common.util.CollectionUtils;
 import me.shedaniel.rei.api.common.util.EntryStacks;
+import me.shedaniel.rei.impl.client.gui.ScreenOverlayImpl;
 import me.shedaniel.rei.impl.client.gui.screen.AbstractDisplayViewingScreen;
 import me.shedaniel.rei.impl.common.InternalLogger;
 import net.fabricmc.api.EnvType;
@@ -69,6 +72,7 @@ public class ScreenRegistryImpl implements ScreenRegistry {
     private List<OverlayDecider> deciders = new CopyOnWriteArrayList<>();
     private Map<Class<?>, List<OverlayDecider>> cache = new HashMap<>();
     private ExclusionZones exclusionZones = new ExclusionZonesImpl();
+    private OverlayRendererProvider lastRendererProvider = null;
     private final ThreadLocal<Class<? extends Screen>> tmpScreen = new ThreadLocal<>();
     
     @Override
@@ -245,12 +249,20 @@ public class ScreenRegistryImpl implements ScreenRegistry {
         draggableProviders.clear();
         draggableVisitors.clear();
         tmpScreen.remove();
+        if (this.lastRendererProvider != null) {
+            this.lastRendererProvider.onRemoved();
+        }
+        this.lastRendererProvider = null;
         
         registerDefault();
     }
     
     @Override
     public void endReload() {
+        if (this.lastRendererProvider != null) {
+            this.lastRendererProvider.onRemoved();
+        }
+        this.lastRendererProvider = null;
         InternalLogger.getInstance().debug("Registered %d overlay deciders and %d exclusion zones", deciders.size(), exclusionZones.getZonesCount());
     }
     
@@ -298,6 +310,22 @@ public class ScreenRegistryImpl implements ScreenRegistry {
                 return -10.0;
             }
         });
+        registerDecider(new OverlayDecider() {
+            @Override
+            public <R extends Screen> boolean isHandingScreen(Class<R> screen) {
+                return true;
+            }
+            
+            @Override
+            public OverlayRendererProvider getRendererProvider() {
+                return DefaultScreenOverlayRenderer.INSTANCE;
+            }
+            
+            @Override
+            public double getPriority() {
+                return -20.0;
+            }
+        });
         registerFocusedStack(new FocusedStackProvider() {
             @Override
             public CompoundEventResult<EntryStack<?>> provide(Screen screen, Point mouse) {
@@ -318,4 +346,55 @@ public class ScreenRegistryImpl implements ScreenRegistry {
         registerDraggableComponentVisitor(DraggableComponentVisitorWidget.from(context ->
                 Widgets.walk(context.getScreen().children(), DraggableComponentVisitorWidget.class::isInstance)));
     }
+    
+    public OverlayRendererProvider getLastRendererProvider(Screen screen) {
+        for (OverlayDecider decider : getDeciders(screen)) {
+            if (decider.isHandingScreen(screen.getClass())) {
+                OverlayRendererProvider provider = decider.getRendererProvider();
+                
+                if (provider != null) {
+                    return applyRendererProvider(provider);
+                }
+            }
+        }
+        
+        return applyRendererProvider(null);
+    }
+    
+    private OverlayRendererProvider applyRendererProvider(OverlayRendererProvider provider) {
+        if (provider != this.lastRendererProvider) {
+            if (this.lastRendererProvider != null) {
+                this.lastRendererProvider.onRemoved();
+            }
+            this.lastRendererProvider = provider;
+            if (this.lastRendererProvider != null) {
+                this.lastRendererProvider.onApplied(new OverlayRendererProvider.Sink() {
+                    @Override
+                    public void render(PoseStack matrices, int mouseX, int mouseY, float delta) {
+                        if (ScreenRegistryImpl.this.lastRendererProvider == provider) {
+                            ScreenOverlayImpl.getInstance().render(matrices, mouseX, mouseY, delta);
+                        } else {
+                            InternalLogger.getInstance().warn("Renderer provider %s still tries to render after being removed!", provider);
+                        }
+                    }
+                    
+                    @Override
+                    public void lateRender(PoseStack matrices, int mouseX, int mouseY, float delta) {
+                        if (ScreenRegistryImpl.this.lastRendererProvider == provider) {
+                            ScreenOverlayImpl.getInstance().lateRender(matrices, mouseX, mouseY, delta);
+                        } else {
+                            InternalLogger.getInstance().warn("Renderer provider %s still tries to render after being removed!", provider);
+                        }
+                    }
+                    
+                    @Override
+                    public ScreenOverlay getOverlay() {
+                        return ScreenOverlayImpl.getInstance();
+                    }
+                });
+            }
+        }
+        
+        return provider;
+    }
 }