Generify mattery drive

src/main/java/ru/dbotthepony/mc/otm/capability/drive/DrivePool.java

@@ -161,6 +161,7 @@ public class DrivePool {
 		base_directory = new File(server.storageSource.getWorldDir().toFile(), DATA_PATH);
 		base_directory.mkdirs();
 
+		stopping = false;
 		thread = new Thread(null, () -> DrivePool.thread(server), "Overdrive That Matters DrivePool IO");
 		thread.start();
 	}

src/main/java/ru/dbotthepony/mc/otm/capability/drive/IItemMatteryDrive.java

@@ -0,0 +1,22 @@
+package ru.dbotthepony.mc.otm.capability.drive;
+
+import net.minecraft.world.item.Item;
+import net.minecraft.world.item.ItemStack;
+import ru.dbotthepony.mc.otm.storage.IStorageTuple;
+import ru.dbotthepony.mc.otm.storage.ItemStackWrapper;
+
+import java.util.List;
+
+public interface IItemMatteryDrive extends IMatteryDrive<ItemStackWrapper> {
+	/**
+	 * @param item
+	 * @return all items belonging to passed class
+	 */
+	List<IStorageTuple<ItemStackWrapper>> findItems(Item item);
+
+	/**
+	 * @param stack
+	 * @return all items that match this itemstack
+	 */
+	List<IStorageTuple<ItemStackWrapper>> findItems(ItemStack stack);
+}

src/main/java/ru/dbotthepony/mc/otm/capability/drive/IMatteryDrive.java

@@ -7,44 +7,19 @@ import net.minecraft.world.item.ItemStack;
 import ru.dbotthepony.mc.otm.storage.*;
 
 import javax.annotation.ParametersAreNonnullByDefault;
+import java.math.BigDecimal;
 import java.util.List;
 import java.util.UUID;
 
 @ParametersAreNonnullByDefault
 @MethodsReturnNonnullByDefault
-public interface IMatteryDrive extends IStorageComponent<ItemStackWrapper> {
+public interface IMatteryDrive<T extends IStorageStack> extends IStorageComponent<T> {
-	record StoredStack(ItemStack stack, UUID id) {};
-
-	List<StoredStack> getItems();
-
 	boolean isDirty();
 	void markDirty();
 	void markClean();
 
-	/**
+	BigDecimal getStoredCount();
-	 * @param item
+	BigDecimal getCapacity();
-	 * @return all items belonging to passed class
-	 */
-	List<StoredStack> findItems(Item item);
-
-	/**
-	 * @param stack
-	 * @return all items that match this itemstack
-	 */
-	List<StoredStack> findItems(ItemStack stack);
-
-	ItemStack getItem(UUID id);
-
-	int getStoredCount();
-	int getCapacity();
-
-	ItemStack insertItem(ItemStack item, boolean simulate);
-
-	default ItemStack extractItem(UUID id, int amount, boolean simulate) {
-		return extractItem(id, amount, true, simulate);
-	}
-
-	ItemStack extractItem(UUID id, int amount, boolean obey_stack_size, boolean simulate);
 
 	// not extending INBTSerializable to avoid serializing it as forgecaps
 	CompoundTag serializeNBT();

src/main/java/ru/dbotthepony/mc/otm/capability/drive/ItemMatteryDrive.java

@@ -0,0 +1,125 @@
+package ru.dbotthepony.mc.otm.capability.drive;
+
+import net.minecraft.MethodsReturnNonnullByDefault;
+import net.minecraft.nbt.CompoundTag;
+import net.minecraft.nbt.ListTag;
+import net.minecraft.resources.ResourceLocation;
+import net.minecraft.world.item.Item;
+import net.minecraft.world.item.ItemStack;
+import net.minecraft.world.item.Items;
+import net.minecraftforge.registries.RegistryManager;
+import ru.dbotthepony.mc.otm.storage.IStorageTuple;
+import ru.dbotthepony.mc.otm.storage.ItemStackWrapper;
+import ru.dbotthepony.mc.otm.storage.StorageObjectRegistry;
+import ru.dbotthepony.mc.otm.storage.StorageObjectTuple;
+
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+import javax.annotation.ParametersAreNonnullByDefault;
+import java.math.BigDecimal;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+
+@MethodsReturnNonnullByDefault
+@ParametersAreNonnullByDefault
+public class ItemMatteryDrive extends MatteryDrive<ItemStackWrapper> implements IItemMatteryDrive {
+	private static StorageObjectTuple<ItemStackWrapper> identity;
+
+	public ItemMatteryDrive(BigDecimal capacity, int max_different_stacks) {
+		super(capacity, max_different_stacks);
+	}
+
+	public ItemMatteryDrive(BigDecimal capacity) {
+		super(capacity);
+	}
+
+	@Override
+	public StorageObjectTuple<ItemStackWrapper> identity() {
+		if (identity == null)
+			identity = StorageObjectRegistry.getOrError(ItemStackWrapper.class);
+
+		return identity;
+	}
+
+	public ItemStack insertObject(ItemStack item, boolean simulate) {
+		return insertObject(new ItemStackWrapper(item), simulate).stack();
+	}
+
+	@Override
+	protected CompoundTag serializeStack(IStorageTuple<ItemStackWrapper> item) {
+		final var tag = new CompoundTag();
+		final var location = Objects.requireNonNull(item.object().stack().getItem().getRegistryName(), "Missing registry name for stored Item").toString();
+
+		tag.putString("item", location);
+		tag.putInt("count", item.object().stack().getCount());
+
+		CompoundTag item_tag;
+
+		if ((item_tag = item.object().stack().getTag()) != null) {
+			tag.put("data", item_tag);
+		}
+
+		return tag;
+	}
+
+	@Nullable
+	@Override
+	protected ItemStackWrapper deserializeStack(CompoundTag tag) {
+		final var item = RegistryManager.ACTIVE.getRegistry(Item.class).getValue(new ResourceLocation(tag.getString("item")));
+
+		if (item != null && item != Items.AIR) {
+			final var count = tag.getInt("count");
+			final var data = tag.get("data");
+
+			final var itemstack = new ItemStack(item, count);
+
+			if (data != null) {
+				itemstack.setTag((CompoundTag) data);
+			}
+
+			return new ItemStackWrapper(itemstack);
+		}
+
+		return null;
+	}
+
+
+	@Override
+	public List<IStorageTuple<ItemStackWrapper>> findItems(Item item) {
+		var list = items.get(item);
+
+		if (list != null)
+			return List.copyOf(list);
+
+		return List.of();
+	}
+
+	@Override
+	public List<IStorageTuple<ItemStackWrapper>> findItems(ItemStack stack) {
+		var list = items.get(stack.getItem());
+
+		if (list == null)
+			return List.of();
+
+		var amount = 0;
+
+		for (var _stack : list) {
+			if (ItemStack.tagMatches(_stack.object().stack(), stack)) {
+				amount++;
+			}
+		}
+
+		var build_list = new ArrayList<IStorageTuple<ItemStackWrapper>>(amount);
+		var i = 0;
+
+		for (var _stack : list) {
+			if (ItemStack.tagMatches(_stack.object().stack(), stack)) {
+				build_list.set(i, _stack);
+				i++;
+			}
+		}
+
+		return build_list;
+	}
+}

src/main/java/ru/dbotthepony/mc/otm/capability/drive/MatteryDrive.java

@@ -4,59 +4,37 @@ import net.minecraft.MethodsReturnNonnullByDefault;
 import net.minecraft.nbt.CompoundTag;
 import net.minecraft.nbt.ListTag;
 import net.minecraft.nbt.Tag;
-import net.minecraft.resources.ResourceLocation;
+import ru.dbotthepony.mc.otm.storage.*;
-import net.minecraft.world.item.Item;
-import net.minecraft.world.item.ItemStack;
-import net.minecraft.world.item.Items;
-import net.minecraftforge.registries.RegistryManager;
-import ru.dbotthepony.mc.otm.storage.IStorageListener;
-import ru.dbotthepony.mc.otm.storage.IStorageTuple;
-import ru.dbotthepony.mc.otm.storage.ItemStackWrapper;
-import ru.dbotthepony.mc.otm.storage.StorageTuple;
 
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import javax.annotation.ParametersAreNonnullByDefault;
 import java.math.BigDecimal;
 import java.util.*;
 
 @MethodsReturnNonnullByDefault
 @ParametersAreNonnullByDefault
-public class MatteryDrive implements IMatteryDrive {
+abstract public class MatteryDrive<T extends IStorageStack> implements IMatteryDrive<T> {
-	protected final HashMap<Item, List<StoredStack>> items = new HashMap<>();
+	protected final HashMap<Object, List<IStorageTuple<T>>> items = new HashMap<>();
-	protected final HashMap<UUID, StoredStack> items_by_id = new HashMap<>();
+	protected final HashMap<UUID, IStorageTuple<T>> items_by_id = new HashMap<>();
+
+	abstract public StorageObjectTuple<T> identity();
 
 	protected boolean dirty = false;
 	protected int different_stacks = 0;
-	protected int stored = 0;
+	protected BigDecimal stored = BigDecimal.ZERO;
 	protected int max_different_stacks;
-	protected int capacity;
+	protected BigDecimal capacity;
 
-	public MatteryDrive(int capacity, int max_different_stacks) {
+	public MatteryDrive(BigDecimal capacity, int max_different_stacks) {
 		this.capacity = capacity;
 		this.max_different_stacks = max_different_stacks;
 	}
 
-	public MatteryDrive(int capacity) {
+	public MatteryDrive(BigDecimal capacity) {
 		this(capacity, 0xFFFF);
 	}
 
-	@Override
-	public List<StoredStack> getItems() {
-		int amount = 0;
-
-		for (var list : items.values()) {
-			amount += list.size();
-		}
-
-		var output = new ArrayList<StoredStack>(amount);
-
-		for (var list : items.values()) {
-			output.addAll(list);
-		}
-
-		return output;
-	}
-
 	@Override
 	public boolean isDirty() {
 		return dirty;
@@ -73,86 +51,40 @@ public class MatteryDrive implements IMatteryDrive {
 	}
 
 	@Override
-	public int getStoredCount() {
+	public BigDecimal getStoredCount() {
 		return stored;
 	}
 
 	@Override
-	public int getCapacity() {
+	public BigDecimal getCapacity() {
 		return capacity;
 	}
 
-	@Override
-	@Nonnull
-	public List<StoredStack> findItems(@Nonnull Item item) {
-		var list = items.get(item);
-
-		if (list != null)
-			return List.copyOf(list);
-
-		return List.of();
-	}
-
-	@Override
-	@Nonnull
-	public List<StoredStack> findItems(@Nonnull ItemStack stack) {
-		var list = items.get(stack.getItem());
-
-		if (list == null)
-			return List.of();
-
-		var amount = 0;
-
-		for (var _stack : list) {
-			if (ItemStack.tagMatches(_stack.stack(), stack)) {
-				amount++;
-			}
-		}
-
-		var build_list = new ArrayList<StoredStack>(amount);
-		var i = 0;
-
-		for (var _stack : list) {
-			if (ItemStack.tagMatches(_stack.stack(), stack)) {
-				build_list.set(i, _stack);
-				i++;
-			}
-		}
-
-		return build_list;
-	}
-
 	@Nonnull
 	@Override
-	public ItemStack getItem(UUID id) {
+	@SuppressWarnings("unchecked")
-		var stack = items_by_id.get(id);
+	public T insertObject(T item, boolean simulate) {
-		return stack != null ? stack.stack() : ItemStack.EMPTY;
+		BigDecimal max_insert = getCapacity().subtract(getStoredCount()).min(item.getCount());
-	}
 
-	@Nonnull
+		if (max_insert.compareTo(BigDecimal.ZERO) <= 0)
-	@Override
-	public ItemStack insertItem(@Nonnull ItemStack item, boolean simulate) {
-		int max_insert = Math.min(getCapacity() - getStoredCount(), item.getCount());
-
-		if (max_insert <= 0)
 			return item;
 
-		var listing = items.computeIfAbsent(item.getItem(), (key) -> new ArrayList<>());
+		final var listing = items.computeIfAbsent(item.partitionKey(), (key) -> new ArrayList<>());
 
 		for (var state : listing) {
-			if (ItemStack.tagMatches(state.stack(), item)) {
+			if (state.object().sameItem(item)) {
 				if (!simulate) {
-					state.stack().grow(max_insert);
+					state.object().grow(max_insert);
-					stored += max_insert;
+					stored = stored.add(max_insert);
 
 					for (var listener : listeners2) {
-						listener.changeObject(state.id(), new BigDecimal(state.stack().getCount()));
+						listener.changeObject(state.id(), state.object().getCount());
 					}
 
 					markDirty();
 				}
 
-				var copy_item = item.copy();
+				final var copy_item = (T) item.copy();
 				copy_item.shrink(max_insert);
 				return copy_item;
 			}
@@ -164,54 +96,49 @@ public class MatteryDrive implements IMatteryDrive {
 
 		if (!simulate) {
 			different_stacks++;
-			stored += max_insert;
+			stored = stored.add(max_insert);
 
-			var copy = item.copy();
+			final var copy = (T) item.copy();
 			copy.setCount(max_insert);
-			var state = new StoredStack(copy, UUID.randomUUID());
+			final var state = new StorageTuple<>(UUID.randomUUID(), copy);
 			listing.add(state);
 			items_by_id.put(state.id(), state);
 
 			for (var listener : listeners2) {
-				listener.addObject(new ItemStackWrapper(state.stack()), state.id(), this);
+				listener.addObject(state.object(), state.id(), this);
 			}
 
 			markDirty();
 		}
 
-		var copy_item = item.copy();
+		final var copy_item = (T) item.copy();
 		copy_item.shrink(max_insert);
 		return copy_item;
 	}
 
-	@Override
-	public ItemStackWrapper insertObject(ItemStackWrapper obj, boolean simulate) {
-		var get = insertItem(obj.stack(), simulate);
-		return get.isEmpty() ? ItemStackWrapper.EMPTY : new ItemStackWrapper(get);
-	}
-
 	@Nonnull
 	@Override
-	public ItemStack extractItem(UUID id, int amount, boolean obey_stack_size, boolean simulate) {
+	@SuppressWarnings("unchecked")
+	public T extractObject(UUID id, BigDecimal amount, boolean simulate) {
 		var get = items_by_id.get(id);
 
 		if (get == null)
-			return ItemStack.EMPTY;
+			return identity().empty();
 
-		var extract = Math.min(amount, get.stack().getCount());
+		if (amount.compareTo(BigDecimal.ZERO) <= 0)
+			amount = get.object().getMaxStackSize().orElse(get.object().getCount());
 
-		if (obey_stack_size)
+		amount = amount.min(get.object().getCount());
-			extract = Math.min(extract, get.stack().getItem().getItemStackLimit(get.stack()));
 
-		if (extract <= 0)
+		if (amount.compareTo(BigDecimal.ZERO) <= 0)
-			return ItemStack.EMPTY;
+			return identity().empty();
 
-		var copy = get.stack().copy();
+		final var copy = (T) get.object().copy();
-		copy.setCount(extract);
+		copy.setCount(amount);
 
 		if (!simulate) {
-			if (extract == get.stack().getCount()) {
+			if (amount.compareTo(get.object().getCount()) == 0) {
-				var listing = items.get(get.stack().getItem());
+				var listing = items.get(get.object().partitionKey());
 				listing.remove(get);
 				different_stacks--;
 
@@ -220,16 +147,16 @@ public class MatteryDrive implements IMatteryDrive {
 				}
 
 				if (listing.size() == 0) {
-					items.remove(get.stack().getItem());
+					items.remove(get.object().partitionKey());
 				}
 			}
 
-			stored -= extract;
+			stored = stored.subtract(amount);
-			get.stack().shrink(extract);
+			get.object().shrink(amount);
 
-			if (get.stack().getCount() != 0) {
+			if (get.object().getCount().compareTo(BigDecimal.ZERO) != 0) {
 				for (var listener : listeners2) {
-					listener.changeObject(get.id(), new BigDecimal(get.stack().getCount()));
+					listener.changeObject(get.id(), get.object().getCount());
 				}
 			}
 
@@ -239,132 +166,99 @@ public class MatteryDrive implements IMatteryDrive {
 		return copy;
 	}
 
+	@Nullable
+	abstract protected CompoundTag serializeStack(IStorageTuple<T> item);
+
 	@Override
 	public CompoundTag serializeNBT() {
 		final var compound = new CompoundTag();
 
-		compound.putInt("capacity", capacity);
+		compound.putString("capacity", capacity.toString());
 		compound.putInt("max_different_stacks", max_different_stacks);
 
 		var list_items = new ListTag();
 		compound.put("items", list_items);
 
 		for (var entry : items.entrySet()) {
-			var item = Objects.requireNonNull(entry.getKey().getRegistryName(), "Missing registry name for stored Item").toString();
-			var entry_compound = new CompoundTag();
-
-			entry_compound.putString("id", item);
-			var stack_list = new ListTag();
-			entry_compound.put("list", stack_list);
-
 			for (var stack : entry.getValue()) {
-				var stack_nbt = new CompoundTag();
+				CompoundTag entry_compound;
-				stack_list.add(stack_nbt);
 
-				stack_nbt.putInt("count", stack.stack().getCount());
+				if ((entry_compound = serializeStack(stack)) != null) {
-				stack_nbt.putLongArray("id", new long[] { stack.id().getMostSignificantBits(), stack.id().getLeastSignificantBits() });
+					entry_compound.putLongArray("id", new long[] { stack.id().getMostSignificantBits(), stack.id().getLeastSignificantBits() });
-
+					list_items.add(entry_compound);
-				if (stack.stack().getTag() != null) {
-					stack_nbt.put("data", stack.stack().getTag());
 				}
 			}
-
-			list_items.add(entry_compound);
 		}
 
 		return compound;
 	}
 
+	@Nullable
+	abstract protected T deserializeStack(CompoundTag tag);
+
 	@Override
 	public void deserializeNBT(CompoundTag nbt) {
 		items.clear();
 		items_by_id.clear();
-		stored = 0;
+		stored = BigDecimal.ZERO;
 		different_stacks = 0;
 
-		capacity = nbt.getInt("capacity");
+		capacity = new BigDecimal(nbt.getString("capacity"));
 		max_different_stacks = nbt.getInt("max_different_stacks");
 
-		var list_items = nbt.getList("items", Tag.TAG_COMPOUND);
+		for (var _entry : nbt.getList("items", Tag.TAG_COMPOUND)) {
-		var registry = RegistryManager.ACTIVE.getRegistry(Item.class);
-
-		for (var _entry : list_items) {
 			if (_entry instanceof CompoundTag entry) {
-				var item = registry.getValue(new ResourceLocation(entry.getString("id")));
+				final var stack = deserializeStack(entry);
 
-				if (item != null && item != Items.AIR) {
+				if (stack != null) {
-					var map_list = items.computeIfAbsent(item, (key) -> new ArrayList<>());
+					stored = stored.add(stack.getCount());
+					different_stacks++;
+					final var id = entry.getLongArray("id");
 
-					var stack_list = entry.getList("list", Tag.TAG_COMPOUND);
+					final var tuple = new StorageTuple<>(id.length == 2 ? new UUID(id[0], id[1]) : UUID.randomUUID(), stack);
-
+					items.computeIfAbsent(stack.partitionKey(), (key) -> new ArrayList<>()).add(tuple);
-					for (var _stack : stack_list) {
+					items_by_id.put(tuple.id(), tuple);
-						if (_stack instanceof CompoundTag stack) {
-							var count = stack.getInt("count");
-							var id = stack.getLongArray("id");
-							var data = stack.get("data");
-
-							var itemstack = new ItemStack(item, count);
-
-							if (data != null) {
-								itemstack.setTag((CompoundTag) data);
-							}
-
-							stored += count;
-							different_stacks += 1;
-							var state = new StoredStack(itemstack, id.length == 2 ? new UUID(id[0], id[1]) : UUID.randomUUID());
-							map_list.add(state);
-							items_by_id.put(state.id(), state);
-						}
-					}
 				}
 			}
 		}
 	}
 
 	@Override
-	public Class<ItemStackWrapper> storageIdentity() {
+	public Class<T> storageIdentity() {
-		return ItemStackWrapper.class;
+		return identity().identity();
 	}
 
 	@Override
-	public ItemStackWrapper getStoredObject(UUID id) {
+	public T getStoredObject(UUID id) {
-		var get = getItem(id);
+		var get = items_by_id.get(id);
-		return get.isEmpty() ? ItemStackWrapper.EMPTY : new ItemStackWrapper(get);
+		return get == null ? identity().empty() : get.object();
 	}
 
 	@Override
-	public ItemStackWrapper extractObject(UUID id, BigDecimal amount, boolean simulate) {
+	public List<IStorageTuple<T>> getStorageObjects() {
-		var get = extractItem(id, amount.intValue(), simulate);
-		return get.isEmpty() ? ItemStackWrapper.EMPTY : new ItemStackWrapper(get);
-	}
-
-	@Override
-	public List<IStorageTuple<ItemStackWrapper>> getStorageObjects() {
 		int amount = 0;
 
 		for (var listing : items.values())
 			amount += listing.size();
 
-		var list = new ArrayList<IStorageTuple<ItemStackWrapper>>(amount);
+		var list = new ArrayList<IStorageTuple<T>>(amount);
 
 		for (var listing : items.values()) {
-			for (var stack : listing) {
+			list.addAll(listing);
-				list.add(new StorageTuple<>(stack.id(), new ItemStackWrapper(stack.stack())));
-			}
 		}
 
 		return list;
 	}
 
-	protected final HashSet<IStorageListener<ItemStackWrapper>> listeners2 = new HashSet<>();
+	protected final HashSet<IStorageListener<T>> listeners2 = new HashSet<>();
 
 	@Override
-	public boolean addListener(IStorageListener<ItemStackWrapper> listener) {
+	public boolean addListener(IStorageListener<T> listener) {
 		return listeners2.add(listener);
 	}
 
 	@Override
-	public boolean removeListener(IStorageListener<ItemStackWrapper> listener) {
+	public boolean removeListener(IStorageListener<T> listener) {
 		return listeners2.remove(listener);
 	}
 }

src/main/java/ru/dbotthepony/mc/otm/item/ItemPortableCondensationDrive.java

@@ -35,11 +35,14 @@ import ru.dbotthepony.mc.otm.OverdriveThatMatters;
 import ru.dbotthepony.mc.otm.capability.MatteryCapability;
 import ru.dbotthepony.mc.otm.capability.drive.DrivePool;
 import ru.dbotthepony.mc.otm.capability.drive.IMatteryDrive;
+import ru.dbotthepony.mc.otm.capability.drive.ItemMatteryDrive;
 import ru.dbotthepony.mc.otm.capability.drive.MatteryDrive;
 import ru.dbotthepony.mc.otm.menu.DriveViewerMenu;
+import ru.dbotthepony.mc.otm.storage.ItemStackWrapper;
 
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
+import java.math.BigDecimal;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -47,11 +50,11 @@ import java.util.UUID;
 import java.util.function.Supplier;
 
 public class ItemPortableCondensationDrive extends Item {
-	public final int capacity;
+	public final BigDecimal capacity;
 
 	public ItemPortableCondensationDrive(int capacity) {
 		super(new Properties().stacksTo(1).tab(OverdriveThatMatters.CREATIVE_TAB));
-		this.capacity = capacity;
+		this.capacity = new BigDecimal(capacity);
 	}
 
 	@SubscribeEvent
@@ -72,12 +75,12 @@ public class ItemPortableCondensationDrive extends Item {
 				var _cap = stack.getCapability(MatteryCapability.DRIVE).resolve();
 
 				if (_cap.isPresent()) {
-					var cap = _cap.get();
+					var cap = (ItemMatteryDrive) _cap.get();
 					var filter = drive.getFilterSettings(stack);
 
 					if (filter.matches(event.getItem().getItem())) {
 						var copy = event.getItem().getItem().copy();
-						var remaining = cap.insertItem(event.getItem().getItem(), false);
+						var remaining = cap.insertObject(event.getItem().getItem(), false);
 
 						if (remaining.getCount() == event.getItem().getItem().getCount()) {
 							continue;
@@ -117,10 +120,10 @@ public class ItemPortableCondensationDrive extends Item {
 		private final ItemStack stack;
 		private final LazyOptional<IMatteryDrive> resolver = LazyOptional.of(() -> {
 			return DrivePool.get(uuid, (tag) -> {
-				var drive = new MatteryDrive(capacity);
+				var drive = new ItemMatteryDrive(capacity);
 				drive.deserializeNBT(tag);
 				return drive;
-			}, () -> new MatteryDrive(capacity));
+			}, () -> new ItemMatteryDrive(capacity));
 		});
 
 		DriveCapability(ItemStack stack) {

src/main/java/ru/dbotthepony/mc/otm/matter/MatterRegistry.java

@@ -20,9 +20,11 @@ import net.minecraftforge.fmlserverevents.FMLServerStartedEvent;
 import net.minecraftforge.fmlserverevents.FMLServerStartingEvent;
 import ru.dbotthepony.mc.otm.OverdriveThatMatters;
 import ru.dbotthepony.mc.otm.capability.MatteryCapability;
+import ru.dbotthepony.mc.otm.capability.drive.IMatteryDrive;
 import ru.dbotthepony.mc.otm.menu.FormattingHelper;
 import ru.dbotthepony.mc.otm.network.MatterRegistryPacket;
 import ru.dbotthepony.mc.otm.network.MatteryNetworking;
+import ru.dbotthepony.mc.otm.storage.ItemStackWrapper;
 
 import javax.annotation.Nullable;
 import java.math.BigDecimal;
@@ -50,7 +52,7 @@ public class MatterRegistry {
 	}
 
 	public static boolean canDecompose(ItemStack stack) {
-		if (stack.getCapability(MatteryCapability.DRIVE).isPresent() && stack.getCapability(MatteryCapability.DRIVE).resolve().get().getStoredCount() > 0)
+		if (stack.getCapability(MatteryCapability.DRIVE).isPresent() && stack.getCapability(MatteryCapability.DRIVE).resolve().get().getStoredCount().compareTo(BigDecimal.ZERO) > 0)
 			return false;
 
 		if (stack.getCapability(MatteryCapability.MATTER).isPresent() && stack.getCapability(MatteryCapability.MATTER).resolve().get().getStoredMatter().compareTo(BigDecimal.ZERO) > 0)
@@ -74,9 +76,9 @@ public class MatterRegistry {
 		var cap1 = stack.getCapability(MatteryCapability.DRIVE).resolve();
 		var cap2 = stack.getCapability(MatteryCapability.MATTER).resolve();
 
-		if (cap1.isPresent()) {
+		if (cap1.isPresent() && cap1.get().storageIdentity() == ItemStackWrapper.class) {
-			for (var stored : cap1.get().getItems()) {
+			for (var stored : ((IMatteryDrive<ItemStackWrapper>) cap1.get()).getStorageObjects()) {
-				matter = matter.add(getMatterValue(stored.stack(), level + 1));
+				matter = matter.add(getMatterValue(stored.object().stack(), level + 1));
 			}
 		}
 

src/main/java/ru/dbotthepony/mc/otm/menu/DriveViewerMenu.java

@@ -22,6 +22,7 @@ import ru.dbotthepony.mc.otm.menu.data.NetworkedItemView;
 import ru.dbotthepony.mc.otm.menu.slot.MatterySlot;
 import ru.dbotthepony.mc.otm.network.MatteryNetworking;
 import ru.dbotthepony.mc.otm.network.SetCarriedPacket;
+import ru.dbotthepony.mc.otm.storage.ItemStackWrapper;
 
 import javax.annotation.Nullable;
 import java.util.function.Supplier;
@@ -29,7 +30,7 @@ import java.util.function.Supplier;
 public class DriveViewerMenu extends PoweredMatteryMenu implements INetworkedItemViewSupplier {
 	public final NetworkedItemView view;
 	public MatterySlot drive_slot;
-	public IMatteryDrive last_drive;
+	public IMatteryDrive<ItemStackWrapper> last_drive;
 
 	@Override
 	public NetworkedItemView getNetworkedItemView() {
@@ -78,6 +79,7 @@ public class DriveViewerMenu extends PoweredMatteryMenu implements INetworkedIte
 	}
 
 	@Override
+	@SuppressWarnings("unchecked")
 	public void broadcastChanges() {
 		super.broadcastChanges();
 
@@ -89,10 +91,10 @@ public class DriveViewerMenu extends PoweredMatteryMenu implements INetworkedIte
 			} else {
 				var get_cap = get.getCapability(MatteryCapability.DRIVE).resolve();
 
-				if (get_cap.isEmpty()) {
+				if (get_cap.isEmpty() || get_cap.get().storageIdentity() != ItemStackWrapper.class) {
 					last_drive = null;
 				} else {
-					last_drive = get_cap.get();
+					last_drive = (IMatteryDrive<ItemStackWrapper>) get_cap.get();
 				}
 			}
 
@@ -139,9 +141,9 @@ public class DriveViewerMenu extends PoweredMatteryMenu implements INetworkedIte
 			return ItemStack.EMPTY;
 
 		if (amount == item.getCount()) {
-			var remaining = last_drive.insertItem(item, false);
+			var remaining = last_drive.insertObject(new ItemStackWrapper(item), false);
 
-			if (remaining.getCount() == item.getCount())
+			if (remaining.getCount().intValue() == item.getCount())
 				return ItemStack.EMPTY;
 
 			if (remaining.isEmpty()) {
@@ -152,8 +154,8 @@ public class DriveViewerMenu extends PoweredMatteryMenu implements INetworkedIte
 			}
 
 			var copy = item.copy();
-			tile.getIOItemCount(item.getCount() - remaining.getCount(), false);
+			tile.getIOItemCount(item.getCount() - remaining.getCount().intValue(), false);
-			item.shrink(remaining.getCount());
+			item.shrink(remaining.getCount().intValue());
 			slot.setChanged();
 
 			return copy;
@@ -162,14 +164,14 @@ public class DriveViewerMenu extends PoweredMatteryMenu implements INetworkedIte
 		var copy_insert = item.copy();
 		copy_insert.setCount(amount);
 
-		var remaining = last_drive.insertItem(copy_insert, false);
+		var remaining = last_drive.insertObject(new ItemStackWrapper(copy_insert), false);
 
-		if (remaining.getCount() == copy_insert.getCount())
+		if (remaining.getCount().intValue() == copy_insert.getCount())
 			return ItemStack.EMPTY;
 
 		var copy = item.copy();
-		tile.getIOItemCount(amount - remaining.getCount(), false);
+		tile.getIOItemCount(amount - remaining.getCount().intValue(), false);
-		item.shrink(amount - remaining.getCount());
+		item.shrink(amount - remaining.getCount().intValue());
 		slot.setChanged();
 
 		return copy;

src/main/java/ru/dbotthepony/mc/otm/storage/IStorageStack.java

@@ -35,8 +35,8 @@ public interface IStorageStack {
 	 * @return Identity utilized to partition view table (if it has any).
 	 *         Defaults to no partitioning; meaning performance will degrade much quicker than it should.
 	 *         Good object is an object that will influence on sameItem() return result, making it return true.
-	 *         It is weakly considered that if !this.itemIdentity().equals(other.itemIdentity());
+	 *         It is strictly considered that if !this.itemIdentity().equals(other.itemIdentity());
-	 *         then this.sameItem(other) will always return false.
+	 *         then this.sameItem(other) will never return true.
 	 *
 	 *         If implemented, and storage is also partitioned properly, then performance will be flat equally to view table's performance.
 	 *

src/main/java/ru/dbotthepony/mc/otm/storage/StorageObjectRegistry.java

@@ -13,12 +13,10 @@ import java.util.Objects;
 public class StorageObjectRegistry {
 	private static final HashMap<Class<? extends IStorageStack>, StorageObjectTuple<? extends IStorageStack>> REGISTRY = new HashMap<>();
 
-	public static <T extends IStorageStack> boolean register(Class<T> identity, T empty) {
+	public static <T extends IStorageStack> StorageObjectTuple<T> register(Class<T> identity, T empty) {
 		final var tuple = new StorageObjectTuple<>(identity, empty);
-		final boolean removed = REGISTRY.remove(identity) != null;
 		REGISTRY.put(identity, tuple);
-
+		return tuple;
-		return removed;
 	}
 
 	@SuppressWarnings("unchecked")