Keep custom models loaded permanently on Fabric

common/src/main/java/org/embeddedt/modernfix/util/LayeredForwardingMap.java

@@ -0,0 +1,124 @@
+package org.embeddedt.modernfix.util;
+
+import it.unimi.dsi.fastutil.objects.ObjectOpenHashSet;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.*;
+
+/**
+ * Simple forwarding map implementation that allows layering multiple maps together, with the last layer being
+ * mutable.
+ */
+public class LayeredForwardingMap<K, V> implements Map<K, V> {
+    private final Map<K, V>[] layers;
+
+    public LayeredForwardingMap(Map<K, V>[] layers) {
+        if(layers.length < 1)
+            throw new IllegalArgumentException();
+        for(Map<K, V> layer : layers) {
+            if(layer == null)
+                throw new IllegalArgumentException();
+        }
+        this.layers = layers;
+    }
+
+    @Override
+    public int size() {
+        return 1;
+    }
+
+    @Override
+    public boolean isEmpty() {
+        for(Map<K, V> map : layers) {
+            if(!map.isEmpty())
+                return false;
+        }
+        return true;
+    }
+
+    @Override
+    public boolean containsKey(Object key) {
+        for(Map<K, V> map : layers) {
+            if(map.containsKey(key))
+                return true;
+        }
+        return false;
+    }
+
+    @Override
+    public boolean containsValue(Object value) {
+        for(Map<K, V> map : layers) {
+            if(map.containsValue(value))
+                return true;
+        }
+        return false;
+    }
+
+    @Override
+    public V get(Object key) {
+        for(Map<K, V> map : layers) {
+            V value = map.get(key);
+            if(value != null)
+                return value;
+        }
+        return null;
+    }
+
+    @Nullable
+    @Override
+    public V put(K key, V value) {
+        if(value == null)
+            throw new IllegalArgumentException();
+        return layers[layers.length - 1].put(key, value);
+    }
+
+    @Override
+    public V remove(Object key) {
+        for(Map<K, V> map : layers) {
+            map.remove(key);
+        }
+        return null;
+    }
+
+    @Override
+    public void putAll(@NotNull Map<? extends K, ? extends V> m) {
+        for(V value : m.values()) {
+            if(value == null)
+                throw new IllegalArgumentException();
+        }
+        layers[layers.length - 1].putAll(m);
+    }
+
+    @Override
+    public void clear() {
+        throw new UnsupportedOperationException();
+    }
+
+    @NotNull
+    @Override
+    public Set<K> keySet() {
+        Set<K> keys = new ObjectOpenHashSet<>();
+        for(Map<K, V> map : layers) {
+            keys.addAll(map.keySet());
+        }
+        return Collections.unmodifiableSet(keys);
+    }
+
+    @NotNull
+    @Override
+    public Collection<V> values() {
+        Set<K> keys = keySet();
+        List<V> vals = new ArrayList<>();
+        for(K key : keys) {
+            vals.add(get(key));
+        }
+        return vals;
+    }
+
+    @NotNull
+    @Override
+    public Set<Entry<K, V>> entrySet() {
+        throw new UnsupportedOperationException();
+    }
+}

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

@@ -8,6 +8,8 @@ import com.google.common.collect.ImmutableList;
 import com.mojang.math.Transformation;
 import net.minecraft.client.renderer.block.model.BlockModel;
 import net.minecraft.client.renderer.block.model.ItemModelGenerator;
+import net.minecraft.client.renderer.block.model.MultiVariant;
+import net.minecraft.client.renderer.block.model.multipart.MultiPart;
 import net.minecraft.client.renderer.texture.AtlasSet;
 import net.minecraft.client.renderer.texture.TextureAtlasSprite;
 import net.minecraft.client.renderer.texture.TextureManager;
@@ -25,6 +27,7 @@ import org.embeddedt.modernfix.annotation.ClientOnlyMixin;
 import org.embeddedt.modernfix.duck.IExtendedModelBakery;
 import org.embeddedt.modernfix.dynamicresources.DynamicBakedModelProvider;
 import org.embeddedt.modernfix.dynamicresources.ModelBakeryHelpers;
+import org.embeddedt.modernfix.util.LayeredForwardingMap;
 import org.jetbrains.annotations.Nullable;
 import org.spongepowered.asm.mixin.*;
 import org.spongepowered.asm.mixin.injection.At;
@@ -145,7 +148,7 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
         this.loadedModels.put(MISSING_MODEL_LOCATION, this.missingModel);
         this.bakedCache = loadedBakedModels.asMap();
         ConcurrentMap<ResourceLocation, UnbakedModel> unbakedCacheBackingMap = loadedModels.asMap();
-        this.unbakedCache = new ForwardingMap<ResourceLocation, UnbakedModel>() {
+        Map<ResourceLocation, UnbakedModel> mutableBackingMap = new ForwardingMap<ResourceLocation, UnbakedModel>() {
             @Override
             protected Map<ResourceLocation, UnbakedModel> delegate() {
                 return unbakedCacheBackingMap;
@@ -157,6 +160,11 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
                 return super.put(key, value);
             }
         };
+        // discard unwrapped models
+        int oldSize = this.unbakedCache.size();
+        this.unbakedCache.entrySet().removeIf(entry -> entry.getValue() instanceof BlockModel || entry.getValue() instanceof MultiVariant || entry.getValue() instanceof MultiPart);
+        ModernFix.LOGGER.info("{} models evicted, {} custom models loaded permanently", oldSize - this.unbakedCache.size(), this.unbakedCache.size());
+        this.unbakedCache = new LayeredForwardingMap<>(new Map[] { this.unbakedCache, mutableBackingMap });
         this.bakedTopLevelModels = new DynamicBakedModelProvider((ModelBakery)(Object)this, bakedCache);
 
         // ensure missing model is a permanent override