Get dynamic model loading working with ModifyBakingResult

Author: embeddedt

common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/dynamic_resources/ModelManagerMixin.java

@@ -1,40 +1,41 @@
 package org.embeddedt.modernfix.common.mixin.perf.dynamic_resources;
 
+import com.google.common.collect.Maps;
+import com.llamalad7.mixinextras.sugar.Local;
 import net.minecraft.client.Minecraft;
 import net.minecraft.client.model.geom.EntityModelSet;
 import net.minecraft.client.renderer.block.model.BlockModel;
 import net.minecraft.client.renderer.item.ClientItem;
 import net.minecraft.client.renderer.item.ItemModel;
+import net.minecraft.client.resources.model.AtlasSet;
 import net.minecraft.client.resources.model.BakedModel;
 import net.minecraft.client.resources.model.BlockStateModelLoader;
 import net.minecraft.client.resources.model.ModelManager;
 import net.minecraft.client.resources.model.ModelResourceLocation;
 import net.minecraft.client.resources.model.UnbakedModel;
 import net.minecraft.resources.ResourceLocation;
 import net.minecraft.server.packs.resources.ResourceManager;
-import net.minecraft.util.profiling.ProfilerFiller;
 import net.minecraft.world.level.block.state.BlockState;
+import org.apache.commons.lang3.ArrayUtils;
 import org.embeddedt.modernfix.annotation.ClientOnlyMixin;
 import org.embeddedt.modernfix.dynamicresources.DynamicModelProvider;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Overwrite;
 import org.spongepowered.asm.mixin.Shadow;
 import org.spongepowered.asm.mixin.Unique;
 import org.spongepowered.asm.mixin.injection.At;
-import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.ModifyArg;
 import org.spongepowered.asm.mixin.injection.Redirect;
-import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 
+import java.lang.ref.WeakReference;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 
 @Mixin(ModelManager.class)
 @ClientOnlyMixin
-public class ModelManagerMixin {
+public class ModelManagerMixin implements DynamicModelProvider.ModelManagerExtension {
     @Shadow private BakedModel missingModel;
-    @Shadow private ItemModel missingItemModel;
-    @Shadow private EntityModelSet entityModelSet;
     @Unique
     private DynamicModelProvider mfix$modelProvider;
 
@@ -57,16 +58,19 @@ private static Map<BlockState, BakedModel> createBlockStateToModelDispatch(Map<M
         return Map.of();
     }
 
-    @Inject(method = "apply", at = @At("RETURN"))
-    private void createModelProvider(ModelManager.ReloadState reloadState, ProfilerFiller profiler, CallbackInfo ci) {
-        this.mfix$modelProvider = new DynamicModelProvider(
-                null, // TODO
-                this.missingModel,
-                this.missingItemModel,
-                Minecraft.getInstance().getResourceManager(),
-                this.entityModelSet,
-                reloadState.atlasPreparations()
-        );
+    @ModifyArg(method = "reload", at = @At(value = "INVOKE", target = "Ljava/util/concurrent/CompletableFuture;allOf([Ljava/util/concurrent/CompletableFuture;)Ljava/util/concurrent/CompletableFuture;", ordinal = 1))
+    private CompletableFuture<?>[] createModelProvider(CompletableFuture<?>[] cfs, @Local(ordinal = 0) CompletableFuture<EntityModelSet> entityModelFuture, @Local(ordinal = 0, argsOnly = true) Executor executor, @Local(ordinal = 0) Map<ResourceLocation, CompletableFuture<AtlasSet.StitchResult>> atlasPreparations) {
+        CompletableFuture<Void> makeModelProviderFuture = CompletableFuture.supplyAsync(() -> {
+            return Map.copyOf(Maps.transformValues(atlasPreparations, CompletableFuture::join));
+        }, executor).thenAcceptBoth(entityModelFuture, (stitchResults, entityModelSet) -> {
+            this.mfix$modelProvider = new DynamicModelProvider(
+                    Minecraft.getInstance().getResourceManager(),
+                    entityModelSet,
+                    stitchResults
+            );
+            DynamicModelProvider.currentReloadingModelProvider = new WeakReference<>(this.mfix$modelProvider);
+        });
+        return ArrayUtils.add(cfs, makeModelProviderFuture);
     }
 
     /**
@@ -99,4 +103,9 @@ public ItemModel getItemModel(ResourceLocation resourceLocation) {
     public ClientItem.Properties getItemProperties(ResourceLocation resourceLocation) {
         return this.mfix$modelProvider.getClientItemProperties(resourceLocation);
     }
+
+    @Override
+    public DynamicModelProvider mfix$getModelProvider() {
+        return this.mfix$modelProvider;
+    }
 }