Reduce memory usage of dynamic CTMPackReloadListener

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

@@ -1,8 +1,11 @@
 package org.embeddedt.modernfix.duck;
 
-
+import com.google.common.collect.ImmutableList;
-import net.minecraft.client.renderer.texture.AtlasSet;
+import net.minecraft.client.resources.model.ModelResourceLocation;
+import net.minecraft.world.level.block.Block;
+import net.minecraft.world.level.block.state.BlockState;
+import net.minecraft.world.level.block.state.StateDefinition;
 
 public interface IExtendedModelBakery {
-    AtlasSet getUnfinishedAtlasSet();
+    ImmutableList<BlockState> getBlockStatesForMRL(StateDefinition<Block, BlockState> stateDefinition, ModelResourceLocation location);
 }

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

@@ -39,6 +39,7 @@ import net.minecraftforge.registries.ForgeRegistryEntry;
 import org.apache.commons.lang3.tuple.Triple;
 import org.apache.logging.log4j.Logger;
 import org.embeddedt.modernfix.ModernFix;
+import org.embeddedt.modernfix.duck.IExtendedModelBakery;
 import org.embeddedt.modernfix.dynamicresources.DynamicBakedModelProvider;
 import org.embeddedt.modernfix.dynamicresources.DynamicModelBakeEvent;
 import org.embeddedt.modernfix.dynamicresources.ModelLocationCache;
@@ -64,7 +65,7 @@ import java.util.stream.Stream;
 
 /* high priority so that our injectors are added before other mods' */
 @Mixin(value = ModelBakery.class, priority = 600)
-public abstract class ModelBakeryMixin {
+public abstract class ModelBakeryMixin implements IExtendedModelBakery {
 
     private static final boolean debugDynamicModelLoading = Boolean.getBoolean("modernfix.debugDynamicModelLoading");
 
@@ -438,6 +439,11 @@ public abstract class ModelBakeryMixin {
         return ImmutableList.copyOf(finalList);
     }
 
+    @Override
+    public ImmutableList<BlockState> getBlockStatesForMRL(StateDefinition<Block, BlockState> stateDefinition, ModelResourceLocation location) {
+        return loadOnlyRelevantBlockState(stateDefinition, location);
+    }
+
     @Inject(method = "getBakedModel", at = @At("HEAD"), cancellable = true)
     public void getOrLoadBakedModelDynamic(ResourceLocation arg, ModelState arg2, Function<Material, TextureAtlasSprite> textureGetter, CallbackInfoReturnable<BakedModel> cir) {
         Triple<ResourceLocation, Transformation, Boolean> triple = Triple.of(arg, arg2.getRotation(), arg2.isUvLocked());

src/main/java/org/embeddedt/modernfix/mixin/perf/dynamic_resources/ctm/CTMPackReloadListenerMixin.java

@@ -1,5 +1,6 @@
 package org.embeddedt.modernfix.mixin.perf.dynamic_resources.ctm;
 
+import com.google.common.collect.ImmutableList;
 import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap;
 import net.minecraft.client.renderer.ItemBlockRenderTypes;
 import net.minecraft.client.renderer.RenderType;
@@ -8,13 +9,17 @@ import net.minecraft.client.resources.model.BakedModel;
 import net.minecraft.client.resources.model.ModelResourceLocation;
 import net.minecraft.client.resources.model.MultiPartBakedModel;
 import net.minecraft.client.resources.model.WeightedBakedModel;
+import net.minecraft.resources.ResourceLocation;
 import net.minecraft.world.level.block.Block;
+import net.minecraft.world.level.block.Blocks;
 import net.minecraft.world.level.block.state.BlockState;
 import net.minecraftforge.common.MinecraftForge;
 import net.minecraftforge.eventbus.api.EventPriority;
 import net.minecraftforge.fml.javafmlmod.FMLJavaModLoadingContext;
 import net.minecraftforge.registries.ForgeRegistries;
 import net.minecraftforge.registries.IRegistryDelegate;
+import org.embeddedt.modernfix.ModernFix;
+import org.embeddedt.modernfix.duck.IExtendedModelBakery;
 import org.embeddedt.modernfix.dynamicresources.DynamicModelBakeEvent;
 import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
@@ -37,8 +42,6 @@ public abstract class CTMPackReloadListenerMixin {
 
     @Shadow protected abstract Predicate<RenderType> getExistingRenderCheck(Block block);
 
-    private Map<ModelResourceLocation, BlockState> locationToState = new Object2ObjectOpenHashMap<>();
-
     @Inject(method = "<init>", at = @At("RETURN"))
     private void onInit(CallbackInfo ci) {
         MinecraftForge.EVENT_BUS.addListener(EventPriority.LOW, this::onModelBake);
@@ -48,28 +51,30 @@ public abstract class CTMPackReloadListenerMixin {
     private void refreshLayerHacks() {
         blockRenderChecks.forEach((b, p) -> ItemBlockRenderTypes.setRenderLayer((Block) b.get(), p));
         blockRenderChecks.clear();
-        if(locationToState.isEmpty()) {
-            for(Block block : ForgeRegistries.BLOCKS.getValues()) {
-                for(BlockState state : block.getStateDefinition().getPossibleStates()) {
-                    locationToState.put(BlockModelShaper.stateToModelLocation(state), state);
-                }
-            }
-        }
     }
 
     private void onModelBake(DynamicModelBakeEvent event) {
         if(!(event.getModel() instanceof AbstractCTMBakedModel || event.getModel() instanceof WeightedBakedModel || event.getModel() instanceof MultiPartBakedModel))
             return;
-        BlockState state = locationToState.get(event.getLocation());
+        /* we construct a new ResourceLocation because an MRL is coming in */
-        if(state == null)
+        Block block = ForgeRegistries.BLOCKS.getValue(new ResourceLocation(event.getLocation().getNamespace(), event.getLocation().getPath()));
+        if(block == null || block == Blocks.AIR || blockRenderChecks.containsKey(block.delegate))
             return;
-        Block block = state.getBlock();
+        /* find all states that match this MRL */
-        if(blockRenderChecks.containsKey(block.delegate))
+        ImmutableList<BlockState> allStates;
+        try {
+            allStates = ((IExtendedModelBakery)(Object)event.getModelLoader()).getBlockStatesForMRL(block.getStateDefinition(), (ModelResourceLocation)event.getLocation());
+        } catch(RuntimeException e) {
+            ModernFix.LOGGER.error("Couldn't get state for MRL " + event.getLocation(), e);
             return;
-        Predicate<RenderType> newPredicate = this.getLayerCheck(state, event.getModel());
+        }
-        if(newPredicate != null) {
+        for(BlockState state : allStates) {
-            blockRenderChecks.put(block.delegate, this.getExistingRenderCheck(block));
+            Predicate<RenderType> newPredicate = this.getLayerCheck(state, event.getModel());
-            ItemBlockRenderTypes.setRenderLayer(block, newPredicate);
+            if(newPredicate != null) {
+                blockRenderChecks.put(block.delegate, this.getExistingRenderCheck(block));
+                ItemBlockRenderTypes.setRenderLayer(block, newPredicate);
+                return;
+            }
         }
     }
 }