[1.20.6] Add ComponentEnergyStorage for using IEnergyStorage with Dat…

…a Components (#985)

src/main/java/net/neoforged/neoforge/energy/ComponentEnergyStorage.java

@@ -0,0 +1,123 @@
+/*
+ * Copyright (c) NeoForged and contributors
+ * SPDX-License-Identifier: LGPL-2.1-only
+ */
+
+package net.neoforged.neoforge.energy;
+
+import net.minecraft.core.component.DataComponentType;
+import net.minecraft.util.Mth;
+import net.minecraft.world.item.ItemStack;
+import net.neoforged.neoforge.capabilities.ICapabilityProvider;
+import net.neoforged.neoforge.capabilities.RegisterCapabilitiesEvent;
+import net.neoforged.neoforge.common.MutableDataComponentHolder;
+
+/**
+ * Variant of {@link EnergyStorage} for use with data components.
+ * <p>
+ * The actual data storage is managed by a data component, and all changes will write back to that component.
+ * <p>
+ * To use this class, register a new {@link DataComponentType} which holds an {@link Integer} for your item.
+ * Then reference that component from your {@link ICapabilityProvider} passed to {@link RegisterCapabilitiesEvent#registerItem} to create an instance of this class.
+ */
+public class ComponentEnergyStorage implements IEnergyStorage {
+    protected final MutableDataComponentHolder parent;
+    protected final DataComponentType<Integer> energyComponent;
+    protected final int capacity;
+    protected final int maxReceive;
+    protected final int maxExtract;
+
+    /**
+     * Creates a new ComponentEnergyStorage with a data component as the backing store for the energy value.
+     * 
+     * @param parent          The parent component holder, such as an {@link ItemStack}
+     * @param energyComponent The data component referencing the stored energy of the item stack
+     * @param capacity        The max capacity of the energy being stored
+     * @param maxReceive      The max per-transfer power input rate
+     * @param maxExtract      The max per-transfer power output rate
+     */
+    public ComponentEnergyStorage(MutableDataComponentHolder parent, DataComponentType<Integer> energyComponent, int capacity, int maxReceive, int maxExtract) {
+        this.parent = parent;
+        this.energyComponent = energyComponent;
+        this.capacity = capacity;
+        this.maxReceive = maxReceive;
+        this.maxExtract = maxExtract;
+    }
+
+    /**
+     * Creates a new ItemEnergyStorage with a unified receive / extract rate.
+     * 
+     * @see ComponentEnergyStorage#ItemEnergyStorage(ItemStack, DataComponentType, int, int, int)
+     */
+    public ComponentEnergyStorage(MutableDataComponentHolder parent, DataComponentType<Integer> energyComponent, int capacity, int maxTransfer) {
+        this(parent, energyComponent, capacity, maxTransfer, maxTransfer);
+    }
+
+    /**
+     * Creates a new ItemEnergyStorage with a transfer rate equivalent to the capacity.
+     * 
+     * @see ComponentEnergyStorage#ItemEnergyStorage(ItemStack, DataComponentType, int, int, int)
+     */
+    public ComponentEnergyStorage(MutableDataComponentHolder parent, DataComponentType<Integer> energyComponent, int capacity) {
+        this(parent, energyComponent, capacity, capacity);
+    }
+
+    @Override
+    public int receiveEnergy(int toReceive, boolean simulate) {
+        if (!canReceive() || toReceive <= 0) {
+            return 0;
+        }
+
+        int energy = this.getEnergyStored();
+        int energyReceived = Mth.clamp(this.capacity - energy, 0, Math.min(this.maxReceive, toReceive));
+        if (!simulate && energyReceived > 0) {
+            this.setEnergy(energy + energyReceived);
+        }
+        return energyReceived;
+    }
+
+    @Override
+    public int extractEnergy(int toExtract, boolean simulate) {
+        if (!canExtract() || toExtract <= 0) {
+            return 0;
+        }
+
+        int energy = this.getEnergyStored();
+        int energyExtracted = Math.min(energy, Math.min(this.maxExtract, toExtract));
+        if (!simulate && energyExtracted > 0) {
+            this.setEnergy(energy - energyExtracted);
+        }
+        return energyExtracted;
+    }
+
+    @Override
+    public int getEnergyStored() {
+        int rawEnergy = this.parent.getOrDefault(this.energyComponent, 0);
+        return Mth.clamp(rawEnergy, 0, this.capacity);
+    }
+
+    @Override
+    public int getMaxEnergyStored() {
+        return this.capacity;
+    }
+
+    @Override
+    public boolean canExtract() {
+        return this.maxExtract > 0;
+    }
+
+    @Override
+    public boolean canReceive() {
+        return this.maxReceive > 0;
+    }
+
+    /**
+     * Writes a new energy value to the data component. Clamps to [0, capacity]
+     * 
+     * @param energy The new energy value
+     */
+    protected void setEnergy(int energy) {
+        int realEnergy = Mth.clamp(energy, 0, this.capacity);
+        this.parent.set(this.energyComponent, realEnergy);
+    }
+}

src/main/java/net/neoforged/neoforge/energy/EnergyStorage.java

@@ -8,6 +8,7 @@
 import net.minecraft.core.HolderLookup;
 import net.minecraft.nbt.IntTag;
 import net.minecraft.nbt.Tag;
+import net.minecraft.util.Mth;
 import net.neoforged.neoforge.common.util.INBTSerializable;
 
 /**
@@ -42,35 +43,37 @@ public EnergyStorage(int capacity, int maxReceive, int maxExtract, int energy) {
     }
 
     @Override
-    public int receiveEnergy(int maxReceive, boolean simulate) {
-        if (!canReceive())
+    public int receiveEnergy(int toReceive, boolean simulate) {
+        if (!canReceive() || toReceive <= 0) {
             return 0;
+        }
 
-        int energyReceived = Math.min(capacity - energy, Math.min(this.maxReceive, maxReceive));
+        int energyReceived = Mth.clamp(this.capacity - this.energy, 0, Math.min(this.maxReceive, toReceive));
         if (!simulate)
-            energy += energyReceived;
+            this.energy += energyReceived;
         return energyReceived;
     }
 
     @Override
-    public int extractEnergy(int maxExtract, boolean simulate) {
-        if (!canExtract())
+    public int extractEnergy(int toExtract, boolean simulate) {
+        if (!canExtract() || toExtract <= 0) {
             return 0;
+        }
 
-        int energyExtracted = Math.min(energy, Math.min(this.maxExtract, maxExtract));
+        int energyExtracted = Math.min(this.energy, Math.min(this.maxExtract, toExtract));
         if (!simulate)
-            energy -= energyExtracted;
+            this.energy -= energyExtracted;
         return energyExtracted;
     }
 
     @Override
     public int getEnergyStored() {
-        return energy;
+        return this.energy;
     }
 
     @Override
     public int getMaxEnergyStored() {
-        return capacity;
+        return this.capacity;
     }
 
     @Override

src/main/java/net/neoforged/neoforge/energy/IEnergyStorage.java

@@ -16,26 +16,22 @@
  */
 public interface IEnergyStorage {
     /**
-     * Adds energy to the storage. Returns quantity of energy that was accepted.
+     * Adds energy to the storage. Returns the amount of energy that was accepted.
      *
-     * @param maxReceive
-     *                   Maximum amount of energy to be inserted.
-     * @param simulate
-     *                   If TRUE, the insertion will only be simulated.
+     * @param toReceive The amount of energy being received.
+     * @param simulate  If true, the insertion will only be simulated, meaning {@link #getEnergyStored()} will not change.
      * @return Amount of energy that was (or would have been, if simulated) accepted by the storage.
      */
-    int receiveEnergy(int maxReceive, boolean simulate);
+    int receiveEnergy(int toReceive, boolean simulate);
 
     /**
-     * Removes energy from the storage. Returns quantity of energy that was removed.
+     * Removes energy from the storage. Returns the amount of energy that was removed.
      *
-     * @param maxExtract
-     *                   Maximum amount of energy to be extracted.
-     * @param simulate
-     *                   If TRUE, the extraction will only be simulated.
+     * @param toExtract The amount of energy being extracted.
+     * @param simulate  If true, the extraction will only be simulated, meaning {@link #getEnergyStored()} will not change.
      * @return Amount of energy that was (or would have been, if simulated) extracted from the storage.
      */
-    int extractEnergy(int maxExtract, boolean simulate);
+    int extractEnergy(int toExtract, boolean simulate);
 
     /**
      * Returns the amount of energy currently stored.

tests/src/main/java/net/neoforged/neoforge/debug/capabilities/ItemEnergyTests.java

@@ -0,0 +1,78 @@
+/*
+ * Copyright (c) NeoForged and contributors
+ * SPDX-License-Identifier: LGPL-2.1-only
+ */
+
+package net.neoforged.neoforge.debug.capabilities;
+
+import com.mojang.serialization.Codec;
+import java.util.function.Supplier;
+import net.minecraft.core.component.DataComponentType;
+import net.minecraft.core.registries.Registries;
+import net.minecraft.gametest.framework.GameTest;
+import net.minecraft.network.codec.ByteBufCodecs;
+import net.minecraft.world.item.Item;
+import net.minecraft.world.item.ItemStack;
+import net.neoforged.neoforge.capabilities.Capabilities.EnergyStorage;
+import net.neoforged.neoforge.capabilities.RegisterCapabilitiesEvent;
+import net.neoforged.neoforge.energy.ComponentEnergyStorage;
+import net.neoforged.neoforge.energy.IEnergyStorage;
+import net.neoforged.neoforge.registries.DeferredItem;
+import net.neoforged.neoforge.registries.DeferredRegister;
+import net.neoforged.testframework.DynamicTest;
+import net.neoforged.testframework.TestFramework;
+import net.neoforged.testframework.annotation.ForEachTest;
+import net.neoforged.testframework.annotation.OnInit;
+import net.neoforged.testframework.annotation.TestHolder;
+import net.neoforged.testframework.gametest.EmptyTemplate;
+import net.neoforged.testframework.registration.DeferredItems;
+import net.neoforged.testframework.registration.RegistrationHelper;
+
+@ForEachTest(groups = "capabilities.itemenergy")
+public class ItemEnergyTests {
+    public static final int MAX_CAPACITY = 16384;
+
+    private static final RegistrationHelper HELPER = RegistrationHelper.create("item_energy_tests");
+
+    private static final DeferredRegister<DataComponentType<?>> COMPONENTS = HELPER.registrar(Registries.DATA_COMPONENT_TYPE);
+    private static final Supplier<DataComponentType<Integer>> ENERGY_COMPONENT = COMPONENTS.register("test_energy", () -> DataComponentType.<Integer>builder()
+            .persistent(Codec.intRange(0, MAX_CAPACITY))
+            .networkSynchronized(ByteBufCodecs.INT)
+            .build());
+
+    private static final DeferredItems ITEMS = HELPER.items();
+    private static final DeferredItem<Item> BATTERY = ITEMS.register("test_battery", () -> new Item(new Item.Properties().component(ENERGY_COMPONENT, MAX_CAPACITY)));
+
+    @OnInit
+    static void init(final TestFramework framework) {
+        COMPONENTS.register(framework.modEventBus());
+        ITEMS.register(framework.modEventBus());
+        framework.modEventBus().<RegisterCapabilitiesEvent>addListener(e -> {
+            e.registerItem(EnergyStorage.ITEM, (stack, ctx) -> {
+                return new ComponentEnergyStorage(stack, ENERGY_COMPONENT.get(), MAX_CAPACITY);
+            }, BATTERY);
+        });
+    }
+
+    @GameTest
+    @EmptyTemplate
+    @TestHolder(description = "Tests that ComponentEnergyStorage can read and write from a data component")
+    public static void testItemEnergy(DynamicTest test, RegistrationHelper reg) {
+        test.onGameTest(helper -> {
+            ItemStack stack = BATTERY.toStack();
+            IEnergyStorage energy = stack.getCapability(EnergyStorage.ITEM);
+            helper.assertValueEqual(energy.getEnergyStored(), MAX_CAPACITY, "Default stored energy should be equal to the max capacity.");
+
+            helper.assertValueEqual(energy.extractEnergy(MAX_CAPACITY, false), MAX_CAPACITY, "Extracted energy should be equal to the target value.");
+            helper.assertValueEqual(energy.getEnergyStored(), 0, "Post-extraction energy stored should be zero.");
+
+            // Sanity check the real component here
+            helper.assertValueEqual(stack.get(ENERGY_COMPONENT), 0, "Post-extraction data component value should be zero.");
+
+            helper.assertValueEqual(energy.receiveEnergy(MAX_CAPACITY, false), MAX_CAPACITY, "Received energy should be equal to the target value.");
+            helper.assertValueEqual(energy.getEnergyStored(), MAX_CAPACITY, "Post-insertion energy stored should be max capacity.");
+
+            helper.succeed();
+        });
+    }
+}