From d9baced6857d339aaa2a330cfc69981553f1b1d0 Mon Sep 17 00:00:00 2001
From: DBotThePony <dbotthepony@yandex.ru>
Date: Sat, 18 Sep 2021 13:51:04 +0700
Subject: [PATCH] Improve some grid logic

---
 .../block/entity/BlockEntityDriveRack.java    | 14 +++++++-------
 .../capability/AbstractStorageGridCell.java   | 19 ++++++++++++++++++-
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/src/main/java/ru/dbotthepony/mc/otm/block/entity/BlockEntityDriveRack.java b/src/main/java/ru/dbotthepony/mc/otm/block/entity/BlockEntityDriveRack.java
index dc7b87056..614330d54 100644
--- a/src/main/java/ru/dbotthepony/mc/otm/block/entity/BlockEntityDriveRack.java
+++ b/src/main/java/ru/dbotthepony/mc/otm/block/entity/BlockEntityDriveRack.java
@@ -14,12 +14,15 @@ import net.minecraft.world.level.Level;
 import net.minecraft.world.level.block.state.BlockState;
 import net.minecraftforge.common.capabilities.Capability;
 import net.minecraftforge.common.util.LazyOptional;
+import ru.dbotthepony.mc.otm.OverdriveThatMatters;
 import ru.dbotthepony.mc.otm.Registry;
 import ru.dbotthepony.mc.otm.capability.AbstractStorageGridCell;
 import ru.dbotthepony.mc.otm.capability.MatteryCapability;
 import ru.dbotthepony.mc.otm.capability.MatteryMachineEnergyStorage;
 import ru.dbotthepony.mc.otm.container.MatteryContainer;
 import ru.dbotthepony.mc.otm.menu.DriveRackMenu;
+import ru.dbotthepony.mc.otm.storage.ItemStackWrapper;
+import ru.dbotthepony.mc.otm.storage.VirtualComponent;
 
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -39,16 +42,13 @@ public class BlockEntityDriveRack extends BlockEntityMatteryPowered {
 		public void setChanged(int slot, ItemStack new_state, ItemStack old_state) {
 			super.setChanged(slot, new_state, old_state);
 
-			/*old_state.getCapability(MatteryCapability.DRIVE).ifPresent((drive) -> {
-				drive.removeListenerAuto(cell.computeIfAbsent(ItemStackObject.class, StorageItemView::new));
+			old_state.getCapability(MatteryCapability.DRIVE).ifPresent((drive) -> {
+				cell.computeIfAbsent(drive.storageIdentity(), VirtualComponent::new).remove(drive);
 			});
 
 			new_state.getCapability(MatteryCapability.DRIVE).ifPresent((drive) -> {
-				drive.addListenerAuto(cell.computeIfAbsent(ItemStackObject.class, StorageItemView::new));
-			});*/
-
-			old_state.getCapability(MatteryCapability.DRIVE).ifPresent(cell::removeStorageComponent);
-			new_state.getCapability(MatteryCapability.DRIVE).ifPresent(cell::addStorageComponent);
+				cell.computeIfAbsent(drive.storageIdentity(), VirtualComponent::new).add(drive);
+			});
 		}
 	};
 
diff --git a/src/main/java/ru/dbotthepony/mc/otm/capability/AbstractStorageGridCell.java b/src/main/java/ru/dbotthepony/mc/otm/capability/AbstractStorageGridCell.java
index f72deba74..a547637d7 100644
--- a/src/main/java/ru/dbotthepony/mc/otm/capability/AbstractStorageGridCell.java
+++ b/src/main/java/ru/dbotthepony/mc/otm/capability/AbstractStorageGridCell.java
@@ -12,6 +12,7 @@ import javax.annotation.ParametersAreNonnullByDefault;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.function.Function;
 import java.util.function.Supplier;
 
 @MethodsReturnNonnullByDefault
@@ -26,6 +27,7 @@ public class AbstractStorageGridCell implements IStorageGridCell {
 		return Collections.unmodifiableList(components);
 	}
 
+	@SuppressWarnings("unchecked")
 	public <T extends IStorageStack, U extends IStorageIdentity<T>> U computeIfAbsent(Class<T> identity, Supplier<U> provider) {
 		for (var component : components) {
 			if (component.storageIdentity() == identity) {
@@ -38,6 +40,19 @@ public class AbstractStorageGridCell implements IStorageGridCell {
 		return factory;
 	}
 
+	@SuppressWarnings("unchecked")
+	public <T extends IStorageStack, U extends IStorageIdentity<T>> U computeIfAbsent(Class<T> identity, Function<Class<T>, U> provider) {
+		for (var component : components) {
+			if (component.storageIdentity() == identity) {
+				return (U) component;
+			}
+		}
+
+		var factory = provider.apply(identity);
+		addStorageComponent(factory);
+		return factory;
+	}
+
 	public void addStorageComponent(IStorageIdentity<?> component) {
 		for (var component1 : components) {
 			if (component == component1 || component1.storageIdentity() == component.storageIdentity()) {
@@ -67,6 +82,7 @@ public class AbstractStorageGridCell implements IStorageGridCell {
 
 		valid = false;
 		resolver.invalidate();
+		// detachStorage();
 	}
 
 	public void revive() {
@@ -75,6 +91,7 @@ public class AbstractStorageGridCell implements IStorageGridCell {
 
 		valid = true;
 		resolver = LazyOptional.of(() -> this);
+		// attachStorage();
 	}
 
 	public LazyOptional<IStorageGridCell> get() {
@@ -90,7 +107,7 @@ public class AbstractStorageGridCell implements IStorageGridCell {
 	}
 
 	@Override
-	public void setStorageGrid(StorageGrid grid) {
+	public void setStorageGrid(@Nullable StorageGrid grid) {
 		storage_grid = grid;
 	}