Make model discovery dynamic

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

@@ -1,5 +1,7 @@
 package org.embeddedt.modernfix.common.mixin.perf.dynamic_resources;
 
+import net.minecraft.client.resources.model.BlockStateModelLoader;
+import net.minecraft.client.resources.model.ClientItemInfoLoader;
 import net.minecraft.client.resources.model.ModelManager;
 import net.minecraft.client.resources.model.UnbakedModel;
 import net.minecraft.resources.Identifier;
@@ -7,6 +9,7 @@ import net.minecraft.server.packs.resources.Resource;
 import org.embeddedt.modernfix.annotation.ClientOnlyMixin;
 import org.embeddedt.modernfix.dynresources.DynamicModelSystem;
 import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Overwrite;
 import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.ModifyArg;
 
@@ -27,4 +30,15 @@ public class MixinModelManager {
     private static Function<Map<Identifier, Resource>, ? extends CompletionStage<Map<Identifier, UnbakedModel>>> skipAOTUnbakedModelLoad(Function<Map<Identifier, Resource>, ? extends CompletionStage<Map<Identifier, UnbakedModel>>> original) {
         return resourceMap -> CompletableFuture.completedFuture(DynamicModelSystem.createDynamicUnbakedModelMap(resourceMap));
     }
+
+    /**
+     * @author embeddedt
+     * @reason Model resolution is not necessary, because we dynamically find models as they are referenced
+     */
+    @Overwrite
+    private static ModelManager.ResolvedModels discoverModelDependencies(
+            Map<Identifier, UnbakedModel> inputModels, BlockStateModelLoader.LoadedModels loadedModels, ClientItemInfoLoader.LoadedClientInfos loadedClientInfos, net.neoforged.neoforge.client.model.standalone.StandaloneModelLoader.LoadedModels standaloneModels
+    ) {
+        return new DynamicModelSystem.DynamicResolver(inputModels, loadedModels, loadedClientInfos, standaloneModels).resolvedModels();
+    }
 }

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

@@ -0,0 +1,15 @@
+package org.embeddedt.modernfix.common.mixin.perf.dynamic_resources;
+
+import it.unimi.dsi.fastutil.objects.Object2ObjectMap;
+import net.minecraft.client.resources.model.ModelDiscovery;
+import net.minecraft.resources.Identifier;
+import org.embeddedt.modernfix.annotation.ClientOnlyMixin;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.gen.Accessor;
+
+@Mixin(ModelDiscovery.class)
+@ClientOnlyMixin
+public interface ModelDiscoveryAccessor {
+    @Accessor("modelWrappers")
+    Object2ObjectMap<Identifier, ModelDiscovery.ModelWrapper> mfix$getModelWrappers();
+}

src/main/java/org/embeddedt/modernfix/dynresources/DynamicModelSystem.java

@@ -4,7 +4,14 @@ import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
 import com.google.common.collect.Maps;
+import net.minecraft.client.renderer.block.model.ItemModelGenerator;
+import net.minecraft.client.renderer.item.ItemModel;
 import net.minecraft.client.resources.model.BlockStateModelLoader;
+import net.minecraft.client.resources.model.ClientItemInfoLoader;
+import net.minecraft.client.resources.model.MissingBlockModel;
+import net.minecraft.client.resources.model.ModelDiscovery;
+import net.minecraft.client.resources.model.ModelManager;
+import net.minecraft.client.resources.model.ResolvedModel;
 import net.minecraft.client.resources.model.UnbakedModel;
 import net.minecraft.resources.FileToIdConverter;
 import net.minecraft.resources.Identifier;
@@ -12,8 +19,10 @@ import net.minecraft.server.packs.resources.Resource;
 import net.minecraft.world.level.block.Block;
 import net.minecraft.world.level.block.state.BlockState;
 import net.neoforged.neoforge.client.model.UnbakedModelParser;
+import net.neoforged.neoforge.client.model.standalone.StandaloneModelLoader;
 import org.embeddedt.modernfix.ModernFix;
 import org.embeddedt.modernfix.common.mixin.perf.dynamic_resources.IdMapperAccessor;
+import org.embeddedt.modernfix.common.mixin.perf.dynamic_resources.ModelDiscoveryAccessor;
 
 import java.io.Reader;
 import java.util.List;
@@ -70,4 +79,37 @@ public class DynamicModelSystem {
             return loadedModels.models().get(state);
         }));
     }
+
+    public record DynamicResolver(Map<Identifier, UnbakedModel> inputModels,
+                                          BlockStateModelLoader.LoadedModels loadedModels,
+                                          ClientItemInfoLoader.LoadedClientInfos loadedClientInfos,
+                                          StandaloneModelLoader.LoadedModels standaloneModels) {
+
+        private ResolvedModel resolveModel(Identifier id) {
+            var discovery = new ModelDiscovery(inputModels, MissingBlockModel.missingModel());
+            discovery.addSpecialModel(ItemModelGenerator.GENERATED_ITEM_MODEL_ID, new ItemModelGenerator());
+            if (!id.equals(ItemModelGenerator.GENERATED_ITEM_MODEL_ID)) {
+                UnbakedModel unbaked = inputModels.get(id);
+                if (unbaked != null) {
+                    var wrapper = discovery.createAndQueueWrapper(id, unbaked);
+                    ((ModelDiscoveryAccessor)discovery).mfix$getModelWrappers().put(id, wrapper);
+                } else {
+                    ModernFix.LOGGER.warn("Cannot find the root model for {}", id);
+                }
+            }
+            var resolved = discovery.resolve();
+            return resolved.getOrDefault(id, discovery.missingModel());
+        }
+
+        public ModelManager.ResolvedModels resolvedModels() {
+            var resolvedMissingModel = new ModelDiscovery(inputModels, MissingBlockModel.missingModel()).missingModel();
+            LoadingCache<Identifier, ResolvedModel> resolvedModelCache = CacheBuilder.newBuilder().softValues().maximumSize(1000).build(new CacheLoader<>() {
+                @Override
+                public ResolvedModel load(Identifier key) {
+                    return resolveModel(key);
+                }
+            });
+            return new ModelManager.ResolvedModels(resolvedMissingModel, Maps.asMap(inputModels.keySet(), resolvedModelCache::getUnchecked));
+        }
+    }
 }

src/main/resources/META-INF/accesstransformer.cfg

@@ -54,4 +54,8 @@ public net.minecraft.world.level.levelgen.DensityFunctions$MulOrAdd$Type
 public net.minecraft.client.renderer.entity.EnderDragonRenderer$DragonModel entity
 public net.minecraft.client.KeyMapping ALL
 
-public net.minecraft.client.resources.model.BlockStateModelLoader$LoadedBlockModelDefinition
+public net.minecraft.client.resources.model.BlockStateModelLoader$LoadedBlockModelDefinition
+public net.minecraft.client.resources.model.ModelManager$ResolvedModels
+public net.minecraft.client.resources.model.ModelDiscovery$ModelWrapper
+public net.minecraft.client.resources.model.ModelDiscovery$ModelWrapper ModelWrapper(Lnet/minecraft/resources/Identifier;Lnet/minecraft/client/resources/model/UnbakedModel;Z)V
+public net.minecraft.client.resources.model.ModelDiscovery createAndQueueWrapper(Lnet/minecraft/resources/Identifier;Lnet/minecraft/client/resources/model/UnbakedModel;)Lnet/minecraft/client/resources/model/ModelDiscovery$ModelWrapper;