Parallelize parts of model loading/baking

src/main/java/org/embeddedt/modernfix/mixin/ModelBakeryMixin.java

@@ -0,0 +1,125 @@
+package org.embeddedt.modernfix.mixin;
+
+import com.google.common.base.Stopwatch;
+import com.google.common.cache.LoadingCache;
+import net.minecraft.client.Minecraft;
+import net.minecraft.client.renderer.BlockModelShapes;
+import net.minecraft.client.renderer.model.*;
+import net.minecraft.client.renderer.texture.SpriteMap;
+import net.minecraft.profiler.IProfiler;
+import net.minecraft.util.ResourceLocation;
+import net.minecraft.util.Util;
+import net.minecraft.util.math.vector.TransformationMatrix;
+import org.apache.commons.lang3.tuple.Pair;
+import org.apache.commons.lang3.tuple.Triple;
+import org.embeddedt.modernfix.ModernFix;
+import org.spongepowered.asm.mixin.Final;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Shadow;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.Redirect;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
+
+import javax.annotation.Nullable;
+import java.io.IOException;
+import java.util.*;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+@Mixin(ModelBakery.class)
+public abstract class ModelBakeryMixin {
+    @Shadow protected abstract BlockModel loadModel(ResourceLocation location) throws IOException;
+
+    @Shadow @Final private Map<ResourceLocation, IUnbakedModel> unbakedModels;
+    @Shadow @Final public static ModelResourceLocation MODEL_MISSING;
+
+    @Shadow public abstract IUnbakedModel getUnbakedModel(ResourceLocation modelLocation);
+
+    @Shadow @Final public static BlockModel MODEL_GENERATED;
+    @Shadow @Final private static ItemModelGenerator ITEM_MODEL_GENERATOR;
+    @Shadow @Nullable private SpriteMap spriteMap;
+    @Shadow @Final private Map<Triple<ResourceLocation, TransformationMatrix, Boolean>, IBakedModel> bakedModels;
+    @Shadow @Final private Map<ResourceLocation, IBakedModel> topBakedModels;
+    private Map<ResourceLocation, BlockModel> deserializedModelCache = null;
+    private boolean useModelCache = false;
+
+    @Inject(method = "loadModel", at = @At("HEAD"), cancellable = true)
+    private void useCachedModel(ResourceLocation location, CallbackInfoReturnable<BlockModel> cir) {
+        if(useModelCache && deserializedModelCache != null) {
+            BlockModel model = deserializedModelCache.get(location);
+            if(model != null)
+                cir.setReturnValue(model);
+        }
+    }
+
+    private BlockModel loadModelSafely(ResourceLocation location) {
+        try {
+            return this.loadModel(location);
+        } catch(Throwable e) {
+            ModernFix.LOGGER.warn("Model " + location + " will not be preloaded", e);
+            return null;
+        }
+    }
+
+    @Inject(method = "processLoading", at = @At(value = "INVOKE", target = "Lnet/minecraft/profiler/IProfiler;endStartSection(Ljava/lang/String;)V", ordinal = 1))
+    private void preloadJsonModels(IProfiler profilerIn, int maxMipmapLevel, CallbackInfo ci) {
+        profilerIn.endStartSection("loadjsons");
+        ModernFix.LOGGER.warn("Preloading JSONs in parallel...");
+        Stopwatch stopwatch = Stopwatch.createStarted();
+        useModelCache = false;
+        deserializedModelCache = Minecraft.getInstance().getResourceManager().getAllResourceLocations("models", p -> p.endsWith(".json"))
+                .parallelStream()
+                .map(location -> new ResourceLocation(location.getNamespace(), location.getPath().substring(7, location.getPath().length() - 5)))
+                .map(location -> Pair.of(location, this.loadModelSafely(location)))
+                .filter(pair -> pair.getValue() != null)
+                .collect(Collectors.toConcurrentMap(Pair::getKey, Pair::getValue));
+        useModelCache = true;
+        ModernFix.LOGGER.warn("Preloading JSONs took " + stopwatch.elapsed(TimeUnit.SECONDS) + " seconds");
+        stopwatch.stop();
+    }
+
+    @Inject(method = "processLoading", at = @At("RETURN"), remap = false)
+    private void clearModelCache(IProfiler profilerIn, int maxMipmapLevel, CallbackInfo ci) {
+        deserializedModelCache.clear();
+        useModelCache = false;
+    }
+
+    @Redirect(method = "uploadTextures", at = @At(value = "INVOKE", target = "Ljava/util/Set;forEach(Ljava/util/function/Consumer;)V", ordinal = 0))
+    private void parallelBake(Set<ResourceLocation> locationSet, Consumer consumer) {
+        final IModelTransform transform = ModelRotation.X0_Y0;
+        if(this.spriteMap == null)
+            throw new IllegalStateException("no sprite map");
+        ModernFix.LOGGER.warn("Baking models in parallel...");
+        Stopwatch stopwatch = Stopwatch.createStarted();
+        locationSet.forEach(this::getUnbakedModel); /* make sure every unbaked model is loaded, should be fast */
+        List<Pair<ResourceLocation, IBakedModel>> models = CompletableFuture.supplyAsync(() -> {
+            return locationSet.parallelStream().map(location -> {
+                IUnbakedModel iunbakedmodel = this.unbakedModels.get(location);
+                if (iunbakedmodel instanceof BlockModel) {
+                    BlockModel blockmodel = (BlockModel)iunbakedmodel;
+                    if (blockmodel.getRootModel() == MODEL_GENERATED) {
+                        return Pair.of(location, ITEM_MODEL_GENERATOR.makeItemModel(this.spriteMap::getSprite, blockmodel).bakeModel((ModelBakery)(Object)this, blockmodel, this.spriteMap::getSprite, transform, location, false));
+                    }
+                }
+
+                IBakedModel ibakedmodel = iunbakedmodel.bakeModel((ModelBakery)(Object)this, this.spriteMap::getSprite, transform, location);
+                return Pair.of(location, ibakedmodel);
+            }).collect(Collectors.toList());
+        }, Util.getServerExecutor()).join();
+        models.forEach(pair -> {
+            Triple<ResourceLocation, TransformationMatrix, Boolean> triple = Triple.of(pair.getKey(), transform.getRotation(), transform.isUvLock());
+            this.bakedModels.put(triple, pair.getValue());
+            if(pair.getValue() != null)
+                this.topBakedModels.put(pair.getKey(), pair.getValue());
+        });
+        ModernFix.LOGGER.warn("Baking in parallel took " + stopwatch.elapsed(TimeUnit.SECONDS) + " seconds");
+        stopwatch.stop();
+    }
+}

src/main/resources/modernfix.mixins.json

@@ -13,7 +13,8 @@
   ],
   "client": [
     "MinecraftMixin",
-    "RenderTypeMixin"
+    "RenderTypeMixin",
+    "ModelBakeryMixin"
   ],
   "injectors": {
     "defaultRequire": 1