Improved "quick move" code in menus

src/main/kotlin/ru/dbotthepony/mc/otm/core/util/ItemStackKey.kt

@@ -2,12 +2,17 @@ package ru.dbotthepony.mc.otm.core.util
 
 import it.unimi.dsi.fastutil.HashCommon
 import net.minecraft.core.component.DataComponentMap
+import net.minecraft.core.component.DataComponentPatch
+import net.minecraft.core.component.PatchedDataComponentMap
+import net.minecraft.core.registries.BuiltInRegistries
 import net.minecraft.world.item.Item
 import net.minecraft.world.item.ItemStack
+import ru.dbotthepony.mc.otm.core.getHolder
 
-class ItemStackKey(val item: Item, val components: DataComponentMap) {
+class ItemStackKey(val item: Item, val components: DataComponentPatch) {
 	// make copy of original itemstack because there is no copy() method on DataComponentMap, which is returned by ItemStack#getComponents
-	constructor(itemStack: ItemStack) : this(itemStack.item, itemStack.copy().components)
+	constructor(itemStack: ItemStack) : this(itemStack.item, itemStack.copy().componentsPatch)
+	constructor(item: Item) : this(item, DataComponentPatch.EMPTY)
 
 	private var hashComputed = false
 	private var hash = 0
@@ -25,6 +30,10 @@ class ItemStackKey(val item: Item, val components: DataComponentMap) {
 		return hash
 	}
 
+	fun asItemStack(count: Int = 1): ItemStack {
+		return ItemStack(BuiltInRegistries.ITEM.getHolder(item)!!, count, components)
+	}
+
 	override fun toString(): String {
 		return "ItemStackKey[$item, $components]"
 	}
@@ -33,3 +42,12 @@ class ItemStackKey(val item: Item, val components: DataComponentMap) {
 fun ItemStack.asKey(): ItemStackKey {
 	return ItemStackKey(this)
 }
+
+fun ItemStack.asKeyOrNull(): ItemStackKey? {
+	if (isEmpty) return null
+	return ItemStackKey(this)
+}
+
+fun Item.asKey(): ItemStackKey {
+	return ItemStackKey(this)
+}

src/main/kotlin/ru/dbotthepony/mc/otm/menu/ExopackInventoryMenu.kt

@@ -177,7 +177,7 @@ class ExopackInventoryMenu(val capability: MatteryPlayer) : MatteryMenu(null, CO
 
 		if (!player.level().isClientSide) {
 			for (slot in craftingGrid.slotIterator()) {
-				val leftover = moveItemStackToSlots(slot.item, playerInventorySlots)
+				val leftover = QuickMoveInput.moveItemStackToSlots(slot.item, playerInventorySlots)
 
 				if (!leftover.isEmpty) {
 					player.drop(leftover, true)
@@ -205,11 +205,11 @@ class ExopackInventoryMenu(val capability: MatteryPlayer) : MatteryMenu(null, CO
 
 		if (slotIndex == craftingResultSlot.index) {
 			val item = craftingResultSlot.item
-			val leftover = moveItemStackToSlots(item, playerInventorySlots, simulate = true)
+			val leftover = QuickMoveInput.moveItemStackToSlots(item, playerInventorySlots, simulate = true)
 
 			if (leftover.isEmpty) {
 				val copy = item.copy()
-				moveItemStackToSlots(item, playerInventorySlots, simulate = false)
+				QuickMoveInput.moveItemStackToSlots(item, playerInventorySlots, simulate = false)
 				item.count = 0
 				craftingResultSlot.onTake(ply, copy)
 				return copy

src/main/kotlin/ru/dbotthepony/mc/otm/menu/MatteryMenu.kt

@@ -452,25 +452,10 @@ abstract class MatteryMenu(
 		val copy = slot.item.copy()
 		var any = false
 
-		if (target.any { it.any { it.containerSlotOrNull() is IFilteredContainerSlot } }) {
-			for (collection in target) {
-				if (moveItemStackTo(ply, slot, collection, onlyFiltered = true)) {
-					any = true
-
-					if (!slot.hasItem()) {
-						return copy
-					}
-				}
-			}
-		}
-
 		for (collection in target) {
-			if (moveItemStackTo(ply, slot, collection)) {
+			if (QuickMoveInput.moveItemStackTo(ply, slot, collection)) {
 				any = true
-
+				if (!slot.hasItem()) return copy
-				if (!slot.hasItem()) {
-					return copy
-				}
 			}
 		}
 
@@ -486,7 +471,7 @@ abstract class MatteryMenu(
 			return stack
 		}
 
-		return moveItemStackToSlots(stack, _playerInventorySlots, simulate = simulate)
+		return QuickMoveInput.moveItemStackToSlots(stack, _playerInventorySlots, simulate = simulate)
 	}
 
 	override fun canTakeItemForPickAll(itemStack: ItemStack, slot: Slot): Boolean {
@@ -521,22 +506,12 @@ abstract class MatteryMenu(
 		require(finalSlot < slots.size) { "Final slot $finalSlot is bigger than total size of array of ${slots.size}" }
 
 		val slots = ArrayList<Slot>(finalSlot - initialSlot + 1)
-		var filters = false
 
 		for (i in (if (inverse) finalSlot downTo initialSlot else initialSlot .. finalSlot)) {
-			val slot = slots[i]
+			slots.add(this.slots[i])
-			slots.add(slot)
-
-			if (slot.containerSlotOrNull() is IFilteredContainerSlot) {
-				filters = true
-			}
 		}
 
-		if (filters) {
+		return QuickMoveInput.moveItemStackToSlots(item, slots, simulate)
-			return moveItemStackToSlots(moveItemStackToSlots(item, slots, simulate, onlyFiltered = true), slots, simulate, onlyFiltered = false)
-		}
-
-		return moveItemStackToSlots(item, slots, simulate)
 	}
 
 	private var armorSlots: ImmutableList<PlayerSlot<EquipmentMenuSlot, Slot>>? = null
@@ -613,96 +588,5 @@ abstract class MatteryMenu(
 			InventoryMenu.EMPTY_ARMOR_SLOT_LEGGINGS,
 			InventoryMenu.EMPTY_ARMOR_SLOT_CHESTPLATE,
 			InventoryMenu.EMPTY_ARMOR_SLOT_HELMET)
-
-		fun moveItemStackTo(
-			player: Player,
-			source: Slot,
-			slots: Collection<Slot>,
-			onlyFiltered: Boolean = false
-		): Boolean {
-			val remainder = moveItemStackToSlots(source.item, slots, onlyFiltered = onlyFiltered)
-
-			if (remainder.count == source.item.count) {
-				return false
-			}
-
-			val copy = source.item.copy()
-
-			if (remainder.isEmpty) {
-				source.setByPlayer(ItemStack.EMPTY)
-				source.onTake(player, copy)
-			} else {
-				copy.count = source.item.count - remainder.count
-				source.item.count = remainder.count
-				source.onTake(player, copy)
-			}
-
-			return true
-		}
-
-
-		fun moveItemStackToSlots(item: ItemStack, slots: Collection<Slot>, simulate: Boolean = false, onlyFiltered: Boolean = false): ItemStack {
-			if (item.isEmpty) {
-				return ItemStack.EMPTY
-			}
-
-			val copy = item.copy()
-
-			// first pass - stack with existing slots
-			if (copy.isStackable) {
-				for (slot in slots) {
-					if (onlyFiltered && slot.containerSlotOrNull().let { it !is IFilteredContainerSlot || it.filter == null || !it.testSlotFilter(item) }) {
-						continue
-					} else if (!onlyFiltered && slot.containerSlotOrNull().let { it is IFilteredContainerSlot && it.filter != null }) {
-						continue
-					}
-
-					val limit = slot.getMaxStackSize(copy)
-
-					if (limit > slot.item.count && slot.mayPlace(item) && ItemStack.isSameItemSameComponents(slot.item, copy)) {
-						val newCount = (slot.item.count + copy.count).coerceAtMost(limit)
-						val diff = newCount - slot.item.count
-						copy.count -= diff
-
-						if (!simulate) {
-							slot.item.count += diff
-							slot.setChanged()
-						}
-
-						if (copy.isEmpty) {
-							return copy
-						}
-					}
-				}
-			}
-
-			// second pass - drop stack into first free slot
-			for (slot in slots) {
-				if (onlyFiltered && slot.containerSlotOrNull().let { it !is IFilteredContainerSlot || it.filter == null || !it.testSlotFilter(item) }) {
-					continue
-				} else if (!onlyFiltered && slot.containerSlotOrNull().let { it is IFilteredContainerSlot && it.filter != null }) {
-					continue
-				}
-
-				val limit = slot.getMaxStackSize(copy)
-
-				if (!slot.hasItem() && slot.mayPlace(item)) {
-					val newCount = copy.count.coerceAtMost(limit)
-
-					if (!simulate) {
-						slot.setByPlayer(copy.copy().also { it.count = newCount })
-						slot.setChanged()
-					}
-
-					copy.count -= newCount
-
-					if (copy.isEmpty) {
-						return copy
-					}
-				}
-			}
-
-			return copy
-		}
 	}
 }

src/main/kotlin/ru/dbotthepony/mc/otm/menu/QuickMoveInput.kt

@@ -0,0 +1,203 @@
+package ru.dbotthepony.mc.otm.menu
+
+import net.minecraft.world.entity.player.Player
+import net.minecraft.world.inventory.Slot
+import net.minecraft.world.item.ItemStack
+import net.minecraft.world.item.Items
+import ru.dbotthepony.mc.otm.container.IFilteredContainerSlot
+import ru.dbotthepony.mc.otm.container.util.containerSlotOrNull
+import ru.dbotthepony.mc.otm.core.isNotEmpty
+import ru.dbotthepony.mc.otm.core.util.ItemStackKey
+import ru.dbotthepony.mc.otm.core.util.asKey
+import ru.dbotthepony.mc.otm.core.util.asKeyOrNull
+
+class QuickMoveInput(private val menu: MatteryMenu, val from: Collection<Slot>, val to: Collection<Slot>, val mode: Mode) {
+	enum class Mode {
+		RESTOCK {
+			override fun move(from: Collection<Slot>, to: Collection<Slot>, player: Player) {
+				val (_, itemsFrom) = computeSlotLists(from, true)
+				val (_, itemsTo) = computeSlotLists(to, false)
+
+				val intersect = if (itemsFrom.size < itemsTo.size) itemsFrom.keys.filter { it in itemsTo.keys } else itemsTo.keys.filter { it in itemsFrom.keys }
+
+				for (key in intersect) {
+					val slotsTo = itemsTo[key]!!
+					itemsFrom[key]!!.forEach { moveItemStackTo(player, it, slotsTo) }
+				}
+			}
+		},
+
+		RESTOCK_WITH_MOVE {
+			override fun move(from: Collection<Slot>, to: Collection<Slot>, player: Player) {
+				val (_, itemsFrom) = computeSlotLists(from, true)
+				val (emptyTo, itemsTo) = computeSlotLists(to, false)
+
+				val intersect = if (itemsFrom.size < itemsTo.size) itemsFrom.keys.filter { it in itemsTo.keys } else itemsTo.keys.filter { it in itemsFrom.keys }
+
+				for (key in intersect) {
+					val slotsTo = prioritySortSlots(itemsTo[key]!!, key.asItemStack())
+					val slotsFrom = itemsFrom[key]!!
+					slotsFrom.removeIf { moveItemStackTo(player, it, slotsTo, sort = false); it.item.isEmpty }
+					var moveAny = false
+					slotsFrom.forEach { moveAny = moveItemStackTo(player, it, emptyTo, sort = false) || moveAny }
+					if (moveAny) emptyTo.removeIf { it.item.isNotEmpty }
+				}
+			}
+		},
+
+		MOVE {
+			override fun move(from: Collection<Slot>, to: Collection<Slot>, player: Player) {
+				val toSorted = prioritySortSlots(to)
+
+				from.forEach {
+					val slot = it.containerSlotOrNull()
+
+					if (slot !is IFilteredContainerSlot || !slot.hasFilter)
+						moveItemStackTo(player, it, toSorted, sort = false)
+				}
+			}
+		};
+
+		abstract fun move(from: Collection<Slot>, to: Collection<Slot>, player: Player)
+	}
+
+	private val input = menu.oneWayInput(handler = ::handle)
+
+	fun trigger() {
+		input.accept(null)
+	}
+
+	private fun handle() {
+		mode.move(from, to, menu.player)
+	}
+
+	companion object {
+		private fun computeSlotLists(slots: Collection<Slot>, skipFilteredSlots: Boolean): Pair<MutableList<Slot>, MutableMap<ItemStackKey, MutableList<Slot>>> {
+			val emptySlots = ArrayList<Slot>()
+			val filledSlots = HashMap<ItemStackKey, MutableList<Slot>>()
+
+			for (slot in slots) {
+				val underlyingSlot = slot.containerSlotOrNull()
+
+				if (underlyingSlot is IFilteredContainerSlot && (underlyingSlot.filter == Items.AIR || underlyingSlot.filter != null && skipFilteredSlots))
+					continue
+
+				val key = slot.item.asKeyOrNull() ?: (underlyingSlot as? IFilteredContainerSlot)?.filter?.asKey()
+
+				if (key == null) {
+					emptySlots.add(slot)
+				} else {
+					filledSlots.computeIfAbsent(key) { ArrayList() }.add(slot)
+				}
+			}
+
+			return emptySlots to filledSlots
+		}
+
+		fun moveItemStackTo(
+			player: Player,
+			source: Slot,
+			slots: Collection<Slot>,
+			sort: Boolean = true
+		): Boolean {
+			if (!source.mayPickup(player) || source.item.isEmpty || slots.isEmpty())
+				return false
+
+			val remainder = moveItemStackToSlots(source.item, slots, sort = sort)
+
+			if (remainder.count == source.item.count)
+				return false
+
+			val copy = source.item.copy()
+
+			if (remainder.isEmpty) {
+				source.setByPlayer(ItemStack.EMPTY)
+				source.onTake(player, copy)
+			} else {
+				copy.count = source.item.count - remainder.count
+				source.item.count = remainder.count
+				source.onTake(player, copy)
+			}
+
+			return true
+		}
+
+		fun prioritySortSlots(slots: Collection<Slot>, filterItem: ItemStack? = null): MutableList<Slot> {
+			val sortedSlots = ArrayList(slots)
+
+			sortedSlots.removeIf {
+				val slot = it.containerSlotOrNull()
+				it.isOverCapacity || filterItem != null && !it.mayPlace(filterItem) || slot is IFilteredContainerSlot && slot.isForbiddenForAutomation
+			}
+
+			sortedSlots.sortWith { a, b ->
+				val hasItemA = a.item.isNotEmpty
+				val hasItemB = b.item.isNotEmpty
+
+				if (hasItemA && hasItemB)
+					return@sortWith 0
+				else if (hasItemA)
+					return@sortWith -1
+				else if (hasItemB)
+					return@sortWith 1
+
+				val slotA = a.containerSlotOrNull()
+				val slotB = b.containerSlotOrNull()
+
+				val hasFilterA = slotA is IFilteredContainerSlot && slotA.hasFilter
+				val hasFilterB = slotB is IFilteredContainerSlot && slotB.hasFilter
+
+				if (hasFilterA && hasFilterB || !hasFilterA && !hasFilterB)
+					return@sortWith 0
+				else if (hasFilterA)
+					return@sortWith -1
+				else
+					return@sortWith 1
+			}
+
+			return sortedSlots
+		}
+
+		fun moveItemStackToSlots(item: ItemStack, slots: Collection<Slot>, simulate: Boolean = false, sort: Boolean = true): ItemStack {
+			if (item.isEmpty)
+				return ItemStack.EMPTY
+			else if (slots.isEmpty())
+				return item.copy()
+
+			val sortedSlots = if (sort) prioritySortSlots(slots, item) else slots
+			val copy = item.copy()
+
+			for (slot in sortedSlots) {
+				val limit = slot.getMaxStackSize(copy)
+
+				if (!slot.hasItem()) {
+					val newCount = copy.count.coerceAtMost(limit)
+
+					if (!simulate) {
+						slot.setByPlayer(copy.copy().also { it.count = newCount })
+						// slot.setChanged()
+					}
+
+					copy.shrink(newCount)
+
+					if (copy.isEmpty)
+						return ItemStack.EMPTY
+				} else if (limit > slot.item.count && ItemStack.isSameItemSameComponents(slot.item, copy)) {
+					val newCount = (slot.item.count + copy.count).coerceAtMost(limit)
+					val diff = newCount - slot.item.count
+					copy.count -= diff
+
+					if (!simulate) {
+						slot.item.count += diff
+						slot.setChanged()
+					}
+
+					if (copy.isEmpty)
+						return ItemStack.EMPTY
+				}
+			}
+
+			return copy
+		}
+	}
+}

src/main/kotlin/ru/dbotthepony/mc/otm/menu/Slots.kt

@@ -26,6 +26,7 @@ import ru.dbotthepony.mc.otm.container.UpgradeContainer
 import ru.dbotthepony.mc.otm.container.util.containerSlotOrNull
 import ru.dbotthepony.mc.otm.core.collect.ConditionalEnumSet
 import ru.dbotthepony.mc.otm.core.immutableList
+import ru.dbotthepony.mc.otm.core.isNotEmpty
 import ru.dbotthepony.mc.otm.core.math.Decimal
 import ru.dbotthepony.mc.otm.menu.input.BooleanInputWithFeedback
 import ru.dbotthepony.mc.otm.menu.input.InstantBooleanInput
@@ -242,6 +243,10 @@ fun MatteryMenu.addFilterControls(slots: KMutableProperty0<ItemFilter>?, amount:
 	return addFilterControls(slots?.let { Delegate.Of(it) }, amount)
 }
 
+val Slot.isOverCapacity: Boolean get() {
+	return item.isNotEmpty && getMaxStackSize(item) <= item.count
+}
+
 /**
  * [openState] **is clientside only**, attempting to use it on server will result
  * in classloading exceptions.

src/main/kotlin/ru/dbotthepony/mc/otm/menu/storage/ItemMonitorMenu.kt

@@ -17,6 +17,7 @@ import ru.dbotthepony.mc.otm.core.collect.reduce
 import ru.dbotthepony.mc.otm.core.isNotEmpty
 import ru.dbotthepony.mc.otm.menu.MatteryPoweredMenu
 import ru.dbotthepony.mc.otm.menu.MatteryMenuSlot
+import ru.dbotthepony.mc.otm.menu.QuickMoveInput
 import ru.dbotthepony.mc.otm.menu.data.INetworkedItemViewProvider
 import ru.dbotthepony.mc.otm.menu.data.NetworkedItemView
 import ru.dbotthepony.mc.otm.menu.input.BooleanInputWithFeedback
@@ -137,9 +138,9 @@ class ItemMonitorMenu(
 
 		if (settings.resultTarget == ItemMonitorPlayerSettings.ResultTarget.ALL_SYSTEM) {
 			remaining = view.insertStack(ItemStorageStack(itemStack), simulate).toItemStack()
-			remaining = moveItemStackToSlots(remaining, playerInventorySlots, simulate)
+			remaining = QuickMoveInput.moveItemStackToSlots(remaining, playerInventorySlots, simulate)
 		} else {
-			remaining = moveItemStackToSlots(itemStack, playerInventorySlots, simulate)
+			remaining = QuickMoveInput.moveItemStackToSlots(itemStack, playerInventorySlots, simulate)
 			remaining = view.insertStack(ItemStorageStack(remaining), simulate).toItemStack()
 		}
 
@@ -184,7 +185,7 @@ class ItemMonitorMenu(
 				else -> {}
 			}
 
-			remainder = moveItemStackToSlots(remainder, playerInventorySlots)
+			remainder = QuickMoveInput.moveItemStackToSlots(remainder, playerInventorySlots)
 
 			slots[slotIndex].set(remainder)
 			return if (remainder.count != item.count) item else ItemStack.EMPTY