From dec1b6da597c073fa1de976b2adf1f22064e6575 Mon Sep 17 00:00:00 2001
From: Technici4n <13494793+Technici4n@users.noreply.github.com>
Date: Wed, 31 Jul 2024 23:06:05 +0200
Subject: [PATCH] Use capability caches to improve energy transfer performance
 (#172)

---
 .../java/owmii/powah/block/cable/CableTile.java     |  9 ++++++++-
 .../powah/lib/block/AbstractEnergyStorage.java      | 13 +++++++++++--
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/src/main/java/owmii/powah/block/cable/CableTile.java b/src/main/java/owmii/powah/block/cable/CableTile.java
index 5698b56..8bd3d50 100644
--- 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
index 7df63d2..961cf16 100644
--- 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,7 +21,6 @@
 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>
@@ -26,6 +28,8 @@ public abstract class AbstractEnergyStorage<C extends IEnergyConfig<Tier>, B ext
     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);
                 }
             }