Fix preparation barrier being waited on for every image (breaking reload listener ordering) (close #173)

1 files changed, 106 insertions, 70 deletions

diff --git a/src/main/java/dev/isxander/yacl3/gui/image/YACLImageReloadListener.java b/src/main/java/dev/isxander/yacl3/gui/image/YACLImageReloadListener.java
index 243f22f..fb0695c 100644
--- a/src/main/java/dev/isxander/yacl3/gui/image/YACLImageReloadListener.java
+++ b/src/main/java/dev/isxander/yacl3/gui/image/YACLImageReloadListener.java
@@ -2,9 +2,6 @@ package dev.isxander.yacl3.gui.image;
 
 import dev.isxander.yacl3.impl.utils.YACLConstants;
 import dev.isxander.yacl3.platform.YACLPlatform;
-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;
@@ -12,12 +9,15 @@ import net.minecraft.server.packs.resources.ResourceManager;
 import net.minecraft.util.profiling.ProfilerFiller;
 import org.jetbrains.annotations.NotNull;
 
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
+import java.util.*;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
+import java.util.function.BiConsumer;
+import java.util.function.BinaryOperator;
+import java.util.function.Function;
+import java.util.function.Supplier;
+import java.util.stream.Collector;
+import java.util.stream.Stream;
 
 public class YACLImageReloadListener
         implements PreparableReloadListener
@@ -27,80 +27,75 @@ public class YACLImageReloadListener
 {
     @Override
     public @NotNull CompletableFuture<Void> reload(
-            PreparationBarrier preparationBarrier, 
-            ResourceManager resourceManager, 
-            ProfilerFiller preparationsProfiler, 
-            ProfilerFiller reloadProfiler, 
-            Executor backgroundExecutor, 
-            Executor gameExecutor
+            PreparationBarrier preparationBarrier,
+            @NotNull ResourceManager resourceManager,
+            @NotNull ProfilerFiller preparationsProfiler,
+            @NotNull ProfilerFiller reloadProfiler,
+            @NotNull Executor backgroundExecutor,
+            @NotNull Executor gameExecutor
     ) {
-        Map<ResourceLocation, Resource> imageResources = resourceManager.listResources(
+        return prepare(resourceManager, preparationsProfiler, backgroundExecutor)
+                .thenCompose(preparationBarrier::wait)
+                .thenCompose(suppliers -> apply(suppliers, reloadProfiler, gameExecutor));
+    }
+
+    private CompletableFuture<List<Optional<SupplierPreparation>>> prepare(
+            ResourceManager manager,
+            ProfilerFiller profiler,
+            Executor executor
+    ) {
+        Map<ResourceLocation, Resource> imageResources = manager.listResources(
                 "textures",
                 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();
+        return imageResources.keySet().stream()
+                .map(location -> {
+                    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(
+                    return 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);
-                }
-            });
-        }
+                            ).map(supplier -> new SupplierPreparation(location, supplier)),
+                            executor
+                    );
+                })
+                .collect(CompletableFutureCollector.allOf());
+    }
 
-        return CompletableFuture.allOf(futures.toArray(CompletableFuture[]::new));
+    private CompletableFuture<Void> apply(
+            List<Optional<SupplierPreparation>> suppliers,
+            ProfilerFiller profiler,
+            Executor executor
+    ) {
+        return CompletableFuture.allOf(suppliers.stream()
+                .flatMap(Optional::stream)
+                .map(prep -> CompletableFuture.supplyAsync(
+                        () -> {
+                            ImageRenderer imageRenderer;
+                            try {
+                                imageRenderer = prep.supplier().completeImage();
+                            } catch (Exception e) {
+                                YACLConstants.LOGGER.error("Failed to create image '{}'", prep.location(), e);
+                                return Optional.empty();
+                            }
+                            ImageRendererManager.PRELOADED_IMAGE_CACHE.put(prep.location(), imageRenderer);
+                            YACLConstants.LOGGER.info("Successfully loaded image '{}'", prep.location());
+                            return Optional.of(imageRenderer);
+                        },
+                        executor
+                ))
+                .toArray(CompletableFuture<?>[]::new));
+    }
+
+    private record SupplierPreparation(ResourceLocation location, ImageRendererFactory.ImageSupplier supplier) {
     }
 
     /*? if fabric {*/
@@ -109,4 +104,45 @@ public class YACLImageReloadListener
         return YACLPlatform.rl("image_reload_listener");
     }
     /*?}*/
+
+    public static class CompletableFutureCollector<X, T extends CompletableFuture<X>> implements Collector<T, List<T>, CompletableFuture<List<X>>> {
+        private CompletableFutureCollector() {
+        }
+
+        public static <X, T extends CompletableFuture<X>> Collector<T, List<T>, CompletableFuture<List<X>>> allOf() {
+            return new CompletableFutureCollector<>();
+        }
+
+        @Override
+        public Supplier<List<T>> supplier() {
+            return ArrayList::new;
+        }
+
+        @Override
+        public BiConsumer<List<T>, T> accumulator() {
+            return List::add;
+        }
+
+        @Override
+        public BinaryOperator<List<T>> combiner() {
+            return (left, right) -> {
+                left.addAll(right);
+                return left;
+            };
+        }
+
+        @Override
+        public Function<List<T>, CompletableFuture<List<X>>> finisher() {
+            return ls -> CompletableFuture.allOf(ls.toArray(CompletableFuture[]::new))
+                    .thenApply(v -> ls
+                            .stream()
+                            .map(CompletableFuture::join)
+                            .toList());
+        }
+
+        @Override
+        public Set<Characteristics> characteristics() {
+            return Collections.emptySet();
+        }
+    }
 }