extreme image preloading hackery

7 files changed, 160 insertions, 22 deletions

diff --git a/common/src/main/java/dev/isxander/yacl3/gui/image/ImageRendererFactory.java b/common/src/main/java/dev/isxander/yacl3/gui/image/ImageRendererFactory.java
index 480d8b1..d9d2e2d 100644
--- a/common/src/main/java/dev/isxander/yacl3/gui/image/ImageRendererFactory.java
+++ b/common/src/main/java/dev/isxander/yacl3/gui/image/ImageRendererFactory.java
@@ -1,21 +1,21 @@
 package dev.isxander.yacl3.gui.image;
 
-public interface ImageRendererFactory<T extends ImageRenderer> {
+public interface ImageRendererFactory {
     /**
      * Prepares the image. This can be run off-thread,
      * and should NOT contain any GL calls whatsoever.
      */
-    ImageSupplier<T> prepareImage() throws Exception;
+    ImageSupplier prepareImage() throws Exception;
 
     default boolean requiresOffThreadPreparation() {
         return true;
     }
 
-    interface ImageSupplier<T extends ImageRenderer> {
-        T completeImage() throws Exception;
+    interface ImageSupplier {
+        ImageRenderer completeImage() throws Exception;
     }
 
-    interface OnThread<T extends ImageRenderer> extends ImageRendererFactory<T> {
+    interface OnThread extends ImageRendererFactory {
         @Override
         default boolean requiresOffThreadPreparation() {
             return false;
diff --git a/common/src/main/java/dev/isxander/yacl3/gui/image/ImageRendererManager.java b/common/src/main/java/dev/isxander/yacl3/gui/image/ImageRendererManager.java
index 680a6da..0c9b8a3 100644
--- a/common/src/main/java/dev/isxander/yacl3/gui/image/ImageRendererManager.java
+++ b/common/src/main/java/dev/isxander/yacl3/gui/image/ImageRendererManager.java
@@ -1,22 +1,52 @@
 package dev.isxander.yacl3.gui.image;
 
 import com.mojang.blaze3d.systems.RenderSystem;
+import dev.isxander.yacl3.gui.image.impl.AnimatedDynamicTextureImage;
 import dev.isxander.yacl3.impl.utils.YACLConstants;
 import net.minecraft.client.Minecraft;
 import net.minecraft.resources.ResourceLocation;
 
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.*;
+import java.util.function.Function;
+import java.util.function.Predicate;
 import java.util.function.Supplier;
 
 public class ImageRendererManager {
     private static final ExecutorService SINGLE_THREAD_EXECUTOR = Executors.newSingleThreadExecutor(task -> new Thread(task, "YACL Image Prep"));
 
     private static final Map<ResourceLocation, CompletableFuture<ImageRenderer>> IMAGE_CACHE = new ConcurrentHashMap<>();
+    static final Map<ResourceLocation, ImageRenderer> PRELOADED_IMAGE_CACHE = new ConcurrentHashMap<>();
+
+    static final List<PreloadedImageFactory> PRELOADED_IMAGE_FACTORIES = List.of(
+            new PreloadedImageFactory(
+                    location -> location.getPath().endsWith(".webp"),
+                    AnimatedDynamicTextureImage::createWEBPFromTexture
+            ),
+            new PreloadedImageFactory(
+                    location -> location.getPath().endsWith(".gif"),
+                    AnimatedDynamicTextureImage::createGIFFromTexture
+            )
+    );
+
+    public static <T extends ImageRenderer> Optional<T> getImage(ResourceLocation id) {
+        if (PRELOADED_IMAGE_CACHE.containsKey(id)) {
+            return Optional.of((T) PRELOADED_IMAGE_CACHE.get(id));
+        }
+
+        if (IMAGE_CACHE.containsKey(id)) {
+            return Optional.ofNullable((T) IMAGE_CACHE.get(id).getNow(null));
+        }
+
+        return Optional.empty();
+    }
 
     @SuppressWarnings("unchecked")
-    public static <T extends ImageRenderer> CompletableFuture<T> registerImage(ResourceLocation id, ImageRendererFactory<T> factory) {
+    public static <T extends ImageRenderer> CompletableFuture<T> registerImage(ResourceLocation id, ImageRendererFactory factory) {
+        System.out.println(PRELOADED_IMAGE_CACHE.get(id));
+
         if (IMAGE_CACHE.containsKey(id)) {
             return (CompletableFuture<T>) IMAGE_CACHE.get(id);
         }
@@ -25,7 +55,7 @@ public class ImageRendererManager {
         IMAGE_CACHE.put(id, future);
 
         SINGLE_THREAD_EXECUTOR.submit(() -> {
-            Supplier<Optional<ImageRendererFactory.ImageSupplier<T>>> supplier =
+            Supplier<Optional<ImageRendererFactory.ImageSupplier>> supplier =
                     factory.requiresOffThreadPreparation()
                             ? new CompletedSupplier<>(safelyPrepareFactory(id, factory))
                             : () -> safelyPrepareFactory(id, factory);
@@ -36,10 +66,10 @@ public class ImageRendererManager {
         return (CompletableFuture<T>) future;
     }
 
-    private static <T extends ImageRenderer> void completeImageFactory(ResourceLocation id, Supplier<Optional<ImageRendererFactory.ImageSupplier<T>>> supplier, CompletableFuture<ImageRenderer> future) {
+    private static <T extends ImageRenderer> void completeImageFactory(ResourceLocation id, Supplier<Optional<ImageRendererFactory.ImageSupplier>> supplier, CompletableFuture<ImageRenderer> future) {
         RenderSystem.assertOnRenderThread();
 
-        ImageRendererFactory.ImageSupplier<T> completableImage = supplier.get().orElse(null);
+        ImageRendererFactory.ImageSupplier completableImage = supplier.get().orElse(null);
         if (completableImage == null) {
             return;
         }
@@ -71,7 +101,7 @@ public class ImageRendererManager {
         });
     }
 
-    private static <T extends ImageRenderer> Optional<ImageRendererFactory.ImageSupplier<T>> safelyPrepareFactory(ResourceLocation id, ImageRendererFactory<T> factory) {
+    static Optional<ImageRendererFactory.ImageSupplier> safelyPrepareFactory(ResourceLocation id, ImageRendererFactory factory) {
         try {
             return Optional.of(factory.prepareImage());
         } catch (Exception e) {
@@ -81,6 +111,9 @@ public class ImageRendererManager {
         }
     }
 
+    public record PreloadedImageFactory(Predicate<ResourceLocation> predicate, Function<ResourceLocation, ImageRendererFactory> factory) {
+    }
+
     private record CompletedSupplier<T>(T get) implements Supplier<T> {
     }
 
diff --git a/common/src/main/java/dev/isxander/yacl3/gui/image/YACLImageReloadListener.java b/common/src/main/java/dev/isxander/yacl3/gui/image/YACLImageReloadListener.java
new file mode 100644
index 0000000..cc259b0
--- /dev/null
+++ b/common/src/main/java/dev/isxander/yacl3/gui/image/YACLImageReloadListener.java
@@ -0,0 +1,99 @@
+package dev.isxander.yacl3.gui.image;
+
+import dev.isxander.yacl3.impl.utils.YACLConstants;
+import net.minecraft.CrashReport;
+import net.minecraft.CrashReportCategory;
+import net.minecraft.ReportedException;
+import net.minecraft.resources.ResourceLocation;
+import net.minecraft.server.packs.resources.PreparableReloadListener;
+import net.minecraft.server.packs.resources.Resource;
+import net.minecraft.server.packs.resources.ResourceManager;
+import net.minecraft.util.profiling.ProfilerFiller;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
+import java.util.function.Function;
+
+public class YACLImageReloadListener implements PreparableReloadListener {
+    @Override
+    public CompletableFuture<Void> reload(
+            PreparationBarrier preparationBarrier, 
+            ResourceManager resourceManager, 
+            ProfilerFiller preparationsProfiler, 
+            ProfilerFiller reloadProfiler, 
+            Executor backgroundExecutor, 
+            Executor gameExecutor
+    ) {
+        Map<ResourceLocation, Resource> imageResources = resourceManager.listResources(
+                "",
+                location -> ImageRendererManager.PRELOADED_IMAGE_FACTORIES
+                        .stream()
+                        .anyMatch(factory -> factory.predicate().test(location))
+        );
+
+        // extreme mojang hackery.
+        // for some reason this wait method needs to be called for the reload
+        // instance to be marked as complete
+        if (imageResources.isEmpty()) {
+            preparationBarrier.wait(null);
+        }
+
+        List<CompletableFuture<?>> futures = new ArrayList<>(imageResources.size());
+
+        for (Map.Entry<ResourceLocation, Resource> entry : imageResources.entrySet()) {
+            ResourceLocation location = entry.getKey();
+            Resource resource = entry.getValue();
+
+            ImageRendererFactory imageFactory = ImageRendererManager.PRELOADED_IMAGE_FACTORIES
+                    .stream()
+                    .filter(factory -> factory.predicate().test(location))
+                    .map(factory -> factory.factory().apply(location))
+                    .findAny()
+                    .orElseThrow();
+
+            CompletableFuture<Optional<ImageRenderer>> imageFuture =
+                    CompletableFuture.supplyAsync(
+                            () -> ImageRendererManager.safelyPrepareFactory(
+                                    location, imageFactory
+                            ),
+                            backgroundExecutor
+                    )
+                            .thenCompose(preparationBarrier::wait)
+                            .thenApplyAsync(imageSupplierOpt -> {
+                                if (imageSupplierOpt.isEmpty()) {
+                                    return Optional.empty();
+                                }
+                                ImageRendererFactory.ImageSupplier supplier = imageSupplierOpt.get();
+
+                                ImageRenderer imageRenderer;
+                                try {
+                                    imageRenderer = supplier.completeImage();
+                                } catch (Exception e) {
+                                    YACLConstants.LOGGER.error("Failed to create image '{}'", location, e);
+                                    return Optional.empty();
+                                }
+
+                                ImageRendererManager.PRELOADED_IMAGE_CACHE.put(location, imageRenderer);
+
+                                return Optional.of(imageRenderer);
+                            }, gameExecutor);
+
+            futures.add(imageFuture);
+
+            imageFuture.whenComplete((result, throwable) -> {
+                if (throwable != null) {
+                    CrashReport crashReport = CrashReport.forThrowable(throwable, "Failed to load image");
+                    CrashReportCategory category = crashReport.addCategory("YACL Gui");
+                    category.setDetail("Image identifier", location.toString());
+                    throw new ReportedException(crashReport);
+                }
+            });
+        }
+
+        return CompletableFuture.allOf(futures.toArray(CompletableFuture[]::new));
+    }
+}
diff --git a/common/src/main/java/dev/isxander/yacl3/gui/image/impl/AnimatedDynamicTextureImage.java b/common/src/main/java/dev/isxander/yacl3/gui/image/impl/AnimatedDynamicTextureImage.java
index cc5a2c3..39ddb55 100644
--- a/common/src/main/java/dev/isxander/yacl3/gui/image/impl/AnimatedDynamicTextureImage.java
+++ b/common/src/main/java/dev/isxander/yacl3/gui/image/impl/AnimatedDynamicTextureImage.java
@@ -94,7 +94,7 @@ public class AnimatedDynamicTextureImage extends DynamicTextureImage {
         return targetHeight;
     }
 
-    public static ImageRendererFactory<AnimatedDynamicTextureImage> createGIFFromTexture(ResourceLocation textureLocation) {
+    public static ImageRendererFactory createGIFFromTexture(ResourceLocation textureLocation) {
         return () -> {
             ResourceManager resourceManager = Minecraft.getInstance().getResourceManager();
             Resource resource = resourceManager.getResource(textureLocation).orElseThrow();
@@ -103,11 +103,11 @@ public class AnimatedDynamicTextureImage extends DynamicTextureImage {
         };
     }
 
-    public static ImageRendererFactory<AnimatedDynamicTextureImage> createGIFFromPath(Path path, ResourceLocation uniqueLocation) {
+    public static ImageRendererFactory createGIFFromPath(Path path, ResourceLocation uniqueLocation) {
         return () -> createGIFSupplier(new FileInputStream(path.toFile()), uniqueLocation);
     }
 
-    public static ImageRendererFactory<AnimatedDynamicTextureImage> createWEBPFromTexture(ResourceLocation textureLocation) {
+    public static ImageRendererFactory createWEBPFromTexture(ResourceLocation textureLocation) {
         return () -> {
             ResourceManager resourceManager = Minecraft.getInstance().getResourceManager();
             Resource resource = resourceManager.getResource(textureLocation).orElseThrow();
@@ -116,11 +116,11 @@ public class AnimatedDynamicTextureImage extends DynamicTextureImage {
         };
     }
 
-    public static ImageRendererFactory<AnimatedDynamicTextureImage> createWEBPFromPath(Path path, ResourceLocation uniqueLocation) {
+    public static ImageRendererFactory createWEBPFromPath(Path path, ResourceLocation uniqueLocation) {
         return () -> createWEBPSupplier(new FileInputStream(path.toFile()), uniqueLocation);
     }
 
-    private static ImageRendererFactory.ImageSupplier<AnimatedDynamicTextureImage> createGIFSupplier(InputStream is, ResourceLocation uniqueLocation) {
+    private static ImageRendererFactory.ImageSupplier createGIFSupplier(InputStream is, ResourceLocation uniqueLocation) {
         try (is) {
             ImageReader reader = ImageIO.getImageReadersBySuffix("gif").next();
             reader.setInput(ImageIO.createImageInputStream(is));
@@ -144,7 +144,7 @@ public class AnimatedDynamicTextureImage extends DynamicTextureImage {
         }
     }
 
-    private static ImageRendererFactory.ImageSupplier<AnimatedDynamicTextureImage> createWEBPSupplier(InputStream is, ResourceLocation uniqueLocation) {
+    private static ImageRendererFactory.ImageSupplier createWEBPSupplier(InputStream is, ResourceLocation uniqueLocation) {
         try (is) {
             ImageReader reader = new WebPImageReaderSpi().createReaderInstance();
             reader.setInput(ImageIO.createImageInputStream(is));
@@ -180,7 +180,7 @@ public class AnimatedDynamicTextureImage extends DynamicTextureImage {
         }
     }
 
-    private static ImageRendererFactory.ImageSupplier<AnimatedDynamicTextureImage> createFromImageReader(ImageReader reader, AnimFrameProvider animationProvider, ResourceLocation uniqueLocation) throws Exception {
+    private static ImageRendererFactory.ImageSupplier createFromImageReader(ImageReader reader, AnimFrameProvider animationProvider, ResourceLocation uniqueLocation) throws Exception {
         if (reader.isSeekForwardOnly()) {
             throw new RuntimeException("Image reader is not seekable");
         }
diff --git a/common/src/main/java/dev/isxander/yacl3/gui/image/impl/DynamicTextureImage.java b/common/src/main/java/dev/isxander/yacl3/gui/image/impl/DynamicTextureImage.java
index 86dc833..2d2abb9 100644
--- a/common/src/main/java/dev/isxander/yacl3/gui/image/impl/DynamicTextureImage.java
+++ b/common/src/main/java/dev/isxander/yacl3/gui/image/impl/DynamicTextureImage.java
@@ -66,7 +66,7 @@ public class DynamicTextureImage implements ImageRenderer {
         textureManager.release(uniqueLocation);
     }
 
-    public static ImageRendererFactory<DynamicTextureImage> fromPath(Path imagePath, ResourceLocation location) {
-        return (ImageRendererFactory.OnThread<DynamicTextureImage>) () -> () -> new DynamicTextureImage(NativeImage.read(new FileInputStream(imagePath.toFile())), location);
+    public static ImageRendererFactory fromPath(Path imagePath, ResourceLocation location) {
+        return (ImageRendererFactory.OnThread) () -> () -> new DynamicTextureImage(NativeImage.read(new FileInputStream(imagePath.toFile())), location);
     }
 }
diff --git a/common/src/main/java/dev/isxander/yacl3/gui/image/impl/ResourceTextureImage.java b/common/src/main/java/dev/isxander/yacl3/gui/image/impl/ResourceTextureImage.java
index 6805611..abbeec7 100644
--- a/common/src/main/java/dev/isxander/yacl3/gui/image/impl/ResourceTextureImage.java
+++ b/common/src/main/java/dev/isxander/yacl3/gui/image/impl/ResourceTextureImage.java
@@ -50,7 +50,7 @@ public class ResourceTextureImage implements ImageRenderer {
 
     }
 
-    public static ImageRendererFactory<ResourceTextureImage> createFactory(ResourceLocation location, float u, float v, int width, int height, int textureWidth, int textureHeight) {
-        return (ImageRendererFactory.OnThread<ResourceTextureImage>) () -> () -> new ResourceTextureImage(location, u, v, width, height, textureWidth, textureHeight);
+    public static ImageRendererFactory createFactory(ResourceLocation location, float u, float v, int width, int height, int textureWidth, int textureHeight) {
+        return (ImageRendererFactory.OnThread) () -> () -> new ResourceTextureImage(location, u, v, width, height, textureWidth, textureHeight);
     }
 }
diff --git a/common/src/main/java/dev/isxander/yacl3/impl/OptionDescriptionImpl.java b/common/src/main/java/dev/isxander/yacl3/impl/OptionDescriptionImpl.java
index 3a3008c..67fa6a6 100644
--- a/common/src/main/java/dev/isxander/yacl3/impl/OptionDescriptionImpl.java
+++ b/common/src/main/java/dev/isxander/yacl3/impl/OptionDescriptionImpl.java
@@ -88,7 +88,13 @@ public record OptionDescriptionImpl(Component text, CompletableFuture<Optional<I
         public Builder webpImage(ResourceLocation image) {
             Validate.isTrue(imageUnset, "Image already set!");
 
-            this.image = ImageRendererManager.registerImage(image, AnimatedDynamicTextureImage.createWEBPFromTexture(image)).thenApply(Optional::of);
+            Optional<ImageRenderer> completedImage = ImageRendererManager.getImage(image);
+            if (completedImage.isPresent()) {
+                this.image = CompletableFuture.completedFuture(completedImage);
+            } else {
+                this.image = ImageRendererManager.registerImage(image, AnimatedDynamicTextureImage.createWEBPFromTexture(image)).thenApply(Optional::of);
+            }
+
             imageUnset = false;
             return this;
         }