Disk drives work again

src/main/java/refinedstorage/storage/NBTStorage.java

@@ -21,7 +21,9 @@ public class NBTStorage implements IStorage {
     private NBTTagCompound tag;
     private int capacity;
     private int priority;
+    private boolean dirty;
 
+    // We use a map here because is much faster than looping over a list
     private Map<ItemGroupMeta, Integer> groups = new HashMap<ItemGroupMeta, Integer>();
 
     public NBTStorage(NBTTagCompound tag, int capacity, int priority) {
@@ -48,6 +50,26 @@ public class NBTStorage implements IStorage {
         }
     }
 
+    public void writeToNBT(NBTTagCompound tag) {
+        NBTTagList list = new NBTTagList();
+
+        for (Map.Entry<ItemGroupMeta, Integer> entry : groups.entrySet()) {
+            NBTTagCompound itemTag = new NBTTagCompound();
+
+            itemTag.setInteger(NBT_ITEM_TYPE, Item.getIdFromItem(entry.getKey().getType()));
+            itemTag.setInteger(NBT_ITEM_QUANTITY, entry.getValue());
+            itemTag.setInteger(NBT_ITEM_DAMAGE, entry.getKey().getDamage());
+
+            if (entry.getKey().hasTag()) {
+                itemTag.setTag(NBT_ITEM_NBT, entry.getKey().getTag());
+            }
+
+            list.appendTag(itemTag);
+        }
+
+        tag.setTag(NBT_ITEMS, list);
+    }
+
     @Override
     public void addItems(List<ItemGroup> items) {
         for (Map.Entry<ItemGroupMeta, Integer> entry : groups.entrySet()) {
@@ -57,6 +79,10 @@ public class NBTStorage implements IStorage {
 
     @Override
     public void push(ItemStack stack) {
+        markDirty();
+
+        tag.setInteger(NBT_STORED, getStored(tag) + stack.stackSize);
+
         for (Map.Entry<ItemGroupMeta, Integer> entry : groups.entrySet()) {
             if (entry.getKey().compareNoQuantity(stack)) {
                 groups.put(entry.getKey(), entry.getValue() + stack.stackSize);
@@ -70,6 +96,36 @@ public class NBTStorage implements IStorage {
 
     @Override
     public ItemStack take(ItemStack stack, int flags) {
+        int quantity = stack.stackSize;
+
+        for (Map.Entry<ItemGroupMeta, Integer> entry : groups.entrySet()) {
+            if (entry.getKey().compare(stack, flags)) {
+                if (quantity > entry.getValue()) {
+                    quantity = entry.getValue();
+                }
+
+                if (entry.getValue() - quantity == 0) {
+                    groups.remove(entry.getKey());
+                } else {
+                    groups.put(entry.getKey(), entry.getValue() - quantity);
+                }
+
+                tag.setInteger(NBT_STORED, getStored(tag) - quantity);
+
+                ItemStack result = new ItemStack(
+                    entry.getKey().getType(),
+                    quantity,
+                    entry.getKey().getDamage()
+                );
+
+                result.setTagCompound(entry.getKey().getTag());
+
+                markDirty();
+
+                return result;
+            }
+        }
+
         return null;
     }
 
@@ -91,6 +147,18 @@ public class NBTStorage implements IStorage {
         return tag;
     }
 
+    public void markDirty() {
+        this.dirty = true;
+    }
+
+    public boolean isDirty() {
+        return dirty;
+    }
+
+    public void markClean() {
+        this.dirty = false;
+    }
+
     public static int getStored(NBTTagCompound tag) {
         return tag.getInteger(NBT_STORED);
     }

src/main/java/refinedstorage/tile/TileDiskDrive.java

@@ -29,12 +29,22 @@ public class TileDiskDrive extends TileMachine implements IStorageProvider, ISto
     private InventorySimple inventory = new InventorySimple("disk_drive", 8, this);
     private InventorySimple filterInventory = new InventorySimple("filters", 9, this);
 
-    private NBTStorage storages[] = new NBTStorage[getSizeInventory()];
+    private NBTStorage storages[] = new NBTStorage[8];
 
     private int priority = 0;
     private int compare = 0;
     private int mode = 0;
 
+    public NBTStorage getStorage(int slot) {
+        if (inventory.getStackInSlot(slot) == null) {
+            storages[slot] = null;
+        } else if (storages[slot] == null) {
+            storages[slot] = new DiskStorage(getStackInSlot(slot), this);
+        }
+
+        return storages[slot];
+    }
+
     @Override
     public int getEnergyUsage() {
         int base = 5;
@@ -50,11 +60,21 @@ public class TileDiskDrive extends TileMachine implements IStorageProvider, ISto
 
     @Override
     public void updateMachine() {
+        for (int i = 0; i < getSizeInventory(); ++i) {
+            NBTStorage storage = getStorage(i);
+
+            if (storage != null && storage.isDirty()) {
+                storage.writeToNBT(getStackInSlot(i).getTagCompound());
+                storage.markClean();
+            }
+        }
     }
 
     @Override
     public void provide(List<IStorage> storages) {
-        for (NBTStorage storage : this.storages) {
+        for (int i = 0; i < getSizeInventory(); ++i) {
+            NBTStorage storage = getStorage(i);
+
             if (storage != null) {
                 storages.add(storage);
             }
@@ -251,12 +271,6 @@ public class TileDiskDrive extends TileMachine implements IStorageProvider, ISto
 
     @Override
     public void setInventorySlotContents(int slot, ItemStack stack) {
-        if (stack != null) {
-            storages[slot] = new DiskStorage(stack, this);
-        } else {
-            storages[slot] = null;
-        }
-
         inventory.setInventorySlotContents(slot, stack);
     }