From 8e2d6968a0d7e51da7d3d2e1a1b89f36d2eb5233 Mon Sep 17 00:00:00 2001
From: embeddedt <42941056+embeddedt@users.noreply.github.com>
Date: Sat, 11 Apr 2026 11:50:05 -0400
Subject: [PATCH] Improve dynamic resources performance when mods iterate the
 whole model map

---
 .../dynresources/DynamicModelSystem.java      | 21 ++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/src/main/java/org/embeddedt/modernfix/dynresources/DynamicModelSystem.java b/src/main/java/org/embeddedt/modernfix/dynresources/DynamicModelSystem.java
index 1a84fe2a..4c6b0454 100644
--- a/src/main/java/org/embeddedt/modernfix/dynresources/DynamicModelSystem.java
+++ b/src/main/java/org/embeddedt/modernfix/dynresources/DynamicModelSystem.java
@@ -20,6 +20,7 @@ import net.minecraft.client.resources.model.ResolvedModel;
 import net.minecraft.client.resources.model.UnbakedModel;
 import net.minecraft.client.resources.model.cuboid.ItemModelGenerator;
 import net.minecraft.client.resources.model.cuboid.MissingCuboidModel;
+import net.minecraft.core.registries.BuiltInRegistries;
 import net.minecraft.resources.FileToIdConverter;
 import net.minecraft.resources.Identifier;
 import net.minecraft.server.packs.resources.Resource;
@@ -73,7 +74,25 @@ public class DynamicModelSystem {
     }
 
     public static Set<BlockState> getAllBlockStates() {
-        return ReferenceSets.unmodifiable(((IdMapperAccessor<BlockState>) Block.BLOCK_STATE_REGISTRY).getReferenceMap().keySet());
+        var blockStateSet = ((IdMapperAccessor<BlockState>) Block.BLOCK_STATE_REGISTRY).getReferenceMap().keySet();
+        return new AbstractSet<>() {
+            @Override
+            public Iterator<BlockState> iterator() {
+                // We explicitly override iterator() and handle it differently so that mods iterating the maps
+                // are likely to work with the same block many times in a row, which hits our caches better
+                return BuiltInRegistries.BLOCK.stream().flatMap(b -> b.getStateDefinition().getPossibleStates().stream()).iterator();
+            }
+
+            @Override
+            public boolean contains(Object o) {
+                return blockStateSet.contains(o);
+            }
+
+            @Override
+            public int size() {
+                return blockStateSet.size();
+            }
+        };
     }
 
     public static BlockStateModelLoader.LoadedModels createDynamicBlockStateLoadedModels(Map<Identifier, List<Resource>> resourceMap, SingleBlockStateEntryLoader entryLoader) {