Refactor how ScreenProcessors are created/registered

src/main/java/dev/isxander/controlify/mixins/feature/screenop/MinecraftMixin.java

@@ -13,7 +13,6 @@ import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 public class MinecraftMixin {
     @Inject(method = "setScreen", at = @At(value = "INVOKE", target = "Lcom/mojang/blaze3d/vertex/BufferUploader;reset()V"))
     private void changeScreen(Screen screen, CallbackInfo ci) {
-        ScreenProcessorProvider.REGISTRY.clearCache();
         ComponentProcessorProvider.REGISTRY.clearCache();
     }
 }

src/main/java/dev/isxander/controlify/mixins/feature/screenop/ScreenMixin.java

@@ -1,5 +1,6 @@
 package dev.isxander.controlify.mixins.feature.screenop;
 
+import dev.isxander.controlify.screenop.ScreenProcessorFactory;
 import dev.isxander.controlify.screenop.ScreenProcessorProvider;
 import dev.isxander.controlify.screenop.ScreenProcessor;
 import net.minecraft.client.Minecraft;
@@ -13,10 +14,10 @@ import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 @Mixin(Screen.class)
 public class ScreenMixin implements ScreenProcessorProvider {
     @Unique
-    private final ScreenProcessor<Screen> controlify$processor = new ScreenProcessor<>((Screen) (Object) this);
+    private final ScreenProcessor<? super Screen> controlify$processor = ScreenProcessorFactory.createForScreen((Screen) (Object) this);
 
     @Override
-    public ScreenProcessor<Screen> screenProcessor() {
+    public ScreenProcessor<?> screenProcessor() {
         return controlify$processor;
     }
 

src/main/java/dev/isxander/controlify/screenop/ScreenProcessorFactory.java

@@ -0,0 +1,32 @@
+package dev.isxander.controlify.screenop;
+
+import net.minecraft.client.gui.screens.Screen;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.function.Function;
+
+@SuppressWarnings("unchecked")
+public final class ScreenProcessorFactory {
+    private static final Map<Class<? extends Screen>, Function<Screen, ScreenProcessor<?>>> factories = new HashMap<>();
+
+    private ScreenProcessorFactory() {
+    }
+
+    public static <T extends Screen> ScreenProcessor<? super T> createForScreen(T screen) {
+        Class<? extends Screen> screenClass = screen.getClass();
+        while (!factories.containsKey(screenClass) && screenClass != Screen.class) {
+            screenClass = (Class<? extends Screen>) screenClass.getSuperclass();
+        }
+
+        return (ScreenProcessor<T>) factories.getOrDefault(screenClass, ScreenProcessorFactory::createDefault).apply(screen);
+    }
+
+    public static <T extends Screen> void registerProvider(Class<T> screenClass, Function<T, ScreenProcessor<? super T>> factory) {
+        factories.put(screenClass, (Function<Screen, ScreenProcessor<?>>) factory);
+    }
+
+    private static <T extends Screen> ScreenProcessor<T> createDefault(T screen) {
+        return new ScreenProcessor<>(screen);
+    }
+}

src/main/java/dev/isxander/controlify/screenop/ScreenProcessorProvider.java

@@ -3,20 +3,16 @@ package dev.isxander.controlify.screenop;
 import net.minecraft.client.gui.screens.Screen;
 import org.jetbrains.annotations.NotNull;
 
-import java.util.Optional;
+import java.util.function.Function;
 
 public interface ScreenProcessorProvider {
     ScreenProcessor<?> screenProcessor();
 
     static ScreenProcessor<?> provide(@NotNull Screen screen) {
-        Optional<ScreenProcessor<?>> optional = REGISTRY.get(screen);
-        if (optional.isPresent()) return optional.get();
-
         return ((ScreenProcessorProvider) screen).screenProcessor();
     }
 
-    /**
-     * Register a screen processor for a screen from an entrypoint
-     */
-    Registry<Screen, ScreenProcessor<?>> REGISTRY = new Registry<>();
+    static <T extends Screen> void registerProvider(@NotNull Class<T> screenClass, @NotNull Function<T, ScreenProcessor<? super T>> factory) {
+        ScreenProcessorFactory.registerProvider(screenClass, factory);
+    }
 }