package moe.kageru.sekwences

/**
 * Returns a sequences that optionally merges adjacent elements.
 * Heavily inspired by Itertools.coalesce() from the Rust crate.
 *
 * The merge function will be passed two adjacent elements in the sequence.
 * If the two can be merged, the result of the merging process should be returned.
 * If null is returned, the first of the two elements is returned,
 * and the second will be used as the first element in the next merge operation.
 *
 * https://docs.rs/itertools/0.8.0/itertools/trait.Itertools.html#method.coalesce
 */
fun <T> Sequence<T>.coalesce(merger: (T, T) -> T?): Sequence<T> {
    return if (this.iterator().hasNext()) {
        CoalescingSequence(this, merger)
    } else {
        emptySequence()
    }
}

/**
 * Sequence that holds a [CoalescingIterator].
 */
internal class CoalescingSequence<T>(private val source: Sequence<T>, private inline val merger: (T, T) -> T?) : Sequence<T> {
    override fun iterator(): Iterator<T> {
        return CoalescingIterator(source.iterator(), merger)
    }
}

/**
 * Iterator that will attempt to merge adjacent elements as described in the KDoc for [coalesce].
 */
internal class CoalescingIterator<T>(private val source: Iterator<T>, private inline val merger: (T, T) -> T?) : Iterator<T> {
    private var previous: T = if (source.hasNext()) source.next() else error("Please don’t pass empty iterators in here")
    // The reason we need this marker is that our iterator can still hold one value, even if the source has been drained.
    private var hasNext = true

    override fun hasNext() = hasNext

    override tailrec fun next(): T {
        if (!source.hasNext()) {
            hasNext = false
            return previous
        }

        val current = source.next()
        val merged = merger(previous, current)
        // If the elements can’t be merged, return the first, then save the second for next time.
        return if (merged == null) {
            val ret = previous
            previous = current
            ret
        } else {
            previous = merged
            next()
        }
    }
}