Get dynamic model loading working with ModifyBakingResult

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

@@ -1,10 +1,13 @@
 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;
@@ -12,8 +15,8 @@ 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;
@@ -21,20 +24,18 @@ 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 @@ public class ModelManagerMixin {
         return Map.of();
     }
 
-    @Inject(method = "apply", at = @At("RETURN"))
+    @ModifyArg(method = "reload", at = @At(value = "INVOKE", target = "Ljava/util/concurrent/CompletableFuture;allOf([Ljava/util/concurrent/CompletableFuture;)Ljava/util/concurrent/CompletableFuture;", ordinal = 1))
-    private void createModelProvider(ModelManager.ReloadState reloadState, ProfilerFiller profiler, CallbackInfo ci) {
+    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) {
-        this.mfix$modelProvider = new DynamicModelProvider(
+        CompletableFuture<Void> makeModelProviderFuture = CompletableFuture.supplyAsync(() -> {
-                null, // TODO
+            return Map.copyOf(Maps.transformValues(atlasPreparations, CompletableFuture::join));
-                this.missingModel,
+        }, executor).thenAcceptBoth(entityModelFuture, (stitchResults, entityModelSet) -> {
-                this.missingItemModel,
+            this.mfix$modelProvider = new DynamicModelProvider(
-                Minecraft.getInstance().getResourceManager(),
+                    Minecraft.getInstance().getResourceManager(),
-                this.entityModelSet,
+                    entityModelSet,
-                reloadState.atlasPreparations()
+                    stitchResults
-        );
+            );
+            DynamicModelProvider.currentReloadingModelProvider = new WeakReference<>(this.mfix$modelProvider);
+        });
+        return ArrayUtils.add(cfs, makeModelProviderFuture);
     }
 
     /**
@@ -99,4 +103,9 @@ public class ModelManagerMixin {
     public ClientItem.Properties getItemProperties(ResourceLocation resourceLocation) {
         return this.mfix$modelProvider.getClientItemProperties(resourceLocation);
     }
+
+    @Override
+    public DynamicModelProvider mfix$getModelProvider() {
+        return this.mfix$modelProvider;
+    }
 }

common/src/main/java/org/embeddedt/modernfix/dynamicresources/DynamicModelProvider.java

@@ -3,17 +3,21 @@ package org.embeddedt.modernfix.dynamicresources;
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
+import com.google.common.collect.Iterators;
+import com.google.common.collect.Maps;
 import com.google.gson.JsonObject;
 import com.google.gson.JsonParser;
 import com.mojang.serialization.JsonOps;
 import it.unimi.dsi.fastutil.objects.ObjectOpenHashSet;
 import net.minecraft.client.model.geom.EntityModelSet;
+import net.minecraft.client.renderer.block.BlockModelShaper;
 import net.minecraft.client.renderer.block.model.BlockModel;
 import net.minecraft.client.renderer.block.model.BlockModelDefinition;
 import net.minecraft.client.renderer.block.model.ItemModelGenerator;
 import net.minecraft.client.renderer.block.model.UnbakedBlockStateModel;
 import net.minecraft.client.renderer.item.ClientItem;
 import net.minecraft.client.renderer.item.ItemModel;
+import net.minecraft.client.renderer.item.MissingItemModel;
 import net.minecraft.client.renderer.texture.TextureAtlas;
 import net.minecraft.client.renderer.texture.TextureAtlasSprite;
 import net.minecraft.client.resources.model.AtlasSet;
@@ -29,6 +33,7 @@ import net.minecraft.client.resources.model.ModelResourceLocation;
 import net.minecraft.client.resources.model.ModelState;
 import net.minecraft.client.resources.model.SpriteGetter;
 import net.minecraft.client.resources.model.UnbakedModel;
+import net.minecraft.core.registries.BuiltInRegistries;
 import net.minecraft.resources.ResourceLocation;
 import net.minecraft.server.packs.resources.Resource;
 import net.minecraft.server.packs.resources.ResourceManager;
@@ -37,14 +42,23 @@ import net.minecraft.world.level.block.Block;
 import net.minecraft.world.level.block.state.BlockState;
 import net.minecraft.world.level.block.state.StateDefinition;
 import org.embeddedt.modernfix.ModernFix;
+import org.jetbrains.annotations.NotNull;
 
 import java.io.Reader;
+import java.lang.ref.WeakReference;
+import java.util.AbstractSet;
 import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
@@ -53,69 +67,22 @@ import java.util.stream.Collectors;
  */
 public class DynamicModelProvider {
     private final LoadingCache<ResourceLocation, Optional<BlockStateModelLoader.LoadedModels>> loadedStateDefinitions =
-            CacheBuilder.newBuilder()
+            this.makeLoadingCache(this::loadBlockStateDefinition);
-                    .expireAfterAccess(3, TimeUnit.MINUTES)
-                    .maximumSize(1000)
-                    .concurrencyLevel(8)
-                    .softValues()
-                    .build(new CacheLoader<>() {
-                        @Override
-                        public Optional<BlockStateModelLoader.LoadedModels> load(ResourceLocation key) {
-                            return loadBlockStateDefinition(key);
-                        }
-                    });
 
     private final LoadingCache<ResourceLocation, Optional<UnbakedModel>> loadedBlockModels =
-            CacheBuilder.newBuilder()
+            this.makeLoadingCache(this::loadBlockModel);
-                    .expireAfterAccess(3, TimeUnit.MINUTES)
-                    .maximumSize(1000)
-                    .concurrencyLevel(8)
-                    .softValues()
-                    .build(new CacheLoader<>() {
-                        @Override
-                        public Optional<UnbakedModel> load(ResourceLocation key) {
-                            return loadBlockModel(key);
-                        }
-                    });
 
     private final LoadingCache<ModelResourceLocation, Optional<BakedModel>> loadedBakedModels =
-            CacheBuilder.newBuilder()
+            this.makeLoadingCache(this::loadBakedModel);
-                    .expireAfterAccess(3, TimeUnit.MINUTES)
-                    .maximumSize(1000)
-                    .concurrencyLevel(8)
-                    .softValues()
-                    .build(new CacheLoader<>() {
-                        @Override
-                        public Optional<BakedModel> load(ModelResourceLocation key) {
-                            return loadBakedModel(key);
-                        }
-                    });
 
     private final LoadingCache<ResourceLocation, Optional<ClientItem>> loadedClientItemProperties =
-            CacheBuilder.newBuilder()
+            this.makeLoadingCache(this::loadClientItemProperties);
-                    .expireAfterAccess(3, TimeUnit.MINUTES)
-                    .maximumSize(1000)
-                    .concurrencyLevel(8)
-                    .softValues()
-                    .build(new CacheLoader<>() {
-                        @Override
-                        public Optional<ClientItem> load(ResourceLocation key) {
-                            return loadClientItemProperties(key);
-                        }
-                    });
 
     private final LoadingCache<ResourceLocation, Optional<ItemModel>> loadedItemModels =
-            CacheBuilder.newBuilder()
+           this.makeLoadingCache(this::loadItemModel);
-                    .expireAfterAccess(3, TimeUnit.MINUTES)
+
-                    .maximumSize(1000)
+    private final LoadingCache<ResourceLocation, Optional<BakedModel>> loadedStandaloneModels =
-                    .concurrencyLevel(8)
+            this.makeLoadingCache(this::loadStandaloneModel);
-                    .softValues()
-                    .build(new CacheLoader<>() {
-                        @Override
-                        public Optional<ItemModel> load(ResourceLocation key) {
-                            return loadItemModel(key);
-                        }
-                    });
 
     private final BakedModel missingModel;
     private final ItemModel missingItemModel;
@@ -127,11 +94,13 @@ public class DynamicModelProvider {
     private final EntityModelSet entityModelSet;
     private final ItemModelGenerator itemModelGenerator;
 
-    public DynamicModelProvider(Map<ModelResourceLocation, BakedModel> initialBakedRegistry, BakedModel missingModel,
+    private final Map<ModelResourceLocation, BakedModel> mrlModelOverrides = new ConcurrentHashMap<>();
-                                ItemModel missingItemModel, ResourceManager resourceManager, EntityModelSet entityModelSet,
+    private final Map<ResourceLocation, ItemModel> itemStackModelOverrides = new ConcurrentHashMap<>();
+    private final Map<ResourceLocation, BakedModel> standaloneModelOverrides = new ConcurrentHashMap<>();
+
+    public DynamicModelProvider(ResourceManager resourceManager, EntityModelSet entityModelSet,
                                 Map<ResourceLocation, AtlasSet.StitchResult> atlasMap) {
-        this.missingModel = missingModel;
+        this.unbakedMissingModel = MissingBlockModel.missingModel();
-        this.missingItemModel = missingItemModel;
         this.entityModelSet = entityModelSet;
         var missing = atlasMap.get(TextureAtlas.LOCATION_BLOCKS).missing();
         this.textureGetter = new ModelBakery.TextureGetter() {
@@ -154,9 +123,189 @@ public class DynamicModelProvider {
         };
         this.stateMapper = BlockStateModelLoader.definitionLocationToBlockMapper();
         this.resourceManager = resourceManager;
-        this.unbakedMissingModel = MissingBlockModel.missingModel();
         this.resolver = new DynamicResolver();
         this.itemModelGenerator = new ItemModelGenerator();
+        this.missingModel = this.bakeModel(this.unbakedMissingModel, () -> "missing");
+        this.missingItemModel = new MissingItemModel(this.missingModel);
+    }
+
+    public BakedModel getMissingBakedModel() {
+        return this.missingModel;
+    }
+
+    public ItemModel getMissingItemModel() {
+        return this.missingItemModel;
+    }
+
+    public Map<ModelResourceLocation, BakedModel> getTopLevelEmulatedRegistry() {
+        Set<ModelResourceLocation> topLevelModelLocations = new HashSet<>();
+        // Skip going through ModelLocationCache because most of the accesses will be misses
+        BuiltInRegistries.BLOCK.entrySet().forEach(entry -> {
+            var location = entry.getKey().location();
+            for(BlockState state : entry.getValue().getStateDefinition().getPossibleStates()) {
+                topLevelModelLocations.add(BlockModelShaper.stateToModelLocation(location, state));
+            }
+        });
+        return new EmulatedRegistry<>(ModelResourceLocation.class, this.loadedBakedModels, topLevelModelLocations, this.mrlModelOverrides);
+    }
+
+    public Map<ResourceLocation, BakedModel> getStandaloneEmulatedRegistry() {
+        return new EmulatedRegistry<>(ResourceLocation.class, this.loadedStandaloneModels, Set.of(), this.standaloneModelOverrides);
+    }
+
+    public Map<ResourceLocation, ItemModel> getItemModelEmulatedRegistry() {
+        return new EmulatedRegistry<>(ResourceLocation.class, this.loadedItemModels, BuiltInRegistries.ITEM.keySet(), this.itemStackModelOverrides);
+    }
+
+    public Map<ResourceLocation, ClientItem.Properties> getItemPropertiesEmulatedRegistry() {
+        return Maps.transformValues(new EmulatedRegistry<>(ResourceLocation.class, this.loadedClientItemProperties, BuiltInRegistries.ITEM.keySet(), Map.of()), ClientItem::properties);
+    }
+
+    private <K, V> LoadingCache<K, Optional<V>> makeLoadingCache(Function<K, Optional<V>> loadingFunction) {
+        return CacheBuilder.newBuilder()
+                .expireAfterAccess(3, TimeUnit.MINUTES)
+                .maximumSize(1000)
+                .concurrencyLevel(8)
+                .softValues()
+                .build(new CacheLoader<>() {
+                    @Override
+                    public Optional<V> load(K key) {
+                        return loadingFunction.apply(key);
+                    }
+                });
+    }
+
+    private static class EmulatedRegistry<K, V> implements Map<K, V> {
+        private final LoadingCache<K, Optional<V>> realCache;
+        private final Set<K> keys;
+        private final Map<K, V> overrides;
+        private final Class<K> keyClass;
+
+        public EmulatedRegistry(Class<K> keyClass, LoadingCache<K, Optional<V>> realCache, Set<K> keys, Map<K, V> overrides) {
+            this.keyClass = keyClass;
+            this.realCache = realCache;
+            this.keys = Collections.unmodifiableSet(keys);
+            this.overrides = overrides;
+        }
+
+        @Override
+        public V get(Object key) {
+            if (this.keyClass.isAssignableFrom(key.getClass())) {
+                return this.realCache.getUnchecked((K)key).orElse(null);
+            } else {
+                return null;
+            }
+        }
+
+        @Override
+        public V put(K key, V value) {
+            V oldValue = this.realCache.getUnchecked(key).orElse(null);
+            this.overrides.put(key, value);
+            this.realCache.invalidate(key);
+            return oldValue;
+        }
+
+        @Override
+        public V remove(Object key) {
+            this.overrides.remove(key);
+            this.realCache.invalidate(key);
+            return null;
+        }
+
+        @Override
+        public void putAll(@NotNull Map<? extends K, ? extends V> m) {
+            m.forEach(this::put);
+        }
+
+        @Override
+        public void clear() {
+            this.overrides.clear();
+            this.realCache.invalidateAll();
+        }
+
+        @Override
+        public @NotNull Set<K> keySet() {
+            return keys;
+        }
+
+        @Override
+        public @NotNull Collection<V> values() {
+            return List.of();
+        }
+
+        @Override
+        public int size() {
+            return keys.size();
+        }
+
+        @Override
+        public boolean isEmpty() {
+            return false;
+        }
+
+        @Override
+        public boolean containsKey(Object key) {
+            return keys.contains(key);
+        }
+
+        @Override
+        public boolean containsValue(Object value) {
+            return false;
+        }
+
+        @Override
+        public @NotNull Set<Entry<K, V>> entrySet() {
+            return new AbstractSet<>() {
+                @Override
+                public Iterator<Entry<K, V>> iterator() {
+                    return Iterators.transform(keys.iterator(), key -> new Entry<>() {
+                        @Override
+                        public K getKey() {
+                            return key;
+                        }
+
+                        @Override
+                        public V getValue() {
+                            return get(key);
+                        }
+
+                        @Override
+                        public V setValue(V value) {
+                            return put(key, value);
+                        }
+                    });
+                }
+
+                @Override
+                public int size() {
+                    return keys.size();
+                }
+            };
+        }
+
+        @Override
+        public void replaceAll(BiFunction<? super K, ? super V, ? extends V> function) {
+            for(K location : keys) {
+                /*
+                 * Fetching every model is insanely slow. So we call the function with a null object first, since it
+                 * probably isn't expecting that. If we get an exception thrown, or it returns nonnull, then we know
+                 * it actually cares about the given model.
+                 */
+                boolean needsReplacement;
+                try {
+                    needsReplacement = function.apply(location, null) != null;
+                } catch(Throwable e) {
+                    needsReplacement = true;
+                }
+                if(needsReplacement) {
+                    V existing = get(location);
+                    V replacement = function.apply(location, existing);
+                    if(replacement != existing) {
+                        put(location, replacement);
+                    }
+                }
+            }
+        }
     }
 
     private Optional<BlockStateModelLoader.LoadedModels> loadBlockStateDefinition(ResourceLocation location) {
@@ -178,29 +327,31 @@ public class DynamicModelProvider {
         return Optional.of(BlockStateModelLoader.loadBlockStateDefinitionStack(location, stateDefinition, loadedDefinitions, this.unbakedMissingModel));
     }
 
-    private BakedModel bakeModel(UnbakedModel model, ModelResourceLocation location) {
+    private BakedModel bakeModel(UnbakedModel model, ModelDebugName name) {
         synchronized (this) {
             this.resolver.clearResolver();
             model.resolveDependencies(this.resolver);
-            var modelBaker = new DynamicBaker(location::toString);
+            var modelBaker = new DynamicBaker(name);
             return UnbakedModel.bakeWithTopModelValues(model, modelBaker, BlockModelRotation.X0_Y0);
         }
     }
 
-    private BakedModel bakeModel(UnbakedBlockStateModel model, ModelResourceLocation location) {
+    private BakedModel bakeModel(UnbakedBlockStateModel model, ModelDebugName name) {
         synchronized (this) {
             this.resolver.clearResolver();
             model.resolveDependencies(this.resolver);
-            var modelBaker = new DynamicBaker(location::toString);
+            var modelBaker = new DynamicBaker(name);
             return model.bake(modelBaker);
         }
     }
 
     private Optional<BakedModel> loadBakedModel(ModelResourceLocation location) {
+        var override = this.mrlModelOverrides.get(location);
+        if (override != null) {
+            return Optional.of(override);
+        }
         if (location.variant().equals("standalone") || location.variant().equals("fabric_resource")) {
-            return this.loadedBlockModels.getUnchecked(location.id()).map(unbakedModel -> {
+            return this.loadStandaloneModel(location.id());
-                return this.bakeModel(unbakedModel, location);
-            });
         } else {
             var optLoadedModels = this.loadedStateDefinitions.getUnchecked(location.id());
             Optional<UnbakedBlockStateModel> unbakedModelOpt = optLoadedModels.map(loadedModels -> {
@@ -212,11 +363,21 @@ public class DynamicModelProvider {
                 }
             });
             return unbakedModelOpt.map(unbakedModel -> {
-                return this.bakeModel(unbakedModel, location);
+                return this.bakeModel(unbakedModel, location::toString);
             });
         }
     }
 
+    private Optional<BakedModel> loadStandaloneModel(ResourceLocation location) {
+        var override = this.standaloneModelOverrides.get(location);
+        if (override != null) {
+            return Optional.of(override);
+        }
+        return this.loadedBlockModels.getUnchecked(location).map(unbakedModel -> {
+            return this.bakeModel(unbakedModel, location::toString);
+        });
+    }
+
     private Optional<UnbakedModel> loadBlockModel(ResourceLocation location) {
         if (location.equals(ItemModelGenerator.GENERATED_ITEM_MODEL_ID)) {
             return Optional.of(this.itemModelGenerator);
@@ -243,7 +404,7 @@ public class DynamicModelProvider {
                 ClientItem clientItem = ClientItem.CODEC.parse(JsonOps.INSTANCE, JsonParser.parseReader(reader)).getOrThrow();
                 return Optional.of(clientItem);
             } catch(Exception e) {
-                ModernFix.LOGGER.error("Failed to load block model {} from '{}'", location, resource.get().sourcePackId(), e);
+                ModernFix.LOGGER.error("Failed to load client item {} from '{}'", location, resource.get().sourcePackId(), e);
                 return Optional.empty();
             }
         } else {
@@ -253,6 +414,10 @@ public class DynamicModelProvider {
     }
 
     private Optional<ItemModel> loadItemModel(ResourceLocation location) {
+        var override = this.itemStackModelOverrides.get(location);
+        if (override != null) {
+            return Optional.of(override);
+        }
         return this.loadedClientItemProperties.getUnchecked(location).map(clientItem -> {
             var bakingContext = new ItemModel.BakingContext(new DynamicBaker(location::toString), this.entityModelSet, this.missingItemModel);
             return clientItem.model().bake(bakingContext);
@@ -271,6 +436,10 @@ public class DynamicModelProvider {
         return this.loadedItemModels.getUnchecked(location).orElse(this.missingItemModel);
     }
 
+    public BakedModel getStandaloneModel(ResourceLocation location) {
+        return this.loadedStandaloneModels.getUnchecked(location).orElse(this.missingModel);
+    }
+
     private class DynamicBaker implements ModelBaker {
         private final ModelDebugName modelDebugName;
 
@@ -331,4 +500,10 @@ public class DynamicModelProvider {
             this.resolvedModels.clear();
         }
     }
+
+    public static WeakReference<DynamicModelProvider> currentReloadingModelProvider = new WeakReference<>(null);
+
+    public interface ModelManagerExtension {
+        DynamicModelProvider mfix$getModelProvider();
+    }
 }

neoforge/src/main/java/org/embeddedt/modernfix/neoforge/dynresources/ModelBakeEventHelper.java

@@ -1,55 +1,24 @@
 package org.embeddedt.modernfix.neoforge.dynresources;
 
-import com.google.common.collect.ForwardingMap;
-import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.Iterators;
-import com.google.common.collect.Sets;
 import com.google.common.graph.GraphBuilder;
 import com.google.common.graph.MutableGraph;
-import net.minecraft.client.renderer.block.BlockModelShaper;
+import net.minecraft.client.resources.model.ModelBakery;
-import net.minecraft.client.resources.model.BakedModel;
-import net.minecraft.client.resources.model.MissingBlockModel;
-import net.minecraft.client.resources.model.ModelResourceLocation;
-import net.minecraft.core.registries.BuiltInRegistries;
-import net.minecraft.resources.ResourceLocation;
-import net.minecraft.world.level.block.state.BlockState;
 import net.neoforged.fml.ModContainer;
 import net.neoforged.fml.ModList;
 import net.neoforged.neoforgespi.language.IModInfo;
-import org.embeddedt.modernfix.ModernFix;
+import org.embeddedt.modernfix.dynamicresources.DynamicModelProvider;
-import org.embeddedt.modernfix.util.ForwardingInclDefaultsMap;
-import org.jetbrains.annotations.Nullable;
 
 import java.util.*;
-import java.util.function.BiFunction;
 
 /**
  * Stores a list of all known default block/item models in the game, and provides a namespaced version
  * of the model registry that emulates vanilla keySet behavior.
  */
 public class ModelBakeEventHelper {
-    // TODO: make into config option
+    private final DynamicModelProvider modelRegistry;
-    private static final Set<String> INCOMPATIBLE_MODS = ImmutableSet.of(
-            "industrialforegoing",
-            "mekanism",
-            "vampirism",
-            "elevatorid",
-            "cfm",
-            "embers");
-    private final Map<ModelResourceLocation, BakedModel> modelRegistry;
-    private final Set<ModelResourceLocation> topLevelModelLocations;
     private final MutableGraph<String> dependencyGraph;
-    public ModelBakeEventHelper(Map<ModelResourceLocation, BakedModel> modelRegistry) {
+    public ModelBakeEventHelper(DynamicModelProvider modelRegistry) {
         this.modelRegistry = modelRegistry;
-        this.topLevelModelLocations = new HashSet<>(modelRegistry.keySet());
-        // Skip going through ModelLocationCache because most of the accesses will be misses
-        BuiltInRegistries.BLOCK.entrySet().forEach(entry -> {
-            var location = entry.getKey().location();
-            for(BlockState state : entry.getValue().getStateDefinition().getPossibleStates()) {
-                topLevelModelLocations.add(BlockModelShaper.stateToModelLocation(location, state));
-            }
-        });
-        BuiltInRegistries.ITEM.keySet().forEach(location -> topLevelModelLocations.add(new ModelResourceLocation(location, "inventory")));
         this.dependencyGraph = GraphBuilder.undirected().build();
         ModList.get().forEachModContainer((id, mc) -> {
             this.dependencyGraph.addNode(id);
@@ -69,177 +38,14 @@ public class ModelBakeEventHelper {
         }
     }
 
-    private static final Set<String> WARNED_MOD_IDS = new HashSet<>();
+    public ModelBakery.BakingResult createDynamicResult() {
-
+        return new ModelBakery.BakingResult(
-    /**
+                this.modelRegistry.getMissingBakedModel(),
-     * Create a model registry that warns if keySet, entrySet, values are accessed.
+                this.modelRegistry.getTopLevelEmulatedRegistry(),
-     * @param modId the mod that the event is being fired for
+                this.modelRegistry.getMissingItemModel(),
-     * @return a wrapper around the model registry
+                this.modelRegistry.getItemModelEmulatedRegistry(),
-     */
+                this.modelRegistry.getItemPropertiesEmulatedRegistry(),
-    private Map<ModelResourceLocation, BakedModel> createWarningRegistry(String modId) {
+                this.modelRegistry.getStandaloneEmulatedRegistry()
-        return new ForwardingInclDefaultsMap<ModelResourceLocation, BakedModel>() {
+        );
-            @Override
-            protected Map<ModelResourceLocation, BakedModel> delegate() {
-                return modelRegistry;
-            }
-
-            private void logWarning() {
-                if(!WARNED_MOD_IDS.add(modId))
-                    return;
-                ModernFix.LOGGER.warn("Mod '{}' is accessing Map#keySet/entrySet/values/replaceAll on the model registry map inside its event handler." +
-                        " This probably won't work as expected with dynamic resources on. Prefer using Map#get/put and constructing ModelResourceLocations another way.", modId);
-            }
-
-            @Override
-            public Set<ModelResourceLocation> keySet() {
-                logWarning();
-                return super.keySet();
-            }
-
-            @Override
-            public Set<Entry<ModelResourceLocation, BakedModel>> entrySet() {
-                logWarning();
-                return super.entrySet();
-            }
-
-            @Override
-            public Collection<BakedModel> values() {
-                logWarning();
-                return super.values();
-            }
-
-            @Override
-            public void replaceAll(BiFunction<? super ModelResourceLocation, ? super BakedModel, ? extends BakedModel> function) {
-                logWarning();
-                super.replaceAll(function);
-            }
-        };
-    }
-
-    public Map<ModelResourceLocation, BakedModel> wrapRegistry(String modId) {
-        final Set<String> modIdsToInclude = new HashSet<>();
-        modIdsToInclude.add(modId);
-        try {
-            modIdsToInclude.addAll(this.dependencyGraph.adjacentNodes(modId));
-        } catch(IllegalArgumentException ignored) { /* sanity check */ }
-        modIdsToInclude.remove("minecraft");
-        if(modIdsToInclude.stream().noneMatch(INCOMPATIBLE_MODS::contains))
-            return createWarningRegistry(modId);
-        Set<ModelResourceLocation> ourModelLocations = Sets.filter(this.topLevelModelLocations, loc -> modIdsToInclude.contains(loc.id().getNamespace()));
-        BakedModel missingModel = modelRegistry.get(MissingBlockModel.VARIANT);
-        return new ForwardingMap<ModelResourceLocation, BakedModel>() {
-            @Override
-            protected Map<ModelResourceLocation, BakedModel> delegate() {
-                return modelRegistry;
-            }
-
-            @Override
-            public BakedModel get(@Nullable Object key) {
-                BakedModel model = super.get(key);
-                if(model == null && key != null && modIdsToInclude.contains(((ResourceLocation)key).getNamespace())) {
-                    ModernFix.LOGGER.warn("Model {} is missing, but was requested in model bake event. Returning missing model", key);
-                    return missingModel;
-                }
-                return model;
-            }
-
-            @Override
-            public Set<ModelResourceLocation> keySet() {
-                return ourModelLocations;
-            }
-
-            @Override
-            public boolean containsKey(@Nullable Object key) {
-                return ourModelLocations.contains(key) || super.containsKey(key);
-            }
-
-            @Override
-            public Set<Entry<ModelResourceLocation, BakedModel>> entrySet() {
-                return new DynamicModelEntrySet(this, ourModelLocations);
-            }
-
-            @Override
-            public void replaceAll(BiFunction<? super ModelResourceLocation, ? super BakedModel, ? extends BakedModel> function) {
-                ModernFix.LOGGER.warn("Mod '{}' is calling replaceAll on the model registry. Some hacks will be used to keep this fast, but they may not be 100% compatible.", modId);
-                List<ModelResourceLocation> locations = new ArrayList<>(keySet());
-                for(ModelResourceLocation location : locations) {
-                    /*
-                     * Fetching every model is insanely slow. So we call the function with a null object first, since it
-                     * probably isn't expecting that. If we get an exception thrown, or it returns nonnull, then we know
-                     * it actually cares about the given model.
-                     */
-                    boolean needsReplacement;
-                    try {
-                        needsReplacement = function.apply(location, null) != null;
-                    } catch(Throwable e) {
-                        needsReplacement = true;
-                    }
-                    if(needsReplacement) {
-                        BakedModel existing = get(location);
-                        BakedModel replacement = function.apply(location, existing);
-                        if(replacement != existing) {
-                            put(location, replacement);
-                        }
-                    }
-                }
-            }
-        };
-    }
-
-    private static class DynamicModelEntrySet extends AbstractSet<Map.Entry<ModelResourceLocation, BakedModel>> {
-        private final Map<ModelResourceLocation, BakedModel> modelRegistry;
-        private final Set<ModelResourceLocation> modelLocations;
-
-        private DynamicModelEntrySet(Map<ModelResourceLocation, BakedModel> modelRegistry, Set<ModelResourceLocation> modelLocations) {
-            this.modelRegistry = modelRegistry;
-            this.modelLocations = modelLocations;
-        }
-
-        @Override
-        public Iterator<Map.Entry<ModelResourceLocation, BakedModel>> iterator() {
-            return Iterators.transform(Iterators.unmodifiableIterator(this.modelLocations.iterator()), DynamicModelEntry::new);
-        }
-
-        @Override
-        public boolean contains(Object o) {
-            if(o instanceof Map.Entry entry) {
-                return modelRegistry.containsKey(entry.getKey());
-            } else {
-                return false;
-            }
-        }
-
-        @Override
-        public int size() {
-            return modelRegistry.size();
-        }
-
-        @Override
-        public boolean removeAll(Collection<?> c) {
-            throw new UnsupportedOperationException();
-        }
-
-        private class DynamicModelEntry implements Map.Entry<ModelResourceLocation, BakedModel> {
-            private final ModelResourceLocation location;
-
-            private DynamicModelEntry(ModelResourceLocation location) {
-                this.location = location;
-            }
-
-            @Override
-            public ModelResourceLocation getKey() {
-                return this.location;
-            }
-
-            @Override
-            public BakedModel getValue() {
-                return modelRegistry.get(this.location);
-            }
-
-            @Override
-            public BakedModel setValue(BakedModel value) {
-                return modelRegistry.put(this.location, value);
-            }
-        }
     }
 }

neoforge/src/main/java/org/embeddedt/modernfix/neoforge/mixin/perf/dynamic_resources/ForgeHooksClientMixin.java

@@ -1,42 +0,0 @@
-package org.embeddedt.modernfix.neoforge.mixin.perf.dynamic_resources;
-
-import net.neoforged.bus.api.Event;
-import net.neoforged.fml.ModLoader;
-import net.neoforged.neoforge.client.ClientHooks;
-import net.neoforged.neoforge.client.event.ModelEvent;
-import org.embeddedt.modernfix.ModernFix;
-import org.spongepowered.asm.mixin.Mixin;
-import org.spongepowered.asm.mixin.injection.At;
-import org.spongepowered.asm.mixin.injection.Redirect;
-
-@Mixin(ClientHooks.class)
-public class ForgeHooksClientMixin {
-    /**
-     * Generate a more realistic keySet that contains every item and block model location, to help with mod compat.
-     */
-    @Redirect(method = "onModifyBakingResult", at = @At(value = "INVOKE", target = "Lnet/neoforged/fml/ModLoader;postEvent(Lnet/neoforged/bus/api/Event;)V"), remap = false)
-    private static void postNamespacedKeySetEvent(Event event) {
-        if(ModLoader.hasErrors())
-            return;
-        ModelEvent.ModifyBakingResult bakeEvent = ((ModelEvent.ModifyBakingResult)event);
-        /*
-        ModelBakeEventHelper helper = new ModelBakeEventHelper(bakeEvent.getModels());
-        Method acceptEv = ObfuscationReflectionHelper.findMethod(ModContainer.class, "acceptEvent", Event.class);
-        ModList.get().forEachModContainer((id, mc) -> {
-            Map<ModelResourceLocation, BakedModel> newRegistry = helper.wrapRegistry(id);
-            ModelEvent.ModifyBakingResult postedEvent = new ModelEvent.ModifyBakingResult(newRegistry, bakeEvent.getTextureGetter(), bakeEvent.getModelBakery());
-            Stopwatch timer = Stopwatch.createStarted();
-            try {
-                acceptEv.invoke(mc, postedEvent);
-            } catch(ReflectiveOperationException e) {
-                e.printStackTrace();
-            }
-            timer.stop();
-            if(timer.elapsed(TimeUnit.SECONDS) >= 1) {
-                ModernFix.LOGGER.warn("Mod '{}' took {} in the model bake event", id, timer);
-            }
-        });
-         */
-        ModernFix.LOGGER.warn("ModifyBakingResult support not reimplemented yet");
-    }
-}

neoforge/src/main/java/org/embeddedt/modernfix/neoforge/mixin/perf/dynamic_resources/ModelManagerMixinNeo.java

@@ -0,0 +1,26 @@
+package org.embeddedt.modernfix.neoforge.mixin.perf.dynamic_resources;
+
+import net.minecraft.client.resources.model.ModelBakery;
+import net.minecraft.client.resources.model.ModelManager;
+import net.neoforged.fml.ModLoader;
+import org.embeddedt.modernfix.ModernFix;
+import org.embeddedt.modernfix.annotation.ClientOnlyMixin;
+import org.embeddedt.modernfix.dynamicresources.DynamicModelProvider;
+import org.embeddedt.modernfix.neoforge.dynresources.ModelBakeEventHelper;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.ModifyArg;
+
+@Mixin(ModelManager.class)
+@ClientOnlyMixin
+public class ModelManagerMixinNeo {
+    @ModifyArg(method = "loadModels", at = @At(value = "INVOKE", target = "Lnet/neoforged/neoforge/client/ClientHooks;onModifyBakingResult(Lnet/minecraft/client/resources/model/ModelBakery$BakingResult;Ljava/util/Map;Lnet/minecraft/client/resources/model/ModelBakery;)V"), remap = false, index = 0)
+    private static ModelBakery.BakingResult useDynamicBakingResult(ModelBakery.BakingResult bakingResult) {
+        var currentReloadingProvider = DynamicModelProvider.currentReloadingModelProvider.get();
+        if(ModLoader.hasErrors() || currentReloadingProvider == null) {
+            ModernFix.LOGGER.error("Errors encountered - not using dynamic model BakingResult");
+            return bakingResult;
+        }
+        return new ModelBakeEventHelper(currentReloadingProvider).createDynamicResult();
+    }
+}