Use capability caches to improve energy transfer performance

Technici4n · Jul 31, 2024 · 2a8ba36 · 2a8ba36
1 parent fe13cf9
commit 2a8ba36
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 3 deletions.
diff --git a/src/main/java/owmii/powah/block/cable/CableTile.java b/src/main/java/owmii/powah/block/cable/CableTile.java
@@ -12,7 +12,9 @@
 import net.minecraft.world.level.ChunkPos;
 import net.minecraft.world.level.Level;
 import net.minecraft.world.level.block.state.BlockState;
+import net.neoforged.neoforge.capabilities.BlockCapabilityCache;
 import net.neoforged.neoforge.capabilities.Capabilities;
+import net.neoforged.neoforge.energy.IEnergyStorage;
 import org.apache.commons.lang3.mutable.MutableBoolean;
 import org.jetbrains.annotations.Nullable;
 import owmii.powah.block.Tier;
@@ -37,6 +39,8 @@ public class CableTile extends AbstractEnergyStorage<CableConfig, CableBlock> im
      */
     protected MutableBoolean netInsertionGuard = new MutableBoolean(false);
     protected int startIndex = 0;
+    @SuppressWarnings("unchecked")
+    private final BlockCapabilityCache<IEnergyStorage, @Nullable Direction>[] capabilityCaches = new BlockCapabilityCache[6];
 
     public CableTile(BlockPos pos, BlockState state, Tier variant) {
         super(Tiles.CABLE.get(), pos, state, variant);
@@ -188,7 +192,10 @@ private long pushEnergy(long maxReceive, boolean simulate, @Nullable Direction d
     }
 
     private long receive(Level level, BlockPos pos, Direction side, long amount, boolean simulate) {
-        var energy = level.getCapability(Capabilities.EnergyStorage.BLOCK, pos, side);
+        if (capabilityCaches[side.ordinal()] == null) {
+            capabilityCaches[side.ordinal()] = BlockCapabilityCache.create(Capabilities.EnergyStorage.BLOCK, (ServerLevel) level, pos, side);
+        }
+        var energy = capabilityCaches[side.ordinal()].getCapability();
         return energy != null ? energy.receiveEnergy(Ints.saturatedCast(amount), simulate) : 0;
     }
 

diff --git a/src/main/java/owmii/powah/lib/block/AbstractEnergyStorage.java b/src/main/java/owmii/powah/lib/block/AbstractEnergyStorage.java
@@ -5,9 +5,12 @@
 import net.minecraft.core.Direction;
 import net.minecraft.core.HolderLookup;
 import net.minecraft.nbt.CompoundTag;
+import net.minecraft.server.level.ServerLevel;
 import net.minecraft.world.level.Level;
 import net.minecraft.world.level.block.entity.BlockEntityType;
 import net.minecraft.world.level.block.state.BlockState;
+import net.neoforged.neoforge.capabilities.BlockCapabilityCache;
+import net.neoforged.neoforge.capabilities.Capabilities;
 import net.neoforged.neoforge.energy.IEnergyStorage;
 import org.jetbrains.annotations.Nullable;
 import owmii.powah.block.Tier;
@@ -18,14 +21,15 @@
 import owmii.powah.lib.logistics.energy.SideConfig;
 import owmii.powah.lib.registry.IVariant;
 import owmii.powah.util.ChargeUtil;
-import owmii.powah.util.EnergyUtil;
 import owmii.powah.util.Util;
 
 public abstract class AbstractEnergyStorage<C extends IEnergyConfig<Tier>, B extends AbstractEnergyBlock<C, B>> extends AbstractTickableTile<Tier, B>
         implements IRedstoneInteract {
     protected final SideConfig sideConfig = new SideConfig(this);
     protected final Energy energy = Energy.create(0);
     private final @Nullable IEnergyStorage[] externalAdapters = new IEnergyStorage[Direction.values().length + 1];
+    @SuppressWarnings("unchecked")
+    private final BlockCapabilityCache<IEnergyStorage, @Nullable Direction>[] capabilityCaches = new BlockCapabilityCache[6];
 
     public AbstractEnergyStorage(BlockEntityType<?> type, BlockPos pos, BlockState state) {
         this(type, pos, state, IVariant.getEmpty());
@@ -87,8 +91,13 @@ protected long extractFromSides(Level world) {
         if (!isRemote()) {
             for (Direction side : Direction.values()) {
                 if (canExtractEnergy(side)) {
+                    if (capabilityCaches[side.ordinal()] == null) {
+                        capabilityCaches[side.ordinal()] = BlockCapabilityCache.create(Capabilities.EnergyStorage.BLOCK, (ServerLevel) world,
+                                worldPosition.relative(side), side.getOpposite());
+                    }
                     long amount = Math.min(getEnergyTransfer(), getEnergy().getStored());
-                    long toExtract = EnergyUtil.pushEnergy(world, worldPosition.relative(side), side.getOpposite(), amount);
+                    var cap = capabilityCaches[side.ordinal()].getCapability();
+                    long toExtract = cap == null ? 0 : cap.receiveEnergy(Ints.saturatedCast(amount), false);
                     extracted += extractEnergy(Util.safeInt(toExtract), false, side);
                 }
             }