Block models are now loaded & baked dynamically

Author: embeddedt

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

@@ -0,0 +1,22 @@
+package org.embeddedt.modernfix.common.mixin.perf.dynamic_resources;
+
+import net.minecraft.client.resources.model.ModelBakery;
+import org.embeddedt.modernfix.annotation.ClientOnlyMixin;
+import org.embeddedt.modernfix.dynresources.DynamicModelSystem;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Redirect;
+
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
+import java.util.function.BiFunction;
+
+@Mixin(ModelBakery.class)
+@ClientOnlyMixin
+public class MixinModelBakery {
+    @Redirect(method = "bakeModels", at = @At(value = "INVOKE", target = "Lnet/minecraft/util/thread/ParallelMapTransform;schedule(Ljava/util/Map;Ljava/util/function/BiFunction;Ljava/util/concurrent/Executor;)Ljava/util/concurrent/CompletableFuture;"))
+    private <K, U, V> CompletableFuture<Map<K, V>> dynamicallyBake(Map<K, U> input, BiFunction<K, U, V> baker, Executor executor) {
+        return CompletableFuture.completedFuture(DynamicModelSystem.createDynamicBakedRegistry(input, baker));
+    }
+}

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

@@ -1,7 +1,9 @@
 package org.embeddedt.modernfix.common.mixin.perf.dynamic_resources;
 
+import com.google.common.collect.Maps;
 import it.unimi.dsi.fastutil.objects.Object2IntMap;
 import net.minecraft.client.color.block.BlockColors;
+import net.minecraft.client.renderer.block.model.BlockStateModel;
 import net.minecraft.client.resources.model.BlockStateModelLoader;
 import net.minecraft.client.resources.model.ClientItemInfoLoader;
 import net.minecraft.client.resources.model.ModelManager;
@@ -53,4 +55,15 @@ private static ModelManager.ResolvedModels discoverModelDependencies(
     private static Object2IntMap<BlockState> buildModelGroups(BlockColors blockColors, BlockStateModelLoader.LoadedModels loadedModels) {
         return new DynamicModelSystem.BlockGroupingMap(blockColors, loadedModels);
     }
+
+    /**
+     * @author embeddedt
+     * @reason avoid copying inner map
+     */
+    @Overwrite
+    private static Map<BlockState, BlockStateModel> createBlockStateToModelDispatch(Map<BlockState, BlockStateModel> blockStateModels, BlockStateModel missingModel) {
+        return Maps.asMap(DynamicModelSystem.getAllBlockStates(), state -> {
+            return blockStateModels.getOrDefault(state, missingModel);
+        });
+    }
 }

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

@@ -32,7 +32,9 @@
 import java.io.Reader;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
+import java.util.function.BiFunction;
 import java.util.stream.Collectors;
 
 public class DynamicModelSystem {
@@ -65,6 +67,10 @@ public interface SingleBlockStateEntryLoader {
         BlockStateModelLoader.LoadedModels loadEntry(Identifier identifier, List<Resource> blockstateResources);
     }
 
+    public static Set<BlockState> getAllBlockStates() {
+        return ((IdMapperAccessor<BlockState>)Block.BLOCK_STATE_REGISTRY).getReferenceMap().keySet();
+    }
+
     public static BlockStateModelLoader.LoadedModels createDynamicBlockStateLoadedModels(Map<Identifier, List<Resource>> resourceMap, SingleBlockStateEntryLoader entryLoader) {
         LoadingCache<Identifier, BlockStateModelLoader.LoadedModels> definitionCache = CacheBuilder.newBuilder().softValues().maximumSize(1000).build(new CacheLoader<>() {
             @Override
@@ -77,8 +83,7 @@ public BlockStateModelLoader.LoadedModels load(Identifier key) throws Exception
                 return entryLoader.loadEntry(file, resources);
             }
         });
-        Set<BlockState> allStates = ((IdMapperAccessor<BlockState>)Block.BLOCK_STATE_REGISTRY).getReferenceMap().keySet();
-        return new BlockStateModelLoader.LoadedModels(Maps.asMap(allStates, state -> {
+        return new BlockStateModelLoader.LoadedModels(Maps.asMap(getAllBlockStates(), state -> {
             var identifier = state.getBlock().builtInRegistryHolder().getKey().identifier();
             var loadedModels = definitionCache.getUnchecked(identifier);
             return loadedModels.models().get(state);
@@ -148,4 +153,23 @@ public int getInt(Object key) {
             return -1;
         }
     }
+
+    public static <K, U, V> Map<K, V> createDynamicBakedRegistry(Map<K, U> input, BiFunction<K, U, V> baker) {
+        // TODO: support persistence of overrides
+        LoadingCache<K, Optional<V>> bakedCache = CacheBuilder.newBuilder().softValues().maximumSize(1000).build(new CacheLoader<>() {
+            @Override
+            public Optional<V> load(K key) throws Exception {
+                var unbaked = input.get(key);
+                if (unbaked != null) {
+                    if (DEBUG_DYNAMIC_MODEL_LOADING) {
+                        ModernFix.LOGGER.info("Baking {}", key);
+                    }
+                    return Optional.ofNullable(baker.apply(key, unbaked));
+                } else {
+                    return Optional.empty();
+                }
+            }
+        });
+        return Maps.asMap(input.keySet(), k -> k != null ? bakedCache.getUnchecked(k).orElse(null) : null);
+    }
 }