roc-stuff/prime-bin-search.roc

app "prime-bin-search"
    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 500_000
        |> List.map Num.toStr
        |> Str.joinWith "\n"

    Stdout.line "$(numStrings)"

primesIn : U64 -> List U64
primesIn = \n ->
    numbers = List.range { start: At 2u64, end: At n }
    List.walk numbers numbers primeWalk

primeWalk : List U64, U64 -> List U64
primeWalk = \l, n ->
    if l |> binarySearchContains n then
        if n > 1 then
            l
            |> List.dropIf \e ->
                e > n && Num.isMultipleOf e n
        else
            l
    else
        l

binarySearchContains : List U64, U64 -> Bool where a implements Eq
binarySearchContains = \l, a ->
    len = List.len l

    if Num.isZero len then
        Bool.false
    else
        half = Num.shiftRightBy len 1
        binarySearchContainsRec l a half half

binarySearchContainsRec : List U64, U64, Nat, Nat -> Bool where a implements Eq
binarySearchContainsRec = \l, n, mid, step ->
    if step <= 1 then
        sublist = List.sublist l { start: (mid - (Num.min step mid)), len: (step * 2 + 1) }

        when sublist is
            [] -> Bool.false
            [a] -> n == a
            [a, b] -> n == a || n == b
            [a, b, c] -> n == a || n == b || n == c
            _ -> crash "This shoulnd't have happened"
    else
        element = List.get l mid |> Result.withDefault 0

        when Num.compare n element is
            EQ -> Bool.true
            ne ->
                newStep = Num.divCeil step 2
                len = l |> List.len

                when ne is
                    GT if mid >= len -> Bool.false
                    GT -> binarySearchContainsRec l n (mid + newStep) newStep
                    LT if mid < newStep -> binarySearchContainsRec l n 0 newStep
                    LT -> binarySearchContainsRec l n (mid - newStep) newStep
                    _ -> crash "This shoulnd't have happened"

orderedContains : List (Num a), Num a -> Bool where a implements Eq
orderedContains = \l, a ->
    when l |> List.findFirst \e -> e >= a is
        Err NotFound -> Bool.false
        Ok f -> f == a

expect orderedContains [0, 1, 2, 3, 4, 5, 6] 5 == Bool.true
expect binarySearchContains [0, 1, 2, 3, 4, 5, 6] 5 == Bool.true
expect binarySearchContains [0, 1, 3, 4, 5, 6] 2 == Bool.false
expect binarySearchContains [0, 1, 3, 4, 5, 6] 7 == Bool.false
expect binarySearchContains [] 7 == Bool.false
expect binarySearchContains [2] 2 == Bool.true
expect binarySearchContains [1, 2, 3] 3 == Bool.true
expect binarySearchContains [1, 2, 3] 1 == Bool.true
expect binarySearchContains [1, 2, 3, 4, 5, 6] 0 == Bool.false

expect binarySearchContains (List.range { start: At 0, end: At 100000 }) 4 == Bool.true
expect binarySearchContains (List.range { start: At 0, end: At 100000 } |> List.dropIf Num.isEven) 4 == Bool.false
expect binarySearchContains (List.range { start: At 0, end: At 100000 } |> List.dropIf Num.isOdd) 97 == Bool.false
expect binarySearchContains (List.range { start: At 0, end: At 100000 }) 99998 == Bool.true
expect binarySearchContains (List.range { start: At 0, end: At 100000 }) 100001 == Bool.false

countPrimesIn : U64 -> Nat
countPrimesIn = \limit ->
    List.range { start: At 2, end: At limit }
    |> List.countIf \n ->
        if n < 4 then
            Bool.true
        else if n % 2 == 0 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

    )