listing.kt

package chapter8.sec3.listing3

import arrow.core.getOrElse
import arrow.core.toOption
import chapter8.Falsified
import chapter8.MaxSize
import chapter8.Passed
import chapter8.RNG
import chapter8.Result
import chapter8.SimpleRNG
import chapter8.double
import chapter8.nonNegativeInt
import kotlin.math.absoluteValue
import kotlin.math.min

typealias State<S, A> = (S) -> Pair<A, S>
typealias Gen<A> = State<RNG, A>
typealias SGen<A> = (Int) -> Gen<A>
typealias TestCases = Int
typealias Prop = (MaxSize, TestCases, RNG) -> Result
private typealias Prop1 = (RNG) -> Result
private typealias Prop2 = (TestCases, RNG) -> Result

fun <S, A, B> State<S, A>.map(f: (A) -> B): State<S, B> =
    flatMap { a -> unit<S, B>(f(a)) }

fun <S, A, B> State<S, A>.flatMap(f: (A) -> State<S, B>): State<S, B> =
    { s: S ->
        val (a: A, s2: S) = this(s)
        f(a)(s2)
    }

fun <S, A> unit(a: A): State<S, A> =
    { s: S -> Pair(a, s) }

fun <S, A, B, C> map2(
    ra: State<S, A>,
    rb: State<S, B>,
    f: (A, B) -> C
): State<S, C> =
    ra.flatMap { a ->
        rb.map { b ->
            f(a, b)
        }
    }

fun <S, A> List<State<S, A>>.sequence(): State<S, List<A>> =
    foldRight(unit(emptyList())) { f, acc ->
        map2(f, acc) { h, t -> listOf(h) + t }
    }

fun choose(start: Int, stopExclusive: Int): Gen<Int> =
    { rng: RNG -> double(rng) }
        .map { start + (it * (stopExclusive - start)).toInt() }

fun <A> Gen<A>.listOfN(n: Int): Gen<List<A>> =
    List(n) { this }.sequence()

fun <A> weighted(
    pga: Pair<Gen<A>, Double>,
    pgb: Pair<Gen<A>, Double>
): Gen<A> {
    val (ga, p1) = pga
    val (gb, p2) = pgb
    val prob =
        p1.absoluteValue /
            (p1.absoluteValue + p2.absoluteValue)
    return { rng: RNG -> double(rng) }
        .flatMap { d ->
            if (d < prob) ga else gb
        }
}

fun <A> Gen<A>.sGen(): SGen<List<A>> =
    { listOfN(it) }

//tag::init[]

fun <A> SGen<A>.forAll(f: (A) -> Boolean): Prop =
    { max, n, rng ->
        val casePerSize: Int = (n + (max - 1)) / max // <2>
        val props: Sequence<Prop2> =
            generateSequence(0) { it + 1 } // <3>
                .take(min(n, max) + 1)
                .map { i ->
                    invoke(i).forAll(f)
                } // <4>

        val prop: Prop1 = props.map { p ->
            { rng: RNG ->
                p(casePerSize, rng)
            }
        }.reduce { p1, p2 -> p1.and(p2) } // <5>

        prop(rng) // <6>
    }

//tag::ignore[]
fun <A> Gen<A>.forAll(f: (A) -> Boolean): Prop2 =
    { n: Int, rng: RNG ->
        randomSequence(rng).mapIndexed { i, a ->
            try {
                if (f(a)) Passed else Falsified(
                    a.toString(),
                    i
                )
            } catch (e: Exception) {
                Falsified(buildMessage(a, e), i)
            }
        }.take(n)
            .find { it.isFalsified() }
            .toOption()
            .getOrElse { Passed }
    }

fun <A> Gen<A>.randomSequence(
    rng: RNG
): Sequence<A> =
    sequence {
        val (a: A, rng2: RNG) = invoke(rng)
        yield(a)
        yieldAll(randomSequence(rng2))
    }

private fun <A> buildMessage(a: A, e: Exception) =
    """
    |test case: $a
    |generated an exception: ${e.message}
    |stacktrace:
    |${e.stackTrace.take(10).joinToString("\n")}
""".trimMargin()
//end::ignore[]

fun Prop1.and(p: Prop1): Prop1 =
    { rng -> // <7>
        when (val result: Result = this(rng)) {
            is Passed -> p(rng)
            is Falsified -> result
        }
    }
//end::init[]

// fun main() {
//     val p = { rng: RNG -> nonNegativeInt(rng) }.sGen()
//         .forAll { it.all { i -> i >= 0 } }
//     val r = p(3, 5, SimpleRNG(0))
//     println(r)
// }