Providers.kt

package org.utbot.fuzzing

import mu.KotlinLogging
import kotlin.random.Random

private val logger by lazy { KotlinLogging.logger {} }

/**
 * Entry point to run fuzzing.
 */
suspend fun <TYPE, RESULT, DESCRIPTION : Description<TYPE>, FEEDBACK : Feedback<TYPE, RESULT>> runFuzzing(
    provider: ValueProvider<TYPE, RESULT, DESCRIPTION>,
    description: DESCRIPTION,
    random: Random = Random(0),
    configuration: Configuration = Configuration(),
    handle: suspend (description: DESCRIPTION, values: List<RESULT>) -> FEEDBACK
) {
    BaseFuzzing(listOf(provider), handle).fuzz(description, random, configuration)
}

/**
 * Implements base concepts that use providers to generate values for some types.
 *
 * @param providers is a list of "type to values" generator
 * @param exec this function is called when fuzzer generates values of type R to run it with target program.
 */
class BaseFuzzing<T, R, D : Description<T>, F : Feedback<T, R>>(
    val providers: List<ValueProvider<T, R, D>>,
    val exec: suspend (description: D, values: List<R>) -> F
) : Fuzzing<T, R, D, F> {

    constructor(vararg providers: ValueProvider<T, R, D>, exec: suspend (description: D, values: List<R>) -> F) : this(providers.toList(), exec)

    override fun enrich(description: D, type: T, scope: Scope) {
        providers.asSequence().forEach {
            it.enrich(description, type, scope)
        }
    }

    override fun generate(description: D, type: T): Sequence<Seed<T, R>> {
        return providers.asSequence().flatMap { provider ->
            try {
                if (provider.accept(type)) {
                    provider.generate(description, type)
                } else {
                    emptySequence()
                }
            } catch (t: Throwable) {
                logger.error(t) { "Error occurs in value provider: $provider" }
                emptySequence()
            }
        }
    }

    override suspend fun handle(description: D, values: List<R>): F {
        return exec(description, values)
    }
}

/**
 * Value provider generates [Seed] and has other methods to combine providers.
 */
fun interface ValueProvider<T, R, D : Description<T>> {

    fun enrich(description: D, type: T, scope: Scope) {}

    /**
     * Generate a sequence of [Seed] that is merged with values generated by other provider.
     */
    fun generate(description: D, type: T): Sequence<Seed<T, R>>

    /**
     * Validates if this provider is applicable to some type.
     */
    fun accept(type: T): Boolean = true

    /**
     * Combines this model provider with `anotherValueProviders` into one instance.
     *
     * This model provider is called before `anotherValueProviders`.
     */
    infix fun with(anotherValueProvider: ValueProvider<T, R, D>): ValueProvider<T, R, D> {
        fun toList(m: ValueProvider<T, R, D>) = if (m is Combined<T, R, D>) m.providers else listOf(m)
        return Combined(toList(this) + toList(anotherValueProvider))
    }

    /**
     * Removes `anotherValueProviders<T, R` from current one.
     */
    fun except(anotherValueProvider: ValueProvider<T, R, D>): ValueProvider<T, R, D> {
        return except { it == anotherValueProvider }
    }

    /**
     * Removes providers matching [filter] from the current one.
     */
    fun except(filter: (ValueProvider<T, R, D>) -> Boolean): ValueProvider<T, R, D> =
        map { if (filter(it)) Combined(emptyList()) else it }

    /**
     * Applies [transform] for current provider
     */
    fun map(transform: (ValueProvider<T, R, D>) -> ValueProvider<T, R, D>): ValueProvider<T, R, D> =
        transform(this)

    /**
     * Uses fallback value provider in case when 'this' one failed to generate any value.
     */
    fun withFallback(fallback: ValueProvider<T, R, D>) : ValueProvider<T, R, D> {
        return Fallback(this, fallback)
    }

    /**
     * Creates a new value provider that creates default value if no values are generated by this provider.
     */
    fun withFallback(fallbackSupplier: (T) -> Seed<T, R>) : ValueProvider<T, R, D> {
        return withFallback { _, type ->
            sequenceOf(fallbackSupplier(type))
        }
    }

    fun letIf(flag: Boolean, block: (ValueProvider<T, R, D>) -> ValueProvider<T, R, D>) : ValueProvider<T, R, D> {
        return if (flag) block(this) else this
    }

    /**
     * Checks if current provider has fallback and return the initial provider.
     *
     * If the initial provider is also fallback, then it is unwrapped too.
     *
     * @return unwrapped provider or this if it is not fallback
     */
    fun unwrapIfFallback(): ValueProvider<T, R, D> {
        return if (this is Fallback<T, R, D>) { provider.unwrapIfFallback() } else { this }
    }

    private class Fallback<T, R, D : Description<T>>(
        val provider: ValueProvider<T, R, D>,
        val fallback: ValueProvider<T, R, D>,
    ): ValueProvider<T, R, D> {

        override fun enrich(description: D, type: T, scope: Scope) {
            provider.enrich(description, type, scope)
            // Enriching scope by fallback value provider in this point is not quite right,
            // but it doesn't look as a problem right now.
            fallback.enrich(description, type, scope)
        }

        override fun generate(description: D, type: T): Sequence<Seed<T, R>> {
            val default = if (provider.accept(type)) provider.generate(description, type) else emptySequence()
            return if (default.iterator().hasNext()) {
                default
            } else if (fallback.accept(type)) {
                fallback.generate(description, type)
            } else {
                emptySequence()
            }
        }

        override fun map(transform: (ValueProvider<T, R, D>) -> ValueProvider<T, R, D>): ValueProvider<T, R, D> =
            transform(Fallback(provider.map(transform), fallback.map(transform)))
    }

    /**
     * Wrapper class that delegates implementation to the [providers].
     */
    private class Combined<T, R, D : Description<T>>(providers: List<ValueProvider<T, R, D>>): ValueProvider<T, R, D> {
        val providers: List<ValueProvider<T, R, D>>

        init {
            // Flattening to avoid Combined inside Combined (for correct work of except, map, etc.)
            this.providers = providers.flatMap {
                if (it is Combined)
                    it.providers
                else
                    listOf(it)
            }
        }

        override fun enrich(description: D, type: T, scope: Scope) {
            providers.forEach { it.enrich(description, type, scope) }
        }

        override fun accept(type: T): Boolean {
            return providers.any { it.accept(type) }
        }

        override fun generate(description: D, type: T): Sequence<Seed<T, R>> = sequence {
            providers.asSequence().filter { it.accept(type) }.forEach { provider ->
                provider.generate(description, type).forEach {
                    yield(it)
                }
            }
        }

        override fun map(transform: (ValueProvider<T, R, D>) -> ValueProvider<T, R, D>): ValueProvider<T, R, D> =
            transform(Combined(providers.map { it.map(transform) }))
    }

    companion object {
        fun <T, R, D : Description<T>> of(valueProviders: List<ValueProvider<T, R, D>>): ValueProvider<T, R, D> {
            return Combined(valueProviders)
        }
    }
}

/**
 * Simple value provider for a concrete type.
 *
 * @param type that is used as a filter to call this provider
 * @param generate yields values for the type
 */
class TypeProvider<T, R, D : Description<T>>(
    val type: T,
    val generate: suspend SequenceScope<Seed<T, R>>.(description: D, type: T) -> Unit
) : ValueProvider<T, R, D> {
    override fun accept(type: T) = this.type == type
    override fun generate(description: D, type: T) = sequence {
        if (accept(type)) {
            this.generate(description, type)
        }
    }
}