Tweak Fabric model bakery to run more vanilla code

common/src/main/resources/modernfix.accesswidener

@@ -36,6 +36,8 @@ accessible class net/minecraft/server/level/ChunkMap$DistanceManager
 accessible class net/minecraft/client/resources/model/ModelBakery$BakedCacheKey
 accessible field net/minecraft/client/resources/model/ModelBakery bakedCache Ljava/util/Map;
 accessible field net/minecraft/client/resources/model/ModelBakery ITEM_MODEL_GENERATOR Lnet/minecraft/client/renderer/block/model/ItemModelGenerator;
+accessible method net/minecraft/client/resources/model/ModelBakery loadTopLevel (Lnet/minecraft/client/resources/model/ModelResourceLocation;)V
+accessible field net/minecraft/client/resources/model/ModelBakery topLevelModels Ljava/util/Map;
 accessible class net/minecraft/client/resources/model/ModelBakery$ModelBakerImpl
 accessible method net/minecraft/client/resources/model/ModelBakery$ModelBakerImpl <init> (Lnet/minecraft/client/resources/model/ModelBakery;Ljava/util/function/BiFunction;Lnet/minecraft/resources/ResourceLocation;)V
 accessible method net/minecraft/client/resources/model/ModelBakery$BakedCacheKey <init> (Lnet/minecraft/resources/ResourceLocation;Lcom/mojang/math/Transformation;Z)V

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

@@ -1,9 +1,13 @@
 package org.embeddedt.modernfix.fabric.mixin.perf.dynamic_resources;
 
+import net.fabricmc.loader.api.FabricLoader;
 import net.minecraft.client.renderer.block.model.BlockModel;
 import net.minecraft.client.renderer.texture.TextureAtlasSprite;
 import net.minecraft.client.resources.model.*;
+import net.minecraft.core.registries.BuiltInRegistries;
 import net.minecraft.resources.ResourceLocation;
+import net.minecraft.world.level.block.Block;
+import net.minecraft.world.level.block.state.BlockState;
 import org.embeddedt.modernfix.ModernFix;
 import org.embeddedt.modernfix.duck.IExtendedModelBakery;
 import org.embeddedt.modernfix.dynamicresources.DynamicBakedModelProvider;
@@ -14,6 +18,9 @@ import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
 
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.util.Optional;
 import java.util.function.Function;
 
 @Mixin(ModelBakery.ModelBakerImpl.class)
@@ -25,6 +32,24 @@ public abstract class ModelBakerImplMixin {
 
     @Shadow @Final private Function<Material, TextureAtlasSprite> modelTextureGetter;
 
+    private static final MethodHandle blockStateLoaderHandle;
+    static {
+        try {
+            blockStateLoaderHandle = MethodHandles.lookup().unreflect(ModelBakery.class.getDeclaredMethod(
+                    FabricLoader.getInstance().getMappingResolver().mapMethodName(
+                            "intermediary",
+                            "net.minecraft.client.resources.model.ModelBakery",
+                            "method_4716",
+                            "(Lnet/minecraft/world/level/block/state/BlockState;)V"
+                    ),
+                    BlockState.class
+            ));
+        } catch(ReflectiveOperationException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+
     @Inject(method = "bake", at = @At("HEAD"), cancellable = true)
     public void getOrLoadBakedModelDynamic(ResourceLocation arg, ModelState arg2, CallbackInfoReturnable<BakedModel> cir) {
         ModelBakery.BakedCacheKey key = new ModelBakery.BakedCacheKey(arg, arg2.getRotation(), arg2.isUvLocked());
@@ -32,12 +57,32 @@ public abstract class ModelBakerImplMixin {
         if (existing != null) {
             cir.setReturnValue(existing);
         } else {
-            synchronized (this) {
+            synchronized (this.field_40571) {
                 if(debugDynamicModelLoading)
                     ModernFix.LOGGER.info("Baking {}", arg);
-                UnbakedModel iunbakedmodel = this.getModel(arg);
+                UnbakedModel iunbakedmodel;
                 IExtendedModelBakery extendedBakery = (IExtendedModelBakery)this.field_40571;
-                if(iunbakedmodel == extendedBakery.mfix$getUnbakedMissingModel() && debugDynamicModelLoading)
+                if(arg instanceof ModelResourceLocation && arg != ModelBakery.MISSING_MODEL_LOCATION) {
+                    /* to emulate vanilla model loading, treat as top-level */
+                    Optional<Block> blockOpt = BuiltInRegistries.BLOCK.getOptional(new ResourceLocation(arg.getNamespace(), arg.getPath()));
+                    if(blockOpt.isPresent()) {
+                        /* load via lambda for mods that expect blockstate to get loaded */
+                        for(BlockState state : extendedBakery.getBlockStatesForMRL(blockOpt.get().getStateDefinition(), (ModelResourceLocation)arg)) {
+                            try {
+                                blockStateLoaderHandle.invokeExact(this.field_40571, state);
+                            } catch(Throwable e) {
+                                ModernFix.LOGGER.error("Error loading model", e);
+                            }
+                        }
+                    } else {
+                        this.field_40571.loadTopLevel((ModelResourceLocation)arg);
+                    }
+                    iunbakedmodel = this.field_40571.topLevelModels.getOrDefault(arg, extendedBakery.mfix$getUnbakedMissingModel());
+                    // avoid leaks
+                    this.field_40571.topLevelModels.clear();
+                } else
+                    iunbakedmodel = this.getModel(arg);
+                if(iunbakedmodel == extendedBakery.mfix$getUnbakedMissingModel() && arg != ModelBakery.MISSING_MODEL_LOCATION && debugDynamicModelLoading)
                     ModernFix.LOGGER.warn("Model {} not present", arg);
                 // TODO: make sure parent resolution doesn't re-run many times
                 iunbakedmodel.resolveParents(this::getModel);

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

@@ -62,16 +62,23 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
 
     @Shadow @Final @Mutable private BlockColors blockColors;
     @Shadow @Final private static Logger LOGGER;
+
+    @Shadow
+    public abstract void loadTopLevel(ModelResourceLocation modelResourceLocation);
+
     private Cache<ModelBakery.BakedCacheKey, BakedModel> loadedBakedModels;
 
     private Cache<ResourceLocation, UnbakedModel> loadedModels;
 
     private HashMap<ResourceLocation, UnbakedModel> smallLoadingCache = new HashMap<>();
 
+    private boolean ignoreModelLoad;
+
 
     @Redirect(method = "<init>", at = @At(value = "FIELD", opcode = Opcodes.PUTFIELD, target = "Lnet/minecraft/client/resources/model/ModelBakery;blockColors:Lnet/minecraft/client/color/block/BlockColors;"))
     private void replaceTopLevelBakedModels(ModelBakery bakery, BlockColors val) {
         this.blockColors = val;
+        this.ignoreModelLoad = true;
         this.loadedBakedModels = CacheBuilder.newBuilder()
                 .expireAfterAccess(3, TimeUnit.MINUTES)
                 .maximumSize(1000)
@@ -84,7 +91,7 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
                 .maximumSize(1000)
                 .concurrencyLevel(8)
                 .removalListener(this::onModelRemoved)
-                .softValues()
+                //.softValues()
                 .build();
         this.bakedCache = loadedBakedModels.asMap();
         ConcurrentMap<ResourceLocation, UnbakedModel> unbakedCacheBackingMap = loadedModels.asMap();
@@ -103,6 +110,11 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
         this.bakedTopLevelModels = new DynamicBakedModelProvider((ModelBakery)(Object)this, bakedCache);
     }
 
+    @Inject(method = "<init>", at = @At("RETURN"))
+    private void stopIgnore(CallbackInfo ci) {
+        this.ignoreModelLoad = false;
+    }
+
     private <K, V> void onModelRemoved(RemovalNotification<K, V> notification) {
         if(!debugDynamicModelLoading)
             return;
@@ -117,6 +129,9 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
             rl = ((ModelBakery.BakedCacheKey)k).id();
             baked = true;
         }
+        /* can fire when a model is replaced */
+        if(!baked && this.loadedModels.getIfPresent(rl) != null)
+            return;
         ModernFix.LOGGER.warn("Evicted {} model {}", baked ? "baked" : "unbaked", rl);
     }
 
@@ -135,13 +150,13 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
 
     /**
      * @author embeddedt
-     * @reason don't actually load the model.
+     * @reason don't actually load most models
      */
-    @Inject(method = "loadTopLevel", at = @At("HEAD"), cancellable = true)
+    @Redirect(method = "*", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/resources/model/ModelBakery;loadTopLevel(Lnet/minecraft/client/resources/model/ModelResourceLocation;)V"))
-    private void addTopLevelFile(ModelResourceLocation location, CallbackInfo ci) {
+    private void addTopLevelFile(ModelBakery bakery, ModelResourceLocation location) {
-        if(location == MISSING_MODEL_LOCATION)
+        if(location == MISSING_MODEL_LOCATION || !this.ignoreModelLoad) {
-            return; /* needed for FAPI compat */
+            loadTopLevel(location);
-        ci.cancel();
+        }
     }
 
     @Redirect(method = "<init>", at = @At(value = "INVOKE", target = "Ljava/util/Map;forEach(Ljava/util/function/BiConsumer;)V", ordinal = 0))