HedgehogHelpers.fs

open Hedgehog

module Gen =

    let control (f : unit -> 'a) : Gen<'a> =
        Random (fun _ _ ->
            let next = f ()
            Node (next, LazyList.empty))
        |> Gen.ofRandom
        
    let ofSequence (source : 'a seq) : Gen<'a> =
        let enumerator = source.GetEnumerator()
        control (fun _ ->
            enumerator.MoveNext() |> ignore
            enumerator.Current)

    let roundRobin (candidates : 'a list) : Gen<'a> =
        Seq.initInfinite <| fun i ->
            candidates
            |> List.item (i % (candidates |> List.length))
        |> ofSequence
        
    let indexed (g : Gen<'a>) : Gen<int * 'a> =
        let indexGenerator =
            Seq.initInfinite id
            |> ofSequence
        Gen.map2 (fun i x -> (i, x)) indexGenerator g
        
    let rec private listDistinctByInternal
        (projection : 'a -> 'b)
        (length : int)
        (generator : Gen<'a>)
        (prevResults : 'a list) : Gen<List<'a>> = gen {
            let! curr = generator
            let results =
                curr :: prevResults
                |> List.distinctBy projection
            if (results |> List.length) < length
            then return! listDistinctByInternal projection length generator results
            else return results }

    let listDistinctBy
        (projection : 'a -> 'b)
        (range : Range<int>)
        (generator : Gen<'a>) : Gen<List<'a>> = gen {
            let! length = range |> Gen.integral
            return! listDistinctByInternal projection length generator [] }

    let listDistinct (range : Range<int>) (generator : Gen<'a>) : Gen<List<'a>> =
        listDistinctBy id range generator