Rework wall block caching to cache collision shapes as well

Author: embeddedt

common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/deduplicate_wall_shapes/WallBlockMixin.java

@@ -1,6 +1,8 @@
 package org.embeddedt.modernfix.common.mixin.perf.deduplicate_wall_shapes;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.mojang.datafixers.util.Pair;
 import net.minecraft.world.level.block.Block;
 import net.minecraft.world.level.block.WallBlock;
 import net.minecraft.world.level.block.state.BlockState;
@@ -12,43 +14,47 @@
 import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
 
-import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
 
+/**
+ * Most wall blocks use the default set of vanilla properties, and the default sizes for their shapes. This means
+ * there is no need to reconstruct a separate VoxelShape instance for each wall, we can just repurpose the
+ * same shape instances. To do this we can cache a mapping between a state (represented only as its prop->value map)
+ * and the desired shape, and generate the BlockState->VoxelShape map from this for each block.
+ */
 @Mixin(WallBlock.class)
 public abstract class WallBlockMixin extends Block {
-    private static Map<ImmutableMap<Property<?>, Comparable<?>>, VoxelShape> CACHE_BY_PROPERTIES = new HashMap<>();
-    private static StateDefinition<Block, BlockState> CACHED_DEFINITION = null;
-    private static float[] CACHED_FLOATS = null;
+    private static Map<ImmutableList<Float>, Pair<Map<ImmutableMap<Property<?>, Comparable<?>>, VoxelShape>, StateDefinition<Block, BlockState>>> CACHE_BY_SHAPE_VALS = new HashMap<>();
 
     public WallBlockMixin(Properties properties) {
         super(properties);
     }
 
     @Inject(method = "makeShapes", at = @At("HEAD"), cancellable = true)
     private synchronized void useCachedShapeMap(float f1, float f2, float f3, float f4, float f5, float f6, CallbackInfoReturnable<Map<BlockState, VoxelShape>> cir) {
-        if(CACHED_DEFINITION != null) {
-            // check if this state container's properties exactly match the one we used for the cache
-            if(CACHED_DEFINITION.getProperties().equals(this.stateDefinition.getProperties()) && Arrays.equals(CACHED_FLOATS, new float[] { f1, f2, f3, f4, f5, f6 })) {
-                ImmutableMap.Builder<BlockState, VoxelShape> builder = ImmutableMap.builder();
-                for(BlockState state : this.stateDefinition.getPossibleStates()) {
-                    builder.put(state, CACHE_BY_PROPERTIES.get(state.getValues()));
-                }
-                cir.setReturnValue(builder.build());
-            }
+        ImmutableList<Float> key = ImmutableList.of(f1, f2, f3, f4, f5, f6);
+        Pair<Map<ImmutableMap<Property<?>, Comparable<?>>, VoxelShape>, StateDefinition<Block, BlockState>> cache = CACHE_BY_SHAPE_VALS.get(key);
+        // require the properties to be identical
+        if(cache == null || !cache.getSecond().getProperties().equals(this.stateDefinition.getProperties()))
+            return;
+        ImmutableMap.Builder<BlockState, VoxelShape> builder = ImmutableMap.builder();
+        for(BlockState state : this.stateDefinition.getPossibleStates()) {
+            builder.put(state, cache.getFirst().get(state.getValues()));
         }
+        cir.setReturnValue(builder.build());
     }
 
     @Inject(method = "makeShapes", at = @At("RETURN"))
     private synchronized void storeCachedShapesByProperty(float f1, float f2, float f3, float f4, float f5, float f6, CallbackInfoReturnable<Map<BlockState, VoxelShape>> cir) {
-        if(CACHE_BY_PROPERTIES.size() == 0) {
+        ImmutableList<Float> key = ImmutableList.of(f1, f2, f3, f4, f5, f6);
+        if(!CACHE_BY_SHAPE_VALS.containsKey(key)) {
+            Map<ImmutableMap<Property<?>, Comparable<?>>, VoxelShape> cacheByProperties = new HashMap<>();
             Map<BlockState, VoxelShape> shapeMap = cir.getReturnValue();
             for(Map.Entry<BlockState, VoxelShape> entry : shapeMap.entrySet()) {
-                CACHE_BY_PROPERTIES.put(entry.getKey().getValues(), entry.getValue());
+                cacheByProperties.put(entry.getKey().getValues(), entry.getValue());
             }
-            CACHED_FLOATS = new float[] { f1, f2, f3, f4, f5, f6 };
-            CACHED_DEFINITION = this.stateDefinition;
+            CACHE_BY_SHAPE_VALS.put(key, Pair.of(cacheByProperties, this.stateDefinition));
         }
     }
 }