Sequences.kt

package kotlinx.benchmarks.json

import kotlinx.benchmarks.model.*
import org.openjdk.jmh.annotations.*
import org.openjdk.jmh.infra.Blackhole
import java.util.concurrent.*
import kotlin.random.Random
import org.openjdk.jmh.annotations.Benchmark
import org.openjdk.jmh.annotations.Fork
import org.openjdk.jmh.annotations.Measurement
import org.openjdk.jmh.annotations.Scope
import org.openjdk.jmh.annotations.State
import org.openjdk.jmh.annotations.Warmup
import java.util.concurrent.TimeUnit

@State(Scope.Benchmark)
@Fork(1)
@Warmup(iterations = 5, time = 1, timeUnit = TimeUnit.SECONDS)
@Measurement(iterations = 5, time = 1, timeUnit = TimeUnit.SECONDS)
open class SequenceBenchmark {
    @Benchmark
    fun list() =
        DataRepository.getItemsList()


    @Benchmark
    fun sequence() =
        DataRepository.getItemsListUsingSequence()

    @Benchmark
    fun pushBased() =
        DataRepository.getItemsListUsingShortCirtuitSequence()
}

private const val ITEM_SIZE = 100

object Db {
    fun getItems(): List<DbModel> {
        return (0 until ITEM_SIZE).map { index ->
            DbModel(
                id = index,
                rank = ITEM_SIZE - index, // reverse
                isEnabled = index.mod(2) == 0,
            )
        }
    }
}

object DataRepository {
    fun getItemsList(): List<UiModel> {
        return Db.getItems()
            .filter { it.isEnabled }
            .map { UiModel(it.id) }
            .filter { true }
            .map { UiModel(it.id) }
            .filter { true }
            .map { UiModel(it.id) }
            .filter { true }
            .map { UiModel(it.id) }
    }

    fun getItemsListUsingSequence(): List<UiModel> {
        return Db.getItems()
            .asSequence()
            .filter { it.isEnabled }
            .map { UiModel(it.id) }
            .filter { true }
            .map { UiModel(it.id) }
            .filter { true }
            .map { UiModel(it.id) }
            .filter { true }
            .map { UiModel(it.id) }
            .toList()
    }

    fun getItemsListUsingShortCirtuitSequence(): List<UiModel> {
        return Db.getItems()
            .asShortCirtuitSequence()
            .filter { it.isEnabled }
            .map { UiModel(it.id) }
            .filter { true }
            .map { UiModel(it.id) }
            .filter { true }
            .map { UiModel(it.id) }
            .filter { true }
            .map { UiModel(it.id) }
            .toList()
    }
}



interface ShortCirtuitSequence<T> : Sequence<T> {
    fun consume(block: (T) -> Boolean) {
        iterator().forEach { if (!block(it)) return }
    }
}

fun <T> List<T>.asShortCirtuitSequence(): ShortCirtuitSequence<T> {
    return object : ShortCirtuitSequence<T> {
        override fun iterator(): Iterator<T> {
            TODO()
        }

        override fun consume(block: (T) -> Boolean) {
           this@asShortCirtuitSequence.forEach { if (!block(it)) return }
        }
    }
}

fun <T> ShortCirtuitSequence<T>.toList(): List<T> {
    val result = mutableListOf<T>()
    consume { result.add(it); true }
    return result
}

fun <T, R> ShortCirtuitSequence<T>.map(transform: (T) -> R): ShortCirtuitSequence<R> {
    val source = this
    return object : ShortCirtuitSequence<R> {
        override fun iterator(): Iterator<R> {
            TODO() // old impl
        }

        override fun consume(block: (R) -> Boolean) {
            source.consume {
                block(transform(it) )
            }
        }
    }
}

fun <T> ShortCirtuitSequence<T>.filter(predicate: (T) -> Boolean): ShortCirtuitSequence<T> {
    val source = this
    return object : ShortCirtuitSequence<T> {
        override fun iterator(): Iterator<T> {
            TODO() // old impl
        }

        override fun consume(block: (T) -> Boolean) {
            source.consume {
                if (predicate(it)) {
                    block(it)
                } else {
                    true
                }
            }
        }
    }
}


data class DbModel(val id: Int, val rank: Int, val isEnabled: Boolean)

data class UiModel(val id: Int)