Close #1421

2 files changed, 191 insertions, 0 deletions

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