Handle vanilla mapping blockstates to fake StateDefinitions

Closes #621

Co-authored-by: coredex-source <samdhi2323@gmail.com>

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

@@ -0,0 +1,21 @@
+package org.embeddedt.modernfix.common.mixin.perf.dynamic_resources;
+
+import net.minecraft.client.resources.model.BlockStateDefinitions;
+import net.minecraft.resources.Identifier;
+import net.minecraft.world.level.block.Block;
+import net.minecraft.world.level.block.state.BlockState;
+import net.minecraft.world.level.block.state.StateDefinition;
+import org.embeddedt.modernfix.annotation.ClientOnlyMixin;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.gen.Accessor;
+
+import java.util.Map;
+
+@Mixin(BlockStateDefinitions.class)
+@ClientOnlyMixin
+public interface BlockStateDefinitionsAccessor {
+    @Accessor("STATIC_DEFINITIONS")
+    static Map<Identifier, StateDefinition<Block, BlockState>> getStaticDefinitions() {
+        throw new AssertionError();
+    }
+}

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

@@ -0,0 +1,57 @@
+package org.embeddedt.modernfix.dynresources;
+
+import com.google.common.collect.Iterators;
+
+import java.util.AbstractSet;
+import java.util.Iterator;
+import java.util.Set;
+
+/**
+ * Optimized alternative to {@link com.google.common.collect.Sets#union(Set, Set)} in cases where the sets
+ * are known to be disjoint.
+ * @param <T> element type
+ */
+public class DisjointSetUnion<T> extends AbstractSet<T> {
+    private final Set<T> set1, set2;
+
+    public DisjointSetUnion(Set<T> set1, Set<T> set2) {
+        this.set1 = set1;
+        this.set2 = set2;
+        this.assertDisjoint();
+    }
+
+    private void assertDisjoint() {
+        Set<T> iterate = set1.size() < set2.size() ? set1 : set2;
+        Set<T> contains = set1 == iterate ? set2 : set1;
+        for (T obj : iterate) {
+            if (contains.contains(obj)) {
+                throw new IllegalArgumentException("Provided sets are not disjoint");
+            }
+        }
+    }
+
+    @Override
+    public Iterator<T> iterator() {
+        return Iterators.concat(set1.iterator(), set2.iterator());
+    }
+
+    @Override
+    public int size() {
+        return set1.size() + set2.size();
+    }
+
+    @Override
+    public boolean remove(Object o) {
+        return set1.remove(o) || set2.remove(o);
+    }
+
+    @Override
+    public boolean contains(Object o) {
+        return set1.contains(o) || set2.contains(o);
+    }
+
+    @Override
+    public int hashCode() {
+        return set1.hashCode() + set2.hashCode();
+    }
+}

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

@@ -4,6 +4,7 @@ import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
 import com.google.common.collect.Maps;
+import com.google.common.collect.Sets;
 import it.unimi.dsi.fastutil.objects.AbstractObject2IntMap;
 import it.unimi.dsi.fastutil.objects.Object2IntMap;
 import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
@@ -28,10 +29,13 @@ import net.minecraft.world.level.block.state.BlockState;
 import net.neoforged.neoforge.client.model.UnbakedModelParser;
 import net.neoforged.neoforge.client.model.standalone.StandaloneModelLoader;
 import org.embeddedt.modernfix.ModernFix;
+import org.embeddedt.modernfix.common.mixin.perf.dynamic_resources.BlockStateDefinitionsAccessor;
 import org.embeddedt.modernfix.common.mixin.perf.dynamic_resources.IdMapperAccessor;
 import org.embeddedt.modernfix.common.mixin.perf.dynamic_resources.ModelDiscoveryAccessor;
 
 import java.io.Reader;
+import java.util.AbstractSet;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -85,8 +89,17 @@ public class DynamicModelSystem {
                 return entryLoader.loadEntry(file, resources);
             }
         });
-        return new BlockStateModelLoader.LoadedModels(Maps.asMap(getAllBlockStates(), state -> {
+        var staticDefinitions = BlockStateDefinitionsAccessor.getStaticDefinitions();
-            var identifier = state.getBlock().builtInRegistryHolder().getKey().identifier();
+        var staticIdentifiers = staticDefinitions.entrySet()
+                .stream()
+                .flatMap(e -> e.getValue().getPossibleStates().stream().map(s -> Map.entry(s, e.getKey())))
+                .collect(Collectors.toUnmodifiableMap(Map.Entry::getKey, Map.Entry::getValue));
+        var blockStateSet = new DisjointSetUnion<>(getAllBlockStates(), staticIdentifiers.keySet());
+        return new BlockStateModelLoader.LoadedModels(Maps.asMap(blockStateSet, state -> {
+            var identifier = staticIdentifiers.get(state);
+            if (identifier == null) {
+                identifier = state.getBlock().builtInRegistryHolder().getKey().identifier();
+            }
             var loadedModels = definitionCache.getUnchecked(identifier);
             return loadedModels.models().get(state);
         }));