From 218bb17cc456772cf28cf15700ee32c7e8e819f7 Mon Sep 17 00:00:00 2001
From: DBotThePony <dbotthepony@yandex.ru>
Date: Tue, 22 Aug 2023 11:35:00 +0700
Subject: [PATCH] Faster filtering and flatmapping iterators

---
 .../mc/otm/core/collect/StreamyIterators.kt   | 189 +++++++++---------
 1 file changed, 89 insertions(+), 100 deletions(-)

diff --git a/src/main/kotlin/ru/dbotthepony/mc/otm/core/collect/StreamyIterators.kt b/src/main/kotlin/ru/dbotthepony/mc/otm/core/collect/StreamyIterators.kt
index 5377de346..3a925efff 100644
--- a/src/main/kotlin/ru/dbotthepony/mc/otm/core/collect/StreamyIterators.kt
+++ b/src/main/kotlin/ru/dbotthepony/mc/otm/core/collect/StreamyIterators.kt
@@ -1,11 +1,15 @@
 package ru.dbotthepony.mc.otm.core.collect
 
+import it.unimi.dsi.fastutil.ints.Int2IntFunction
+import it.unimi.dsi.fastutil.ints.IntIterator
+import it.unimi.dsi.fastutil.ints.IntIterators
 import it.unimi.dsi.fastutil.objects.ObjectIterators
 import ru.dbotthepony.mc.otm.core.addAll
 import java.util.Optional
 import java.util.Spliterator
 import java.util.Spliterators
 import java.util.function.BinaryOperator
+import java.util.function.IntPredicate
 import java.util.function.Predicate
 import java.util.function.Supplier
 import java.util.stream.Collector
@@ -21,68 +25,43 @@ import java.util.stream.StreamSupport
 
 // Aside parallel work, unimplemented Stream API elements can be easily implemented when required
 
-/**
- * Filters elements of [parent] iterator
- *
- * Resulting [Iterator] is [MutableIterator] if [parent] is
- */
-class FilteringIterator<T>(private val parent: Iterator<T>, private val predicate: Predicate<in T>) : MutableIterator<T> {
-	private var foundValue: Any? = Companion
+private class FilteringIterator<T>(private val parent: Iterator<T>, private val predicate: Predicate<in T>) : MutableIterator<T> {
+	private var value: T = parent.next()
+	private var hasValue = true
 	private var returned = false
 
 	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
+		return hasValue
 	}
 
-	@Suppress("unchecked_cast")
 	override fun next(): T {
-		var foundValue = foundValue
+		if (!hasValue) throw NoSuchElementException()
+		hasValue = false
+		returned = true
 
-		if (foundValue === Companion) {
-			if (!hasNext()) {
-				throw NoSuchElementException()
-			}
+		val value = this.value
 
-			foundValue = this.foundValue
+		while (parent.hasNext()) {
+			val next = parent.next()
 
-			if (foundValue === Companion) {
-				throw ConcurrentModificationException()
+			if (predicate.test(next)) {
+				hasValue = true
+				this.value = next
+				break
 			}
 		}
 
-		this.foundValue = Companion
-		returned = true
-		return foundValue as T
+		return value
 	}
 
 	override fun remove() {
 		if (!returned) throw NoSuchElementException()
-		(parent as MutableIterator<T>).remove()
 		returned = false
+		(parent as MutableIterator<T>).remove()
 	}
-
-	private companion object
 }
 
-/**
- * Maps elements of [parent] iterator from values of [T] to [R] using function [mapper]
- *
- * Resulting [Iterator] is [MutableIterator] if [parent] is
- */
-class MappingIterator<T, R>(private val parent: Iterator<T>, private val mapper: (T) -> R) : MutableIterator<R> {
+private class MappingIterator<T, R>(private val parent: Iterator<T>, private val mapper: (T) -> R) : MutableIterator<R> {
 	override fun hasNext(): Boolean {
 		return parent.hasNext()
 	}
@@ -96,29 +75,28 @@ class MappingIterator<T, R>(private val parent: Iterator<T>, private val mapper:
 	}
 }
 
-/**
- * Maps elements of [parent] iterator from type [T] to other iterators of type [R] using function [mapper]
- *
- * Resulting [Iterator] is [MutableIterator] if [parent] is
- */
-class FlatMappingIterator<T, R>(private val parent: Iterator<T>, private val mapper: (T) -> Iterator<R>) : MutableIterator<R> {
-	private var current: Iterator<R>? = null
+private class FlatMappingIterator<T, R>(private val parent: Iterator<T>, private val mapper: (T) -> Iterator<R>) : MutableIterator<R> {
+	private var current: Iterator<R> = mapper.invoke(parent.next())
 	private var last: Iterator<R>? = null
+	private var hasNext = current.hasNext()
 
 	override fun hasNext(): Boolean {
-		while (current?.hasNext() != true && parent.hasNext()) {
-			current = mapper.invoke(parent.next())
-		}
-
-		return current?.hasNext() == true
+		return hasNext
 	}
 
 	override fun next(): R {
-		if (!hasNext())
+		if (!hasNext)
 			throw NoSuchElementException()
 
-		val v = current!!.next()
+		val v = current.next()
 		last = current
+		hasNext = false
+
+		while (!hasNext && parent.hasNext()) {
+			current = mapper.invoke(parent.next())
+			hasNext = current.hasNext()
+		}
+
 		return v
 	}
 
@@ -128,12 +106,7 @@ class FlatMappingIterator<T, R>(private val parent: Iterator<T>, private val map
 	}
 }
 
-/**
- * Limits amount of values returned by [parent] iterator to return at most [limit] values
- *
- * Resulting [Iterator] is [MutableIterator] if [parent] is
- */
-class LimitingIterator<T>(private val parent: Iterator<T>, private val limit: Long) : Iterator<T> {
+private class LimitingIterator<T>(private val parent: Iterator<T>, private val limit: Long) : Iterator<T> {
 	init {
 		require(limit > 0) { "Invalid limit $limit" }
 	}
@@ -152,12 +125,7 @@ class LimitingIterator<T>(private val parent: Iterator<T>, private val limit: Lo
 	}
 }
 
-/**
- * Skips (discards) up to [skip] values returned by [parent] iterator
- *
- * Resulting [Iterator] is [MutableIterator] if [parent] is
- */
-class SkippingIterator<T>(private val parent: Iterator<T>, skip: Long) : MutableIterator<T> {
+private class SkippingIterator<T>(private val parent: Iterator<T>, skip: Long) : MutableIterator<T> {
 	init {
 		require(skip >= 0) { "Invalid skip amount $skip" }
 	}
@@ -214,71 +182,92 @@ fun <T> concatIterators(vararg iterators: Iterator<T>): MutableIterator<T> {
  *
  * Resulting [Iterator] is [MutableIterator] if [this] is
  */
-fun <T> Iterator<T>.filter(condition: Predicate<in T>) = FilteringIterator(this, condition)
+fun <T> Iterator<T>.filter(condition: Predicate<in T>): MutableIterator<T> {
+	if (!hasNext()) {
+		return emptyIterator()
+	}
+
+	return FilteringIterator(this, condition)
+}
 
 /**
  * Maps elements of [this] iterator from values of [T] to [R] using function [mapper]
  *
  * Resulting [Iterator] is [MutableIterator] if [this] is
  */
-fun <T, R> Iterator<T>.map(mapper: (T) -> R) = MappingIterator(this, mapper)
+fun <T, R> Iterator<T>.map(mapper: (T) -> R): MutableIterator<R> {
+	if (!hasNext()) {
+		return emptyIterator()
+	}
 
-/**
- * Maps elements of [this] iterator from values of [T] to [R] using function [mapper]
- *
- * Resulting [Iterator] is [MutableIterator] if [this] is
- */
-fun <T, R> Iterator<T>.map(mapper: java.util.function.Function<T, R>) = MappingIterator(this, mapper::apply)
+	return MappingIterator(this, mapper)
+}
 
 /**
  * Maps elements of [this] iterator from type [T] to other iterators of type [R] using function [mapper]
  *
  * Resulting [Iterator] is [MutableIterator] if [this] is
  */
-fun <T, R> Iterator<T>.flatMap(mapper: (T) -> Iterator<R>) = FlatMappingIterator(this, mapper)
+fun <T, R> Iterator<T>.flatMap(mapper: (T) -> Iterator<R>): MutableIterator<R> {
+	if (!hasNext()) {
+		return emptyIterator()
+	}
 
-/**
- * Maps elements of [this] iterator from type [T] to other iterators of type [R] using function [mapper]
- *
- * Resulting [Iterator] is [MutableIterator] if [this] is
- */
-fun <T, R> Iterator<T>.flatMap(mapper: java.util.function.Function<T, Iterator<R>>) = FlatMappingIterator(this, mapper::apply)
+	return FlatMappingIterator(this, mapper)
+}
 
-fun <T> Iterator<T>.reduce(identity: T, reducer: (T, T) -> T): T {
+inline fun <T> Iterator<T>.reduce(identity: T, reducer: (T, T) -> T): T {
 	var result = identity
-	for (value in this) result = reducer.invoke(result, value)
+	while (hasNext()) result = reducer.invoke(result, next())
 	return result
 }
 
-fun <T> Iterator<T>.reduce(identity: T, reducer: BinaryOperator<T>): T = reduce(identity, reducer::apply)
 fun <T> Iterator<T?>.filterNotNull(): MutableIterator<T> = filter { it != null } as MutableIterator<T>
 
 fun <T> Iterator<T>.any(predicate: Predicate<in T>): Boolean {
-	for (value in this) {
-		if (predicate.test(value)) {
+	while (hasNext())
+		if (predicate.test(next()))
+			return true
+
+	return false
+}
+
+inline fun <T> Iterator<T>.any(predicate: (T) -> Boolean): Boolean {
+	while (hasNext())
+		if (predicate.invoke(next()))
 			return true
-		}
-	}
 
 	return false
 }
 
 fun <T> Iterator<T>.all(predicate: Predicate<in T>): Boolean {
-	for (value in this) {
-		if (!predicate.test(value)) {
+	while (hasNext())
+		if (!predicate.test(next()))
+			return false
+
+	return true
+}
+
+inline fun <T> Iterator<T>.all(predicate: (T) -> Boolean): Boolean {
+	while (hasNext())
+		if (!predicate.invoke(next()))
 			return false
-		}
-	}
 
 	return true
 }
 
 fun <T> Iterator<T>.none(predicate: Predicate<in T>): Boolean {
-	for (value in this) {
-		if (predicate.test(value)) {
+	while (hasNext())
+		if (predicate.test(next()))
+			return false
+
+	return true
+}
+
+inline fun <T> Iterator<T>.none(predicate: (T) -> Boolean): Boolean {
+	while (hasNext())
+		if (predicate.invoke(next()))
 			return false
-		}
-	}
 
 	return true
 }
@@ -308,8 +297,8 @@ fun <T : Any> Iterator<T>.find(): Optional<T> {
 	return Optional.empty<T>()
 }
 
-fun <T> Iterator<T>.limit(limit: Long) = LimitingIterator(this, limit)
-fun <T> Iterator<T>.skip(skip: Long) = if (skip == 0L) this else SkippingIterator(this, skip)
+fun <T> Iterator<T>.limit(limit: Long): Iterator<T> = LimitingIterator(this, limit)
+fun <T> Iterator<T>.skip(skip: Long): Iterator<T> = if (skip == 0L) this else SkippingIterator(this, skip)
 
 inline fun <T> Iterator<T>.forEach(action: (T) -> Unit) {
 	for (value in this) {