Fire model bake events on worker thread

src/main/java/org/embeddedt/modernfix/duck/IExtendedModelBakery.java

@@ -0,0 +1,8 @@
+package org.embeddedt.modernfix.duck;
+
+
+import net.minecraft.client.renderer.texture.SpriteMap;
+
+public interface IExtendedModelBakery {
+    SpriteMap getUnfinishedAtlasSet();
+}

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

@@ -18,6 +18,7 @@ import org.apache.commons.lang3.tuple.Triple;
 import org.apache.logging.log4j.Logger;
 import org.embeddedt.modernfix.ModernFix;
 import org.embeddedt.modernfix.ModernFixClient;
+import org.embeddedt.modernfix.duck.IExtendedModelBakery;
 import org.embeddedt.modernfix.models.LazyBakedModel;
 import org.spongepowered.asm.mixin.*;
 import org.spongepowered.asm.mixin.injection.At;
@@ -34,7 +35,7 @@ import java.util.stream.Collectors;
 import static org.embeddedt.modernfix.ModernFix.LOGGER;
 
 @Mixin(ModelBakery.class)
-public abstract class ModelBakeryMixin {
+public abstract class ModelBakeryMixin implements IExtendedModelBakery {
     @Shadow @Final private Map<ResourceLocation, IUnbakedModel> topLevelModels;
 
     @Shadow @Final private Map<ResourceLocation, IBakedModel> bakedTopLevelModels;
@@ -82,7 +83,6 @@ public abstract class ModelBakeryMixin {
                 }
             }));
         });
-        this.atlasSet = null;
     }
 
     /**
@@ -92,7 +92,6 @@ public abstract class ModelBakeryMixin {
     @Overwrite
     public SpriteMap uploadTextures(TextureManager pResourceManager, IProfiler pProfiler) {
         pProfiler.push("atlas_upload");
-
         for(Pair<AtlasTexture, AtlasTexture.SheetData> pair : this.atlasPreparations.values()) {
             AtlasTexture atlastexture = pair.getFirst();
             AtlasTexture.SheetData atlastexture$sheetdata = pair.getSecond();
@@ -100,8 +99,12 @@ public abstract class ModelBakeryMixin {
             pResourceManager.bind(atlastexture.location());
             atlastexture.updateFilter(atlastexture$sheetdata);
         }
-        this.atlasSet = new SpriteMap(this.atlasPreparations.values().stream().map(Pair::getFirst).collect(Collectors.toList()));
         pProfiler.pop();
         return this.atlasSet;
     }
+
+    @Override
+    public SpriteMap getUnfinishedAtlasSet() {
+        return this.atlasSet;
+    }
 }

src/main/java/org/embeddedt/modernfix/mixin/perf/faster_baking/ModelManagerMixin.java

@@ -1,21 +1,77 @@
 package org.embeddedt.modernfix.mixin.perf.faster_baking;
 
+import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import net.minecraft.block.BlockState;
+import net.minecraft.client.Minecraft;
+import net.minecraft.client.renderer.BlockModelShapes;
+import net.minecraft.client.renderer.model.IBakedModel;
 import net.minecraft.client.renderer.model.ModelBakery;
 import net.minecraft.client.renderer.model.ModelManager;
+import net.minecraft.client.renderer.texture.SpriteMap;
+import net.minecraft.client.renderer.texture.TextureManager;
 import net.minecraft.profiler.IProfiler;
 import net.minecraft.resources.IResourceManager;
+import net.minecraft.util.ResourceLocation;
+import net.minecraftforge.client.model.ModelLoader;
+import net.minecraftforge.fml.common.ObfuscationReflectionHelper;
+import org.embeddedt.modernfix.duck.IExtendedModelBakery;
 import org.embeddedt.modernfix.models.LazyBakedModel;
+import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Overwrite;
+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.callback.CallbackInfo;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
+import org.spongepowered.asm.mixin.injection.callback.LocalCapture;
+
+import javax.annotation.Nullable;
+import java.util.Map;
 
 @Mixin(ModelManager.class)
 public class ModelManagerMixin {
-    @Inject(method = "apply(Lnet/minecraft/client/renderer/model/ModelBakery;Lnet/minecraft/resources/IResourceManager;Lnet/minecraft/profiler/IProfiler;)V",
+    @Shadow @Nullable private SpriteMap atlases;
-            at = @At(value = "INVOKE", target = "Lnet/minecraftforge/client/ForgeHooksClient;onModelBake(Lnet/minecraft/client/renderer/model/ModelManager;Ljava/util/Map;Lnet/minecraftforge/client/model/ModelLoader;)V", shift = At.Shift.AFTER)
+
-    )
+    @Shadow private Map<ResourceLocation, IBakedModel> bakedRegistry;
-    private void allowBaking(ModelBakery pObject, IResourceManager pResourceManager, IProfiler pProfiler, CallbackInfo ci) {
+
+    @Shadow private Object2IntMap<BlockState> modelGroups;
+
+    @Shadow @Final private TextureManager textureManager;
+
+    @Shadow private IBakedModel missingModel;
+
+    @Shadow @Final private BlockModelShapes blockModelShaper;
+
+    @Inject(method = "prepare(Lnet/minecraft/resources/IResourceManager;Lnet/minecraft/profiler/IProfiler;)Lnet/minecraft/client/renderer/model/ModelBakery;", at = @At(value = "INVOKE", target = "Lnet/minecraft/profiler/IProfiler;endTick()V"), locals = LocalCapture.CAPTURE_FAILHARD)
+    private void fireModelBakeEvent(IResourceManager pResourceManager, IProfiler pProfiler, CallbackInfoReturnable<ModelBakery> cir, ModelLoader pObject) {
+        pProfiler.push("modelevent");
+        if (this.atlases != null) {
+            Minecraft.getInstance().executeBlocking(() -> {
+                this.atlases.close();
+            });
+        }
+        this.atlases = ((IExtendedModelBakery)(Object)pObject).getUnfinishedAtlasSet();
+        this.bakedRegistry = pObject.getBakedTopLevelModels();
+        this.modelGroups = pObject.getModelGroups();
+        this.missingModel = this.bakedRegistry.get(ModelBakery.MISSING_MODEL_LOCATION);
+        net.minecraftforge.client.ForgeHooksClient.onModelBake((ModelManager)(Object)this, this.bakedRegistry, pObject);
+        pProfiler.popPush("cache");
+        this.blockModelShaper.rebuildCache();
+        pProfiler.pop();
+    }
+
+    /**
+     * @author embeddedt
+     * @reason most of the code is moved to prepare()
+     */
+    @Overwrite
+    protected void apply(ModelBakery pObject, IResourceManager pResourceManager, IProfiler pProfiler) {
+        pProfiler.startTick();
+        pProfiler.push("upload");
+        this.atlases = pObject.uploadTextures(this.textureManager, pProfiler);
+        pProfiler.pop();
         LazyBakedModel.allowBakeForFlags = true;
+        pProfiler.endTick();
     }
 }