app "prime2"
    packages { pf: "https://github.com/roc-lang/basic-cli/releases/download/0.7.1/Icc3xJoIixF3hCcfXrDwLCu4wQHtNdPyoJkEbkgIElA.tar.br" }
    imports [pf.Stdout, pf.Task]
    provides [main] to pf

main : Task.Task {} I32
main =
    numStrings =
        primesIn 50_000_000
        |> List.map Num.toStr
        |> Str.joinWith "\n"

    Stdout.line "$(numStrings)"

primesIn : U64 -> List U64
primesIn = \n ->
    List.range { start: At 0, end: At n }
    |> List.map Yes
    |> primeSieve 0 n
    |> List.keepIf isYes
    |> List.map extractValueUnsafe

PrimeResult : [No, Yes U64]

primeSieve : List PrimeResult, U64, U64 -> List PrimeResult U64
primeSieve = \l, idx, max ->
    len = List.len l
    when nextPrime l idx is
        Ok p if (Num.powInt p 2) > max -> l
        Ok p ->
            setNo l (List.range { start: At (Num.toNat (p * 2)), step: Num.toNat p, end: At len })
            |> primeSieve (idx + 1) max

        Err _ -> l

isYes : PrimeResult -> Bool
isYes = \r ->
    when r is
        No -> Bool.false
        Yes _ -> Bool.true

setNo : List PrimeResult, List Nat -> List PrimeResult U64
setNo = \l, idxs ->
    when idxs is
        [] -> l
        [idx] -> l |> List.set idx No
        [head, .. as tail] -> l |> List.set head No |> setNo tail

nextPrime : List PrimeResult, U64 -> Result U64 [NotFound]
nextPrime = \l, startIdx ->
    l
    |> List.dropFirst (Num.toNat (startIdx + 1))
    |> List.findFirst isYes
    |> Result.map extractValueUnsafe

extractValueUnsafe : PrimeResult -> U64
extractValueUnsafe = \res ->
    when res is
        Yes n -> n
        No -> crash "No value found"

countPrimesIn : U64 -> Nat
countPrimesIn = \limit ->
    List.range { start: At 2, end: At limit }
    |> List.countIf \n ->
        if n < 4 then
            Bool.true
        else if Num.isEven n then
            Bool.false
        else
            res =
                List.range { start: At 3, end: At (Num.ceiling (Num.sqrt (Num.toF64 n))), step: 2 }
                |> List.any \a -> Num.isMultipleOf n a
                |> Bool.not

            res

expect
    (
        countPrimesIn 100000
        |> \expected ->
            dbg expected

            expected
    )
    == (
        List.len (primesIn 100000)
        |> \actual ->
            dbg actual

            actual

    )