Speed up blockstate cache deduplication

For most blocks the cache does not change with tags, so checking the previous cache if it exists will generally save a Map call

common/src/main/java/malte0811/ferritecore/impl/BlockStateCacheImpl.java

@@ -1,78 +1,125 @@
 package malte0811.ferritecore.impl;
 
-import com.mojang.datafixers.util.Pair;
 import it.unimi.dsi.fastutil.objects.Object2ObjectOpenCustomHashMap;
 import malte0811.ferritecore.hash.VoxelShapeArrayHash;
 import malte0811.ferritecore.hash.VoxelShapeHash;
+import malte0811.ferritecore.mixin.blockstatecache.BlockStateCacheAccess;
 import malte0811.ferritecore.mixin.blockstatecache.VSArrayAccess;
+import malte0811.ferritecore.mixin.blockstatecache.VSSplitAccess;
 import malte0811.ferritecore.mixin.blockstatecache.VoxelShapeAccess;
-import malte0811.ferritecore.util.LastAccessedCache;
-import net.minecraft.block.AbstractBlock;
-import net.minecraft.block.BlockState;
-import net.minecraft.util.Direction;
-import net.minecraft.util.math.BlockPos;
-import net.minecraft.util.math.shapes.ISelectionContext;
+import malte0811.ferritecore.util.Constants;
+import net.minecraft.util.LazyValue;
+import net.minecraft.util.math.shapes.SplitVoxelShape;
 import net.minecraft.util.math.shapes.VoxelShape;
 import net.minecraft.util.math.shapes.VoxelShapeArray;
-import net.minecraft.util.math.shapes.VoxelShapes;
-import net.minecraft.world.IBlockReader;
+import org.apache.commons.lang3.tuple.Pair;
 
+import javax.annotation.Nullable;
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.reflect.Field;
+import java.util.Map;
 import java.util.function.Function;
 
+import static net.minecraft.block.AbstractBlock.AbstractBlockState;
+
 public class BlockStateCacheImpl {
-    private static final Direction[] DIRECTIONS = Direction.values();
-    public static final Object2ObjectOpenCustomHashMap<VoxelShapeArray, VoxelShapeArray> CACHE_COLLIDE =
-            new Object2ObjectOpenCustomHashMap<>(VoxelShapeArrayHash.INSTANCE);
-    public static final LastAccessedCache<VoxelShape, VoxelShape[]> CACHE_PROJECT = new LastAccessedCache<>(
-            VoxelShapeHash.INSTANCE, vs -> {
-        VoxelShape[] result = new VoxelShape[DIRECTIONS.length];
-        for (Direction d : DIRECTIONS) {
-            result[d.ordinal()] = VoxelShapes.getFaceShape(vs, d);
+    public static final Map<VoxelShapeArray, VoxelShapeArray> CACHE_COLLIDE = new Object2ObjectOpenCustomHashMap<>(
+            VoxelShapeArrayHash.INSTANCE
+    );
+    // Maps a shape to the "canonical instance" of that shape and its side projections
+    public static final Map<VoxelShape, Pair<VoxelShape, VoxelShape[]>> CACHE_PROJECT =
+            new Object2ObjectOpenCustomHashMap<>(VoxelShapeHash.INSTANCE);
+
+    // Get the cache from a blockstate. Mixin does not handle private inner classes too well, so method handles and
+    // manual remapping it is
+    private static final LazyValue<Function<AbstractBlockState, BlockStateCacheAccess>> GET_CACHE =
+            new LazyValue<>(() -> {
+                try {
+                    Field cacheField = AbstractBlockState.class.getDeclaredField(Constants.blockstateCacheFieldName);
+                    cacheField.setAccessible(true);
+                    MethodHandle getter = MethodHandles.lookup().unreflectGetter(cacheField);
+                    return state -> {
+                        try {
+                            return (BlockStateCacheAccess) getter.invoke(state);
+                        } catch (Throwable throwable) {
+                            throw new RuntimeException(throwable);
+                        }
+                    };
+                } catch (NoSuchFieldException | IllegalAccessException e) {
+                    throw new RuntimeException(e);
+                }
+            });
+    // Is set to the previous cache used by a state before updating the cache. If the new cache has shapes equivalent to
+    // the ones in the old cache, we don't need to go through the map since the old one already had deduplicated shapes
+    private static final ThreadLocal<BlockStateCacheAccess> LAST_CACHE = new ThreadLocal<>();
+
+    // Calls before the cache for <code>state</code> is (re-)populated
+    public static void deduplicateCachePre(AbstractBlockState state) {
+        LAST_CACHE.set(GET_CACHE.getValue().apply(state));
+    }
+
+    // Calls after the cache for <code>state</code> is (re-)populated
+    public static void deduplicateCachePost(AbstractBlockState state) {
+        BlockStateCacheAccess newCache = GET_CACHE.getValue().apply(state);
+        if (newCache != null) {
+            final BlockStateCacheAccess oldCache = LAST_CACHE.get();
+            deduplicateCollisionShape(newCache, oldCache);
+            deduplicateRenderShapes(newCache, oldCache);
+            LAST_CACHE.set(null);
         }
-        return result;
     }
-    );
-    public static int collideCalls = 0;
-    public static int projectCalls = 0;
 
-    public static void resetCaches() {
-        CACHE_COLLIDE.clear();
-        CACHE_COLLIDE.trim();
-        collideCalls = 0;
-        CACHE_PROJECT.clear();
-        projectCalls = 0;
+    private static void deduplicateCollisionShape(
+            BlockStateCacheAccess newCache, @Nullable BlockStateCacheAccess oldCache
+    ) {
+        VoxelShape dedupedCollisionShape;
+        if (oldCache != null && VoxelShapeHash.INSTANCE.equals(
+                oldCache.getCollisionShape(), newCache.getCollisionShape()
+        )) {
+            dedupedCollisionShape = oldCache.getCollisionShape();
+        } else {
+            dedupedCollisionShape = newCache.getCollisionShape();
+            if (dedupedCollisionShape instanceof VoxelShapeArray) {
+                dedupedCollisionShape = CACHE_COLLIDE.computeIfAbsent(
+                        (VoxelShapeArray) dedupedCollisionShape, Function.identity()
+                );
+            }
+        }
+        replaceInternals(dedupedCollisionShape, newCache.getCollisionShape());
+        newCache.setCollisionShape(dedupedCollisionShape);
     }
 
-    /**
-     * Returns an interned/deduplicated version of the voxel shape returned by
-     * {@link AbstractBlock#getCollisionShape(BlockState, IBlockReader, BlockPos, ISelectionContext)}
-     * and replaces the internals of the original returned shape (see {@link BlockStateCacheImpl#replaceInternals(VoxelShapeArray, VoxelShapeArray)})
-     */
-    public static VoxelShape redirectGetCollisionShape(
-            AbstractBlock block, BlockState state, IBlockReader worldIn, BlockPos pos, ISelectionContext context
+    private static void deduplicateRenderShapes(
+            BlockStateCacheAccess newCache, @Nullable BlockStateCacheAccess oldCache
     ) {
-        VoxelShape baseResult = block.getCollisionShape(state, worldIn, pos, context);
-        if (!(baseResult instanceof VoxelShapeArray)) {
-            return baseResult;
+        final VoxelShape newRenderShape = getRenderShape(newCache.getRenderShapes());
+        if (newRenderShape == null) {
+            return;
+        }
+        Pair<VoxelShape, VoxelShape[]> dedupedRenderShapes = null;
+        if (oldCache != null) {
+            final VoxelShape oldRenderShape = getRenderShape(oldCache.getRenderShapes());
+            if (VoxelShapeHash.INSTANCE.equals(newRenderShape, oldRenderShape)) {
+                dedupedRenderShapes = Pair.of(oldRenderShape, oldCache.getRenderShapes());
+            }
+        }
+        if (dedupedRenderShapes == null) {
+            // Who thought that this was a good interface for putIfAbsent…
+            Pair<VoxelShape, VoxelShape[]> newPair = Pair.of(newRenderShape, newCache.getRenderShapes());
+            dedupedRenderShapes = CACHE_PROJECT.putIfAbsent(newRenderShape, newPair);
+            if (dedupedRenderShapes == null) {
+                dedupedRenderShapes = newPair;
+            }
         }
-        VoxelShapeArray baseArray = (VoxelShapeArray) baseResult;
-        ++collideCalls;
-        VoxelShapeArray resultArray = CACHE_COLLIDE.computeIfAbsent(baseArray, Function.identity());
-        replaceInternals(resultArray, baseArray);
-        return resultArray;
+        replaceInternals(dedupedRenderShapes.getLeft(), newRenderShape);
+        newCache.setRenderShapes(dedupedRenderShapes.getRight());
     }
 
-    /**
-     * Returns the interned/deduplicated "face shape" of the given shape, and replaces the internals of the original
-     * shape if necessary/appropriate
-     */
-    public static VoxelShape redirectFaceShape(VoxelShape shape, Direction face) {
-        ++projectCalls;
-        Pair<VoxelShape, VoxelShape[]> sides = CACHE_PROJECT.get(shape);
-        if (sides.getFirst() instanceof VoxelShapeArray && shape instanceof VoxelShapeArray) {
-            replaceInternals((VoxelShapeArray) sides.getFirst(), (VoxelShapeArray) shape);
+    private static void replaceInternals(VoxelShape toKeep, VoxelShape toReplace) {
+        if (toKeep instanceof VoxelShapeArray && toReplace instanceof VoxelShapeArray) {
+            replaceInternals((VoxelShapeArray) toKeep, (VoxelShapeArray) toReplace);
         }
-        return sides.getSecond()[face.ordinal()];
     }
 
     public static void replaceInternals(VoxelShapeArray toKeep, VoxelShapeArray toReplace) {
@@ -100,4 +147,16 @@ public class BlockStateCacheImpl {
     private static VoxelShapeAccess accessVS(VoxelShape a) {
         return (VoxelShapeAccess) a;
     }
+
+    @Nullable
+    private static VoxelShape getRenderShape(@Nullable VoxelShape[] projected) {
+        if (projected != null) {
+            for (VoxelShape side : projected) {
+                if (side instanceof SplitVoxelShape) {
+                    return ((VSSplitAccess) side).getShape();
+                }
+            }
+        }
+        return null;
+    }
 }

common/src/main/java/malte0811/ferritecore/mixin/blockstatecache/AbstractBlockStateMixin.java

@@ -0,0 +1,21 @@
+package malte0811.ferritecore.mixin.blockstatecache;
+
+import malte0811.ferritecore.impl.BlockStateCacheImpl;
+import net.minecraft.block.AbstractBlock;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+
+@Mixin(AbstractBlock.AbstractBlockState.class)
+public class AbstractBlockStateMixin {
+    @Inject(method = "cacheState", at = @At("HEAD"))
+    public void cacheStateHead(CallbackInfo ci) {
+        BlockStateCacheImpl.deduplicateCachePre((AbstractBlock.AbstractBlockState) (Object) this);
+    }
+
+    @Inject(method = "cacheState", at = @At("TAIL"))
+    public void cacheStateTail(CallbackInfo ci) {
+        BlockStateCacheImpl.deduplicateCachePost((AbstractBlock.AbstractBlockState) (Object) this);
+    }
+}

common/src/main/java/malte0811/ferritecore/mixin/blockstatecache/BlockStateCacheAccess.java

@@ -0,0 +1,26 @@
+package malte0811.ferritecore.mixin.blockstatecache;
+
+import net.minecraft.util.math.shapes.VoxelShape;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Mutable;
+import org.spongepowered.asm.mixin.gen.Accessor;
+
+import javax.annotation.Nullable;
+
+@Mixin(targets = "net.minecraft.block.AbstractBlock$AbstractBlockState$Cache")
+public interface BlockStateCacheAccess {
+    @Accessor
+    VoxelShape getCollisionShape();
+
+    @Accessor
+    @Mutable
+    void setCollisionShape(VoxelShape newShape);
+
+    @Accessor
+    @Nullable
+    VoxelShape[] getRenderShapes();
+
+    @Accessor
+    @Mutable
+    void setRenderShapes(@Nullable VoxelShape[] newShapes);
+}

common/src/main/java/malte0811/ferritecore/mixin/blockstatecache/BlockStateCacheMixin.java

@@ -1,40 +0,0 @@
-package malte0811.ferritecore.mixin.blockstatecache;
-
-import malte0811.ferritecore.impl.BlockStateCacheImpl;
-import net.minecraft.block.Block;
-import net.minecraft.block.BlockState;
-import net.minecraft.util.Direction;
-import net.minecraft.util.math.BlockPos;
-import net.minecraft.util.math.shapes.ISelectionContext;
-import net.minecraft.util.math.shapes.VoxelShape;
-import net.minecraft.world.IBlockReader;
-import org.spongepowered.asm.mixin.Mixin;
-import org.spongepowered.asm.mixin.injection.At;
-import org.spongepowered.asm.mixin.injection.Redirect;
-
-@Mixin(targets = "net.minecraft.block.AbstractBlock$AbstractBlockState$Cache")
-public class BlockStateCacheMixin {
-    @Redirect(
-            method = "<init>",
-            at = @At(
-                    value = "INVOKE",
-                    target = "Lnet/minecraft/util/math/shapes/VoxelShapes;getFaceShape(Lnet/minecraft/util/math/shapes/VoxelShape;Lnet/minecraft/util/Direction;)Lnet/minecraft/util/math/shapes/VoxelShape;"
-            )
-    )
-    private VoxelShape redirectFaceShape(VoxelShape shape, Direction face) {
-        return BlockStateCacheImpl.redirectFaceShape(shape, face);
-    }
-
-    @Redirect(
-            method = "<init>",
-            at = @At(
-                    value = "INVOKE",
-                    target = "Lnet/minecraft/block/Block;getCollisionShape(Lnet/minecraft/block/BlockState;Lnet/minecraft/world/IBlockReader;Lnet/minecraft/util/math/BlockPos;Lnet/minecraft/util/math/shapes/ISelectionContext;)Lnet/minecraft/util/math/shapes/VoxelShape;"
-            )
-    )
-    private VoxelShape redirectGetCollisionShape(
-            Block block, BlockState state, IBlockReader worldIn, BlockPos pos, ISelectionContext context
-    ) {
-        return BlockStateCacheImpl.redirectGetCollisionShape(block, state, worldIn, pos, context);
-    }
-}

common/src/main/java/malte0811/ferritecore/mixin/blockstatecache/Config.java

@@ -10,8 +10,8 @@ public class Config extends FerriteMixinConfig {
     @Override
     protected List<String> getAllMixins() {
         return ImmutableList.of(
-                "BlockStateCacheMixin", "VoxelShapeAccess", "VSArrayAccess", "VSPBitSetAccess", "VSPSplitAccess",
-                "VSSplitAccess"
+                "AbstractBlockStateMixin", "BlockStateCacheAccess", "VoxelShapeAccess", "VSArrayAccess",
+                "VSPBitSetAccess", "VSPSplitAccess", "VSSplitAccess"
         );
     }
 

common/src/main/java/malte0811/ferritecore/util/Constants.java

@@ -2,4 +2,5 @@ package malte0811.ferritecore.util;
 
 public class Constants {
     public static final String MODID = "ferritecore";
+    public static String blockstateCacheFieldName;
 }

common/src/main/java/malte0811/ferritecore/util/LastAccessedCache.java

@@ -1,46 +0,0 @@
-package malte0811.ferritecore.util;
-
-import com.mojang.datafixers.util.Pair;
-import it.unimi.dsi.fastutil.Hash;
-import it.unimi.dsi.fastutil.objects.Object2ObjectOpenCustomHashMap;
-
-import java.util.function.Function;
-
-/**
- * A cache which checks if the last accessed value is being accessed again before checking the main map
- *
- * @param <K> Key type
- * @param <V> Value type
- */
-public class LastAccessedCache<K, V> {
-    private final Object2ObjectOpenCustomHashMap<K, V> mainMap;
-    private final Function<K, V> createValue;
-    private final Hash.Strategy<K> strategy;
-    private Pair<K, V> lastAccessed;
-
-    public LastAccessedCache(Hash.Strategy<K> strategy, Function<K, V> createValue) {
-        this.strategy = strategy;
-        this.mainMap = new Object2ObjectOpenCustomHashMap<>(strategy);
-        this.createValue = createValue;
-    }
-
-    public Pair<K, V> get(K key) {
-        final Pair<K, V> last = lastAccessed;
-        if (last != null && strategy.equals(last.getFirst(), key)) {
-            return last;
-        } else {
-            final V result = mainMap.computeIfAbsent(key, createValue);
-            return lastAccessed = Pair.of(key, result);
-        }
-    }
-
-    public void clear() {
-        lastAccessed = null;
-        mainMap.clear();
-        mainMap.trim();
-    }
-
-    public int size() {
-        return mainMap.size();
-    }
-}

common/src/main/resources/ferritecore.blockstatecache.mixin.json

@@ -10,7 +10,8 @@
   "minVersion": "0.8",
   "plugin": "malte0811.ferritecore.mixin.blockstatecache.Config",
   "mixins": [
-    "BlockStateCacheMixin",
+    "AbstractBlockStateMixin",
+    "BlockStateCacheAccess",
     "VoxelShapeAccess",
     "VSArrayAccess",
     "VSPBitSetAccess",

fabric/src/main/java/malte0811/ferritecore/ModMain.java

@@ -0,0 +1,22 @@
+package malte0811.ferritecore;
+
+import malte0811.ferritecore.util.Constants;
+import net.fabricmc.api.ModInitializer;
+import net.fabricmc.loader.api.FabricLoader;
+
+public class ModMain implements ModInitializer {
+    @Override
+    public void onInitialize() {
+        Constants.blockstateCacheFieldName = FabricLoader.getInstance()
+                .getMappingResolver()
+                .mapFieldName(
+                        "intermediary",
+                        // AbstractBlockState
+                        "net.minecraft.class_4970$class_4971",
+                        // cache
+                        "field_23166",
+                        // AbstractBlockState.Cache
+                        "Lnet/minecraft/class_4970$class_4971$class_3752;"
+                );
+    }
+}

fabric/src/main/java/malte0811/ferritecore/mixin/fabric/CacheCallbackMixin.java

@@ -1,16 +0,0 @@
-package malte0811.ferritecore.mixin.fabric;
-
-import malte0811.ferritecore.impl.BlockStateCacheImpl;
-import net.minecraft.block.Blocks;
-import org.spongepowered.asm.mixin.Mixin;
-import org.spongepowered.asm.mixin.injection.At;
-import org.spongepowered.asm.mixin.injection.Inject;
-import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
-
-@Mixin(Blocks.class)
-public class CacheCallbackMixin {
-    @Inject(method = "cacheBlockStates", at = @At("TAIL"))
-    private static void afterCacheRebuild(CallbackInfo ci) {
-        BlockStateCacheImpl.resetCaches();
-    }
-}

fabric/src/main/resources/fabric.mod.json

@@ -9,7 +9,11 @@
   ],
   "license": "MIT",
   "environment": "*",
-  "entrypoints": {},
+  "entrypoints": {
+    "main": [
+      "malte0811.ferritecore.ModMain"
+    ]
+  },
   "depends": {
     "fabricloader": ">=0.7.4",
     "fabric": "*",

fabric/src/main/resources/ferritecore.fabric.mixin.json

@@ -2,15 +2,11 @@
   "required": true,
   "package": "malte0811.ferritecore.mixin.fabric",
   "compatibilityLevel": "JAVA_8",
-  "refmap": "ferritecore-fabric-refmap.json",
   "client": [
+    "MinecraftMixin"
   ],
   "injectors": {
     "defaultRequire": 1
   },
-  "minVersion": "0.8",
-  "mixins": [
-    "CacheCallbackMixin",
-    "MinecraftMixin"
-  ]
+  "minVersion": "0.8"
 }

forge/src/main/java/malte0811/ferritecore/ModMainForge.java

@@ -1,24 +1,17 @@
 package malte0811.ferritecore;
 
-import malte0811.ferritecore.impl.BlockStateCacheImpl;
+import cpw.mods.modlauncher.api.INameMappingService;
 import malte0811.ferritecore.util.Constants;
-import net.minecraftforge.event.TagsUpdatedEvent;
-import net.minecraftforge.eventbus.api.SubscribeEvent;
 import net.minecraftforge.fml.common.Mod;
-import net.minecraftforge.fml.event.lifecycle.FMLModIdMappingEvent;
+import net.minecraftforge.fml.common.ObfuscationReflectionHelper;
 
 @Mod(Constants.MODID)
 @Mod.EventBusSubscriber(bus = Mod.EventBusSubscriber.Bus.FORGE)
 public class ModMainForge {
-    // Caches are populated in two places: a) In ITagCollectionSupplier#updateTags (which triggers this event)
-    @SubscribeEvent
-    public static void onTagReloadVanilla(TagsUpdatedEvent.VanillaTagTypes ignored) {
-        BlockStateCacheImpl.resetCaches();
-    }
 
-    // b) Via ForgeRegistry#bake, which usually triggers this event
-    @SubscribeEvent
-    public static void onModIdMapping(FMLModIdMappingEvent ignored) {
-        BlockStateCacheImpl.resetCaches();
+    public ModMainForge() {
+        Constants.blockstateCacheFieldName = ObfuscationReflectionHelper.remapName(
+                INameMappingService.Domain.FIELD, "field_215707_c"
+        );
     }
 }