From f6d67a61f35eacfc578116cc41312d90291e6bd2 Mon Sep 17 00:00:00 2001
From: Luca_S <luca@acul.me>
Date: Mon, 4 May 2020 23:23:40 +0200
Subject: [PATCH] Handle locked chests

---
 milo/apis/storage.lua | 37 +++++++++++++++++++++++++++++--------
 1 file changed, 29 insertions(+), 8 deletions(-)

diff --git a/milo/apis/storage.lua b/milo/apis/storage.lua
index 34c0d34..143b801 100644
--- a/milo/apis/storage.lua
+++ b/milo/apis/storage.lua
@@ -327,7 +327,7 @@ function Storage:listProviders(throttle)
 			if not res[key] then
 				res[key] = {}
 			end
-			table.insert(res[key], {item = item, device = device[chest], slot = slot})
+			table.insert(res[key], {item = item, device = device[chest], lockedToThis = Util.contains(Util.keys(self.nodes[chest].lock or {}), key) and true or false, slot = slot})
 		end
 	end
 	return res
@@ -338,10 +338,29 @@ function Storage:defrag(throttle)
 	local items = self:listProviders(throttle)
 	local slotsSaved = 0
 
-	for _, providers in pairs(items) do
-		table.sort(providers, function(a, b)
+	-- This will make sure the table is sorted in the following order:
+	-- Unlocked stacks with less than maxCount items | Locked stacks with less than maxCount items | stacks with more than maxCount items
+	-- This way the locked stacks will be filled first
+	local function sortFunction(a, b)
+		local preferenceA, preferenceB
+		preferenceA = (a.item.count == a.item.maxCount and 3)
+									or (a.lockedToThis and 2)
+									or 1
+		preferenceB = (b.item.count == b.item.maxCount and 3)
+									or (b.lockedToThis and 2)
+									or 1
+
+		if preferenceA < preferenceB then
+			return true
+		elseif preferenceB > preferenceA then
+			return false
+		else
 			return a.item.count < b.item.count
-		end)
+		end
+	end
+
+	for _, providers in pairs(items) do
+		table.sort(providers, sortFunction)
 
 		-- We're done when we either compressed the stacks so far, that there's only one left (#providers == 1)
 		-- Or when we've compressed so far, that the there's only one stack which has a lower count than the maxCount
@@ -354,10 +373,14 @@ function Storage:defrag(throttle)
 			-- This loop is guarenteed to assign a value to "to", as the only cases where it wouldn't (#providers == 1 or no provider with less than maxCount)
 			-- are ruled out by the condition of the outer while loop
 			for i = 2, #providers do
-				if providers[i].item.count < providers[i].item.maxCount then
+				-- Give preference to locked chests
+				if not to.lockedToThis and providers[i].lockedToThis then
+					to = providers[i]
+				elseif (to.lockedToThis == providers[i].lockedToThis) and providers[i].item.count < providers[i].item.maxCount then
 					to = providers[i]
 				else
 					-- As this slot is already at maxCount, all the remaining ones will also be due to sorting
+					-- If any of the remaining providers is locked that doesn't matter. We wouldn't have been able to push there anyways
 					break
 				end
 			end
@@ -389,9 +412,7 @@ function Storage:defrag(throttle)
 				slotsSaved = slotsSaved + 1
 			end
 
-			table.sort(providers, function(a, b)
-				return a.item.count < b.item.count
-			end)
+			table.sort(providers, sortFunction)
 		end
 	end