From 234dc6f77f3870fb605d8047fe1d68e2e06dd224 Mon Sep 17 00:00:00 2001
From: raoulvdberge <raoulvdberge@gmail.com>
Date: Tue, 29 Oct 2019 17:54:16 +0100
Subject: [PATCH] Fix controller not persisting energy across drops

---
 .../apiimpl/network/Network.java              |  9 ++++++-
 .../loottable/ControllerLootFunction.java     | 25 +++++++++++++++++++
 .../ControllerLootFunctionSerializer.java     | 24 ++++++++++++++++++
 .../refinedstorage/setup/CommonSetup.java     |  2 ++
 .../refinedstorage/tile/ControllerTile.java   |  9 +++++++
 .../loot_tables/blocks/controller.json        | 19 +-------------
 .../blocks/creative_controller.json           | 19 ++++++++++++++
 7 files changed, 88 insertions(+), 19 deletions(-)
 create mode 100644 src/main/java/com/raoulvdberge/refinedstorage/loottable/ControllerLootFunction.java
 create mode 100644 src/main/java/com/raoulvdberge/refinedstorage/loottable/ControllerLootFunctionSerializer.java
 create mode 100644 src/main/resources/data/refinedstorage/loot_tables/blocks/creative_controller.json

diff --git a/src/main/java/com/raoulvdberge/refinedstorage/apiimpl/network/Network.java b/src/main/java/com/raoulvdberge/refinedstorage/apiimpl/network/Network.java
index 3051470f2..a11be80e4 100644
--- a/src/main/java/com/raoulvdberge/refinedstorage/apiimpl/network/Network.java
+++ b/src/main/java/com/raoulvdberge/refinedstorage/apiimpl/network/Network.java
@@ -35,6 +35,7 @@ import com.raoulvdberge.refinedstorage.tile.config.RedstoneMode;
 import com.raoulvdberge.refinedstorage.util.StackUtils;
 import net.minecraft.item.ItemStack;
 import net.minecraft.nbt.CompoundNBT;
+import net.minecraft.tileentity.TileEntity;
 import net.minecraft.util.math.BlockPos;
 import net.minecraft.world.World;
 import net.minecraft.world.server.ServerWorld;
@@ -82,7 +83,13 @@ public class Network implements INetwork, IRedstoneConfigurable {
         this.world = world;
         this.type = type;
         this.root = new RootNetworkNode(this, world, pos);
-        this.nodeGraph.addListener(() -> ((ControllerTile) world.getTileEntity(pos)).getDataManager().sendParameterToWatchers(ControllerTile.NODES));
+        this.nodeGraph.addListener(() -> {
+            TileEntity tile = world.getTileEntity(pos);
+
+            if (tile instanceof ControllerTile) {
+                ((ControllerTile) tile).getDataManager().sendParameterToWatchers(ControllerTile.NODES);
+            }
+        });
     }
 
     public RootNetworkNode getRoot() {
diff --git a/src/main/java/com/raoulvdberge/refinedstorage/loottable/ControllerLootFunction.java b/src/main/java/com/raoulvdberge/refinedstorage/loottable/ControllerLootFunction.java
new file mode 100644
index 000000000..da0ff5d67
--- /dev/null
+++ b/src/main/java/com/raoulvdberge/refinedstorage/loottable/ControllerLootFunction.java
@@ -0,0 +1,25 @@
+package com.raoulvdberge.refinedstorage.loottable;
+
+import com.raoulvdberge.refinedstorage.api.network.INetwork;
+import com.raoulvdberge.refinedstorage.tile.ControllerTile;
+import net.minecraft.item.ItemStack;
+import net.minecraft.tileentity.TileEntity;
+import net.minecraft.world.storage.loot.LootContext;
+import net.minecraft.world.storage.loot.LootParameters;
+import net.minecraft.world.storage.loot.functions.ILootFunction;
+import net.minecraftforge.energy.CapabilityEnergy;
+
+public class ControllerLootFunction implements ILootFunction {
+    @Override
+    public ItemStack apply(ItemStack itemStack, LootContext lootContext) {
+        TileEntity tile = lootContext.get(LootParameters.BLOCK_ENTITY);
+
+        if (tile instanceof ControllerTile) {
+            INetwork network = ((ControllerTile) tile).getRemovedNetwork() == null ? ((ControllerTile) tile).getNetwork() : ((ControllerTile) tile).getRemovedNetwork();
+
+            itemStack.getCapability(CapabilityEnergy.ENERGY).ifPresent(energy -> energy.receiveEnergy(network.getEnergyStorage().getEnergyStored(), false));
+        }
+
+        return itemStack;
+    }
+}
diff --git a/src/main/java/com/raoulvdberge/refinedstorage/loottable/ControllerLootFunctionSerializer.java b/src/main/java/com/raoulvdberge/refinedstorage/loottable/ControllerLootFunctionSerializer.java
new file mode 100644
index 000000000..9882e50ce
--- /dev/null
+++ b/src/main/java/com/raoulvdberge/refinedstorage/loottable/ControllerLootFunctionSerializer.java
@@ -0,0 +1,24 @@
+package com.raoulvdberge.refinedstorage.loottable;
+
+import com.google.gson.JsonDeserializationContext;
+import com.google.gson.JsonObject;
+import com.google.gson.JsonSerializationContext;
+import com.raoulvdberge.refinedstorage.RS;
+import net.minecraft.util.ResourceLocation;
+import net.minecraft.world.storage.loot.functions.ILootFunction;
+
+public class ControllerLootFunctionSerializer extends ILootFunction.Serializer<ControllerLootFunction> {
+    public ControllerLootFunctionSerializer() {
+        super(new ResourceLocation(RS.ID, "controller"), ControllerLootFunction.class);
+    }
+
+    @Override
+    public void serialize(JsonObject jsonObject, ControllerLootFunction controllerLootFunction, JsonSerializationContext jsonSerializationContext) {
+
+    }
+
+    @Override
+    public ControllerLootFunction deserialize(JsonObject jsonObject, JsonDeserializationContext jsonDeserializationContext) {
+        return new ControllerLootFunction();
+    }
+}
diff --git a/src/main/java/com/raoulvdberge/refinedstorage/setup/CommonSetup.java b/src/main/java/com/raoulvdberge/refinedstorage/setup/CommonSetup.java
index 9d7d8f9cd..ea800645c 100644
--- a/src/main/java/com/raoulvdberge/refinedstorage/setup/CommonSetup.java
+++ b/src/main/java/com/raoulvdberge/refinedstorage/setup/CommonSetup.java
@@ -34,6 +34,7 @@ import com.raoulvdberge.refinedstorage.container.*;
 import com.raoulvdberge.refinedstorage.container.factory.*;
 import com.raoulvdberge.refinedstorage.item.*;
 import com.raoulvdberge.refinedstorage.item.blockitem.*;
+import com.raoulvdberge.refinedstorage.loottable.ControllerLootFunctionSerializer;
 import com.raoulvdberge.refinedstorage.loottable.CrafterLootFunctionSerializer;
 import com.raoulvdberge.refinedstorage.loottable.PortableGridBlockLootFunctionSerializer;
 import com.raoulvdberge.refinedstorage.loottable.StorageBlockLootFunctionSerializer;
@@ -132,6 +133,7 @@ public class CommonSetup {
         LootFunctionManager.registerFunction(new StorageBlockLootFunctionSerializer());
         LootFunctionManager.registerFunction(new PortableGridBlockLootFunctionSerializer());
         LootFunctionManager.registerFunction(new CrafterLootFunctionSerializer());
+        LootFunctionManager.registerFunction(new ControllerLootFunctionSerializer());
     }
 
     private INetworkNode readAndReturn(CompoundNBT tag, NetworkNode node) {
diff --git a/src/main/java/com/raoulvdberge/refinedstorage/tile/ControllerTile.java b/src/main/java/com/raoulvdberge/refinedstorage/tile/ControllerTile.java
index e31e87081..3bb9cf060 100644
--- a/src/main/java/com/raoulvdberge/refinedstorage/tile/ControllerTile.java
+++ b/src/main/java/com/raoulvdberge/refinedstorage/tile/ControllerTile.java
@@ -44,6 +44,9 @@ public class ControllerTile extends BaseTile implements INetworkNodeProxy<RootNe
     private final LazyOptional<INetworkNodeProxy<RootNetworkNode>> networkNodeProxyCap = LazyOptional.of(() -> this);
 
     private final NetworkType type;
+
+    private INetwork removedNetwork;
+
     private Network dummyNetwork;
 
     public ControllerTile(NetworkType type) {
@@ -117,6 +120,8 @@ public class ControllerTile extends BaseTile implements INetworkNodeProxy<RootNe
 
             INetwork network = manager.getNetwork(pos);
 
+            removedNetwork = network;
+
             manager.removeNetwork(pos);
             manager.markForSaving();
 
@@ -124,6 +129,10 @@ public class ControllerTile extends BaseTile implements INetworkNodeProxy<RootNe
         }
     }
 
+    public INetwork getRemovedNetwork() {
+        return removedNetwork;
+    }
+
     @Override
     @Nonnull
     public RootNetworkNode getNode() {
diff --git a/src/main/resources/data/refinedstorage/loot_tables/blocks/controller.json b/src/main/resources/data/refinedstorage/loot_tables/blocks/controller.json
index d719001b0..d44206853 100644
--- a/src/main/resources/data/refinedstorage/loot_tables/blocks/controller.json
+++ b/src/main/resources/data/refinedstorage/loot_tables/blocks/controller.json
@@ -9,24 +9,7 @@
           "name": "refinedstorage:controller",
           "functions": [
             {
-              "function": "minecraft:copy_nbt",
-              "source": "block_entity",
-              "ops": [
-                {
-                  "source": "Energy",
-                  "target": "Energy",
-                  "op": "replace"
-                }
-              ]
-            },
-            {
-              "function": "minecraft:set_contents",
-              "entries": [
-                {
-                  "type": "minecraft:dynamic",
-                  "name": "minecraft:contents"
-                }
-              ]
+              "function": "refinedstorage:controller"
             }
           ]
         }
diff --git a/src/main/resources/data/refinedstorage/loot_tables/blocks/creative_controller.json b/src/main/resources/data/refinedstorage/loot_tables/blocks/creative_controller.json
new file mode 100644
index 000000000..3f9584d6c
--- /dev/null
+++ b/src/main/resources/data/refinedstorage/loot_tables/blocks/creative_controller.json
@@ -0,0 +1,19 @@
+{
+  "type": "minecraft:block",
+  "pools": [
+    {
+      "rolls": 1,
+      "entries": [
+        {
+          "type": "minecraft:item",
+          "name": "refinedstorage:creative_controller"
+        }
+      ],
+      "conditions": [
+        {
+          "condition": "minecraft:survives_explosion"
+        }
+      ]
+    }
+  ]
+}
\ No newline at end of file