Aware items stream

src/main/kotlin/ru/dbotthepony/mc/otm/android/AndroidResearch.kt

@@ -12,6 +12,7 @@ import net.minecraft.world.item.ItemStack
 import net.minecraftforge.common.util.INBTSerializable
 import ru.dbotthepony.mc.otm.OverdriveThatMatters
 import ru.dbotthepony.mc.otm.capability.MatteryPlayerCapability
+import ru.dbotthepony.mc.otm.capability.awareItemsStream
 import ru.dbotthepony.mc.otm.capability.itemsStream
 import ru.dbotthepony.mc.otm.client.render.SkinElement
 import ru.dbotthepony.mc.otm.container.iterator
@@ -136,20 +137,11 @@ class AndroidResearch(val type: AndroidResearchType, val capability: MatteryPlay
 
 		for ((tag, count) in type.items) {
 			var required = count
-			val iterator = capability.ply.inventory.iterator().nonEmpty()
 
-			for (invItem in iterator) {
-				if (tag.test(invItem)) {
-					val toExtract = required.coerceAtMost(invItem.count)
-					required -= toExtract
-
-					if (!simulate) {
-						if (toExtract == invItem.count) {
-							iterator.remove()
-						} else {
-							invItem.count -= toExtract
-						}
-					}
+			for (invItem in ply.awareItemsStream(false)) {
+				if (tag.test(invItem.itemStack)) {
+					val toExtract = required.coerceAtMost(invItem.itemStack.count)
+					required -= invItem.extract(toExtract, simulate).count
 
 					if (required <= 0) {
 						break
@@ -336,11 +328,11 @@ class AndroidResearch(val type: AndroidResearchType, val capability: MatteryPlay
 
 		for ((tag, count) in type.items) {
 			var required = count
-			val iterator = capability.ply.inventory.iterator().nonEmpty()
 
-			for (invItem in iterator) {
-				if (tag.test(invItem)) {
-					required -= required.coerceAtMost(invItem.count)
+			for (invItem in ply.awareItemsStream(false)) {
+				if (tag.test(invItem.itemStack)) {
+					val toExtract = required.coerceAtMost(invItem.itemStack.count)
+					required -= invItem.extract(toExtract, true).count
 
 					if (required <= 0) {
 						break

src/main/kotlin/ru/dbotthepony/mc/otm/capability/Ext.kt

@@ -9,15 +9,19 @@ import net.minecraftforge.common.capabilities.ICapabilityProvider
 import net.minecraftforge.common.util.LazyOptional
 import net.minecraftforge.energy.IEnergyStorage
 import net.minecraftforge.fml.ModList
+import ru.dbotthepony.mc.otm.compat.curios.curiosAwareStream
 import ru.dbotthepony.mc.otm.compat.curios.curiosStream
 import ru.dbotthepony.mc.otm.compat.curios.isCuriosLoaded
 import ru.dbotthepony.mc.otm.compat.mekanism.getMekanismEnergySided
 import ru.dbotthepony.mc.otm.compat.mekanism.mekanismEnergy
-import ru.dbotthepony.mc.otm.container.iterator
+import ru.dbotthepony.mc.otm.container.awareStream
 import ru.dbotthepony.mc.otm.container.stream
+import ru.dbotthepony.mc.otm.core.AwareItemStack
+import ru.dbotthepony.mc.otm.core.ContainerItemStackEntry
 import ru.dbotthepony.mc.otm.core.ImpreciseFraction
 import ru.dbotthepony.mc.otm.core.orNull
 import java.util.IdentityHashMap
+import java.util.LinkedList
 import java.util.stream.Stream
 
 val ICapabilityProvider.matteryPlayer: MatteryPlayerCapability? get() = getCapability(MatteryCapability.MATTERY_PLAYER).orNull()
@@ -156,7 +160,7 @@ fun ICapabilityProvider.getMatteryEnergySided(side: Direction? = null): LazyOpti
  * Contains all items that player might carry
  */
 fun Player.itemsStream(includeCosmetics: Boolean = true): Stream<out ItemStack> {
-	val streams = ArrayList<Stream<out ItemStack>>()
+	val streams = LinkedList<Stream<out ItemStack>>()
 	streams.add(inventory.stream())
 
 	matteryPlayer?.let {
@@ -178,8 +182,8 @@ fun Player.itemsStream(includeCosmetics: Boolean = true): Stream<out ItemStack>
  * Contains all items that player might see/access ([itemsStream] + open container's items)
  */
 fun Player.allItemsStream(includeCosmetics: Boolean = true): Stream<out ItemStack> {
-	if (containerMenu == inventoryMenu) {
-		return itemsStream()
+	if (containerMenu == inventoryMenu || containerMenu == matteryPlayer?.exoSuitMenu) {
+		return itemsStream(includeCosmetics)
 	}
 
 	val seen = IdentityHashMap<ItemStack, Unit>(containerMenu.slots.size)
@@ -192,3 +196,47 @@ fun Player.allItemsStream(includeCosmetics: Boolean = true): Stream<out ItemStac
 		itemsStream(includeCosmetics).filter { it.isEmpty || it !in seen },
 		containerMenu.slots.stream().map { it.item })
 }
+
+/**
+ * Modify returned [ItemStack]s through [AwareItemStack.extract]
+ *
+ * Contains all items that player might carry
+ */
+fun Player.awareItemsStream(includeCosmetics: Boolean = true): Stream<out AwareItemStack> {
+	val streams = LinkedList<Stream<out AwareItemStack>>()
+	streams.add(inventory.awareStream())
+
+	matteryPlayer?.let {
+		if (it.hasExoSuit) {
+			streams.add(it.exoSuitContainer.awareStream())
+		}
+	}
+
+	if (isCuriosLoaded) {
+		streams.add(curiosAwareStream(includeCosmetics))
+	}
+
+	return Streams.concat(*streams.toTypedArray())
+}
+
+/**
+ * Modify returned [ItemStack]s through [AwareItemStack.extract]
+ *
+ * Contains all items that player might see/access ([itemsStream] + open container's items)
+ */
+fun Player.awareAllItemsStream(includeCosmetics: Boolean = true): Stream<out AwareItemStack> {
+	if (containerMenu == inventoryMenu || containerMenu == matteryPlayer?.exoSuitMenu) {
+		return awareItemsStream(includeCosmetics)
+	}
+
+	val seen = HashSet<AwareItemStack>(containerMenu.slots.size)
+
+	for (slot in containerMenu.slots) {
+		seen.add(ContainerItemStackEntry(slot.slotIndex, slot.container))
+	}
+
+	return Streams.concat(
+		awareItemsStream(includeCosmetics).filter { it !in seen },
+		seen.stream())
+}
+

src/main/kotlin/ru/dbotthepony/mc/otm/compat/curios/CuriosCompat.kt

@@ -1,5 +1,6 @@
 package ru.dbotthepony.mc.otm.compat.curios
 
+import com.google.common.collect.Iterators
 import com.google.common.collect.Streams
 import it.unimi.dsi.fastutil.ints.Int2ObjectArrayMap
 import net.minecraft.world.entity.player.Player
@@ -7,11 +8,18 @@ import net.minecraft.world.inventory.Slot
 import net.minecraft.world.item.ItemStack
 import net.minecraftforge.fml.ModList
 import ru.dbotthepony.mc.otm.capability.MatteryCapability
+import ru.dbotthepony.mc.otm.container.awareIterator
+import ru.dbotthepony.mc.otm.container.awareStream
+import ru.dbotthepony.mc.otm.container.iterator
 import ru.dbotthepony.mc.otm.container.stream
+import ru.dbotthepony.mc.otm.core.AwareItemStack
+import ru.dbotthepony.mc.otm.core.EmptyMutableIterator
+import ru.dbotthepony.mc.otm.core.IAwareItemStackIterator
 import ru.dbotthepony.mc.otm.core.orNull
 import top.theillusivec4.curios.api.CuriosApi
 import top.theillusivec4.curios.common.inventory.CosmeticCurioSlot
 import top.theillusivec4.curios.common.inventory.CurioSlot
+import java.util.LinkedList
 import java.util.stream.Stream
 
 val isCuriosLoaded by lazy {
@@ -67,7 +75,7 @@ val Player.curiosSlots: Collection<Pair<Slot, Slot?>> get() {
 private fun Player.curiosStreamImpl(includeCosmetics: Boolean): Stream<out ItemStack> {
 	val handler = getCapability(MatteryCapability.CURIOS_INVENTORY).orNull() ?: return Stream.empty()
 
-	val result = ArrayList<Stream<out ItemStack>>()
+	val result = LinkedList<Stream<out ItemStack>>()
 
 	for ((identifier, curio) in handler.curios) {
 		result.add(curio.stacks.stream())
@@ -87,3 +95,27 @@ fun Player.curiosStream(includeCosmetics: Boolean = true): Stream<out ItemStack>
 
 	return curiosStreamImpl(includeCosmetics)
 }
+
+private fun Player.curiosAwareStreamImpl(includeCosmetics: Boolean): Stream<out AwareItemStack> {
+	val handler = getCapability(MatteryCapability.CURIOS_INVENTORY).orNull() ?: return Stream.empty()
+
+	val result = LinkedList<Stream<out AwareItemStack>>()
+
+	for ((identifier, curio) in handler.curios) {
+		result.add(curio.stacks.awareStream())
+
+		if (includeCosmetics && curio.hasCosmetic()) {
+			result.add(curio.cosmeticStacks.awareStream())
+		}
+	}
+
+	return Streams.concat(*result.toTypedArray())
+}
+
+fun Player.curiosAwareStream(includeCosmetics: Boolean = true): Stream<out AwareItemStack> {
+	if (!isCuriosLoaded) {
+		return Stream.empty()
+	}
+
+	return curiosAwareStreamImpl(includeCosmetics)
+}

src/main/kotlin/ru/dbotthepony/mc/otm/container/ContainerIterator.kt

@@ -1,30 +1,58 @@
 package ru.dbotthepony.mc.otm.container
 
+import it.unimi.dsi.fastutil.objects.ObjectIterators
 import net.minecraft.world.Container
 import net.minecraft.world.item.ItemStack
+import ru.dbotthepony.mc.otm.core.ContainerItemStackEntry
+import ru.dbotthepony.mc.otm.core.IAwareItemStackIterator
+import ru.dbotthepony.mc.otm.core.AwareItemStack
 
-class ContainerIterator(private val container: Container) : MutableIterator<ItemStack> {
-	private var index = 0
-
-	override fun hasNext(): Boolean {
-		return index < container.containerSize
+class ContainerIterator(private val container: Container, initialPosition: Int = 0) : ObjectIterators.AbstractIndexBasedListIterator<ItemStack>(0, initialPosition), MutableListIterator<ItemStack> {
+	init {
+		require(initialPosition in 0 .. container.containerSize) { "Invalid initial position: $initialPosition" }
 	}
 
-	override fun next(): ItemStack {
-		if (index >= container.containerSize) {
-			throw IllegalStateException("Already finished iterating")
-		}
-
-		return container[index++]
+	override fun add(location: Int, k: ItemStack) {
+		throw UnsupportedOperationException()
 	}
 
-	override fun remove() {
-		if (index == 0) {
-			throw IllegalStateException("Never called next()")
-		}
+	override fun set(location: Int, k: ItemStack) {
+		container[location] = k
+	}
 
-		container[index - 1] = ItemStack.EMPTY
+	override fun remove(location: Int) {
+		container[location] = ItemStack.EMPTY
+	}
+
+	override fun get(location: Int): ItemStack {
+		return container[location]
+	}
+
+	override fun getMaxPos(): Int {
+		return container.containerSize
 	}
 }
 
-fun Container.iterator() = ContainerIterator(this)
+class AwareContainerIterator(private val container: Container, initialPosition: Int = 0) : ObjectIterators.AbstractIndexBasedIterator<AwareItemStack>(0, initialPosition), IAwareItemStackIterator {
+	init {
+		require(initialPosition in 0 .. container.containerSize) { "Invalid initial position: $initialPosition" }
+	}
+
+	override fun remove(location: Int) {
+		container[location] = ItemStack.EMPTY
+	}
+
+	override fun get(location: Int): AwareItemStack {
+		return ContainerItemStackEntry(location, container)
+	}
+
+	override fun getMaxPos(): Int {
+		return container.containerSize
+	}
+}
+
+@JvmOverloads
+fun Container.iterator(initialPosition: Int = 0) = ContainerIterator(this, initialPosition)
+
+@JvmOverloads
+fun Container.awareIterator(initialPosition: Int = 0) = AwareContainerIterator(this, initialPosition)

src/main/kotlin/ru/dbotthepony/mc/otm/container/ContainerSpliterator.kt

@@ -4,6 +4,8 @@ import it.unimi.dsi.fastutil.objects.ObjectSpliterator
 import it.unimi.dsi.fastutil.objects.ObjectSpliterators
 import net.minecraft.world.Container
 import net.minecraft.world.item.ItemStack
+import ru.dbotthepony.mc.otm.core.AwareItemStack
+import ru.dbotthepony.mc.otm.core.ContainerItemStackEntry
 import java.util.stream.Stream
 import java.util.stream.StreamSupport
 
@@ -26,5 +28,26 @@ class ContainerSpliterator(private val container: Container, offset: Int = 0, pr
 	}
 }
 
+class AwareContainerSpliterator(private val container: Container, offset: Int = 0, private val maxPos: Int = container.containerSize) : ObjectSpliterators.AbstractIndexBasedSpliterator<AwareItemStack>(offset) {
+	init {
+		require(offset >= 0) { "Invalid offset $offset" }
+		require(offset + maxPos <= container.containerSize) { "$offset -> $maxPos while having only size of ${container.containerSize}!" }
+	}
+
+	override fun get(location: Int): AwareItemStack {
+		return ContainerItemStackEntry(location, container)
+	}
+
+	override fun getMaxPos(): Int {
+		return maxPos
+	}
+
+	override fun makeForSplit(pos: Int, maxPos: Int): ObjectSpliterator<AwareItemStack> {
+		return AwareContainerSpliterator(container, pos, maxPos)
+	}
+}
+
 fun Container.spliterator() = ContainerSpliterator(this)
+fun Container.awareSpliterator() = AwareContainerSpliterator(this)
 fun Container.stream(): Stream<out ItemStack> = StreamSupport.stream(spliterator(), false)
+fun Container.awareStream(): Stream<out AwareItemStack> = StreamSupport.stream(awareSpliterator(), false)

src/main/kotlin/ru/dbotthepony/mc/otm/container/ItemHandlerSpliterator.kt

@@ -4,6 +4,8 @@ import it.unimi.dsi.fastutil.objects.ObjectSpliterator
 import it.unimi.dsi.fastutil.objects.ObjectSpliterators
 import net.minecraft.world.item.ItemStack
 import net.minecraftforge.items.IItemHandler
+import ru.dbotthepony.mc.otm.core.AwareItemStack
+import ru.dbotthepony.mc.otm.core.ItemHandlerItemStackEntry
 import java.util.Spliterator
 import java.util.stream.Stream
 import java.util.stream.StreamSupport
@@ -27,5 +29,26 @@ class ItemHandlerSpliterator(private val handler: IItemHandler, offset: Int = 0,
 	}
 }
 
+class ItemHandlerAwareSpliterator(private val handler: IItemHandler, offset: Int = 0, private val maxPos: Int = handler.slots) : ObjectSpliterators.AbstractIndexBasedSpliterator<AwareItemStack>(offset) {
+	init {
+		require(offset >= 0) { "Invalid offset $offset" }
+		require(offset + maxPos <= handler.slots) { "$offset -> $maxPos while having only size of ${handler.slots}!" }
+	}
+
+	override fun get(location: Int): AwareItemStack {
+		return ItemHandlerItemStackEntry(location, handler)
+	}
+
+	override fun getMaxPos(): Int {
+		return maxPos
+	}
+
+	override fun makeForSplit(pos: Int, maxPos: Int): ObjectSpliterator<AwareItemStack> {
+		return ItemHandlerAwareSpliterator(handler, pos, maxPos)
+	}
+}
+
 fun IItemHandler.spliterator(): Spliterator<out ItemStack> = ItemHandlerSpliterator(this)
+fun IItemHandler.awareSpliterator(): Spliterator<out AwareItemStack> = ItemHandlerAwareSpliterator(this)
 fun IItemHandler.stream(): Stream<out ItemStack> = StreamSupport.stream(spliterator(), false)
+fun IItemHandler.awareStream(): Stream<out AwareItemStack> = StreamSupport.stream(awareSpliterator(), false)

src/main/kotlin/ru/dbotthepony/mc/otm/core/AwareItemStack.kt

@@ -0,0 +1,56 @@
+package ru.dbotthepony.mc.otm.core
+
+import net.minecraft.world.Container
+import net.minecraft.world.item.ItemStack
+import net.minecraftforge.items.IItemHandler
+import ru.dbotthepony.mc.otm.container.get
+import ru.dbotthepony.mc.otm.container.set
+
+/**
+ * Allows to see the contents of container and extract the item from it
+ */
+interface AwareItemStack {
+	operator fun component0(): ItemStack
+	val itemStack: ItemStack get() = component0()
+
+	/**
+	 * Extracts item from underlying container
+	 */
+	fun extract(amount: Int, simulate: Boolean = true): ItemStack
+}
+
+typealias IAwareItemStackIterator = Iterator<AwareItemStack>
+
+data class ContainerItemStackEntry(val index: Int, val container: Container) : AwareItemStack {
+	override fun component0(): ItemStack {
+		return container[index]
+	}
+
+	override fun extract(amount: Int, simulate: Boolean): ItemStack {
+		if (simulate) {
+			return container[index].let { if (it.isEmpty) it else it.copy().also { it.count = it.count.coerceAtMost(amount) } }
+		}
+
+		val item = container[index]
+		val newCount = (item.count - amount).coerceAtLeast(0)
+
+		if (newCount == 0) {
+			container[index] = ItemStack.EMPTY
+			return item
+		}
+
+		val copy = item.copy().also { it.count = amount }
+		container[index] = item.copy().also { it.count -= amount }
+		return copy
+	}
+}
+
+data class ItemHandlerItemStackEntry(val index: Int, val handler: IItemHandler) : AwareItemStack {
+	override fun component0(): ItemStack {
+		return handler[index]
+	}
+
+	override fun extract(amount: Int, simulate: Boolean): ItemStack {
+		return handler.extractItem(index, amount, simulate)
+	}
+}

src/main/kotlin/ru/dbotthepony/mc/otm/core/EmptyMutableIterator.kt

@@ -0,0 +1,44 @@
+package ru.dbotthepony.mc.otm.core
+
+object EmptyMutableIterator : MutableIterator<@UnsafeVariance Nothing>, MutableListIterator<@UnsafeVariance Nothing> {
+	override fun hasPrevious(): Boolean {
+		return false
+	}
+
+	override fun nextIndex(): Int {
+		return 0
+	}
+
+	override fun previous(): Nothing {
+		throw NoSuchElementException()
+	}
+
+	override fun previousIndex(): Int {
+		return -1
+	}
+
+	override fun add(element: Nothing) {
+		throw UnsupportedOperationException()
+	}
+
+	override fun hasNext(): Boolean {
+		return false
+	}
+
+	override fun next(): Nothing {
+		throw NoSuchElementException()
+	}
+
+	override fun remove() {
+		throw NoSuchElementException()
+	}
+
+	override fun set(element: Nothing) {
+		throw UnsupportedOperationException()
+	}
+
+	@Suppress("unchecked_cast")
+	fun <T> cast(): MutableListIterator<T> {
+		return this as MutableListIterator<T>
+	}
+}

src/main/kotlin/ru/dbotthepony/mc/otm/core/ItemStackIterators.kt

@@ -2,50 +2,5 @@ package ru.dbotthepony.mc.otm.core
 
 import net.minecraft.world.item.ItemStack
 
-open class NonEmptyItemStackIterator(protected open val parent: Iterator<ItemStack>) : Iterator<ItemStack> {
-	private var itemStack: ItemStack? = null
-	private var searched = false
-
-	private fun search() {
-		searched = true
-
-		if (itemStack != null) {
-			return
-		}
-
-		for (stack in parent) {
-			if (!stack.isEmpty) {
-				itemStack = stack
-				return
-			}
-		}
-	}
-
-	override fun hasNext(): Boolean {
-		if (!searched) {
-			search()
-		}
-
-		return itemStack != null
-	}
-
-	override fun next(): ItemStack {
-		if (!searched) {
-			search()
-		}
-
-		val itemStack = itemStack ?: throw IllegalStateException("No next element")
-		this.itemStack = null
-		this.searched = false
-		return itemStack
-	}
-}
-
-class NonEmptyMutableItemStackIterator(override val parent: MutableIterator<ItemStack>) : NonEmptyItemStackIterator(parent), MutableIterator<ItemStack> {
-	override fun remove() {
-		parent.remove()
-	}
-}
-
-fun Iterator<ItemStack>.nonEmpty() = NonEmptyItemStackIterator(this)
-fun MutableIterator<ItemStack>.nonEmpty() = NonEmptyMutableItemStackIterator(this)
+fun Iterator<ItemStack>.nonEmpty() = PredicateIterator(this) { !it.isEmpty }
+fun MutableIterator<ItemStack>.nonEmpty() = MutablePredicateIterator(this) { !it.isEmpty }

src/main/kotlin/ru/dbotthepony/mc/otm/core/PredicateIterator.kt

@@ -0,0 +1,91 @@
+package ru.dbotthepony.mc.otm.core
+
+import java.util.function.Consumer
+import java.util.function.Predicate
+
+class PredicateIterator<T> : Iterator<T> {
+	private val parent: Iterator<T>
+	private val predicate: Predicate<in T>
+	private val consumer: Consumer<in T>?
+
+	constructor(parent: Iterator<T>, predicate: Predicate<in T>) {
+		this.parent = parent
+		this.predicate = predicate
+		this.consumer = null
+	}
+
+	constructor(parent: Iterator<T>, predicate: Predicate<in T>, consumer: Consumer<in T>) {
+		this.parent = parent
+		this.predicate = predicate
+		this.consumer = consumer
+	}
+
+	private var foundValue: Any? = Companion
+
+	override fun hasNext(): Boolean {
+		if (foundValue === Companion) {
+			while (parent.hasNext()) {
+				val next = parent.next()
+
+				if (predicate.test(next)) {
+					foundValue = next
+					return true
+				}
+			}
+
+			return false
+		}
+
+		return true
+	}
+
+	@Suppress("unchecked_cast")
+	override fun next(): T {
+		if (!hasNext()) {
+			throw NoSuchElementException()
+		}
+
+		val foundValue = foundValue
+
+		if (foundValue === Companion) {
+			throw ConcurrentModificationException()
+		}
+
+		this.foundValue = Companion
+		consumer?.accept(foundValue as T)
+		return foundValue as T
+	}
+
+	private companion object
+}
+
+fun <T> Iterator<T>.filter(condition: Predicate<in T>) = PredicateIterator(this, condition)
+
+class MutablePredicateIterator<T> : MutableIterator<T> {
+	private val parent: MutableIterator<T>
+	private val predicateParent: PredicateIterator<T>
+
+	constructor(parent: MutableIterator<T>, predicate: Predicate<in T>) {
+		this.parent = parent
+		this.predicateParent = PredicateIterator(parent, predicate)
+	}
+
+	constructor(parent: MutableIterator<T>, predicate: Predicate<in T>, consumer: Consumer<in T>) {
+		this.parent = parent
+		this.predicateParent = PredicateIterator(parent, predicate, consumer)
+	}
+
+	override fun hasNext(): Boolean {
+		return predicateParent.hasNext()
+	}
+
+	override fun next(): T {
+		return predicateParent.next()
+	}
+
+	override fun remove() {
+		return parent.remove()
+	}
+}
+
+fun <T> MutableIterator<T>.filter(condition: Predicate<in T>) = MutablePredicateIterator(this, condition)