How can I make this prime check function more succinct in Haskell?

Question

So I spent the last hour or so trying to write up an effecient haskell function that checks if a number is prime or not. I came up with this algorithm that I'm rather happy with:

prime :: Int -> Bool
prime x
    | x == 1 = False
    | x == 2 = True
    | x == 3 = True
    | x `mod` 2 == 0 = False
    | x `mod` 3 == 0 = False
    | otherwise = all (\y -> x `mod` y /= 0) $ dividends x
    where
        dividends z =
            takeWhile ((<=) . floor . sqrt . fromIntegral $ z)
                $ concatMap (\x -> [x-1, x+1]) [6 * x | x <- [1..]]

I also uploaded a notebook that checks the run time of the algorithm and compares it to the sieve method in case anyone is interested that is here: https://anaconda.org/freemo/primes/notebook

My question is, as someone new to haskell how can I make this algorithm more idomatic. I have a feeling the anonymous functions I used can be eliminated, and there are probably other ways I can make it more concise without having a negative effect on the run time. How can I write this in an more idiomatic haskell way?

Answer 1

So I modified @leftaroundabout 's answer slightly and am using this. Just putting it up here in case anyone winds up using it.

import Control.Applicative

prime'' :: Int -> Bool
prime'' 1 = False
prime'' 2 = True
prime'' 3 = True
prime'' x
    | x `mod` 2 == 0 = False
    | x `mod` 3 == 0 = False
    | otherwise = all
                     -- check if x is not divisibile by the divisors
                     ((/=  0) . (x `mod`))

                     -- only take divisors less than or equal to sqrt(x)
                     . takeWhile (<=(floor . sqrt . fromIntegral $ x))

                     -- generate divisors as an infinite list 6n +/- 1
                     $ [6,12..] <**> [(-1+), (1+)]

Answer 2

The thing I'd consider most unidiomatic here are the equality comparisons. In Haskell, we generally prefer pattern matching, which you can directly apply to the first three cases:

prime 1 = False
prime 2 = True
prime 3 = True
prime x
    | ...

(The difference is cosmetic in this example, but often pattern matching makes code much safer, more succinct and sometimes also significantly more performant.)

Next, you're using both a list comprehension and concatMap. Both constructs do largely the same thing, and are equivalent to monadic binds. But a) it's usually benefitial to use only one of these syntaxes in a place, b) monad is actually stronger than necessary: applicative will do as well, and can be written quite nicely here.

import Control.Applicative
...
       (6*) <$> [1..] <**> [pred, succ]

Then, I find this ((<=) . floor . sqrt . fromIntegral $ z) awkward. More natural would be in flipped form:

    takeWhile (>= floor (sqrt $ fromIntegral z)) ...

...but generally speaking, better yet is to avoid the square root entirely:

    takeWhile ((>=z) . (^2)) ...

The reason being, square-root is pretty expensive and has floating-point inaccuracies, so it's often better to square the other side of the equation instead. Unfortunately, this is actually not reliable here, because the squaring might lead to Int overflows, and Integer would make the performance probably worse than with float-sqrt, because it's done for every element of the list rather than just once. So, a square-root is actually sensible here.

The dividends function is actually superfluous since you only ever call it with x as the argument. (It might be argued that it's sensible to still give this a name, but at any rate you don't need to give it an argument.)

The predicate for the all is something I would consider making point-free.

Final version would then be:

prime :: Int -> Bool
prime 1 = False
prime 2 = True
prime 3 = True
prime x
    | x `mod` 2 == 0 = False
    | x `mod` 3 == 0 = False
    | otherwise = all ((/= 0) . (x`mod`))
                 . takeWhile (>= floor (sqrt $ fromIntegral x))
                 $ (6*) <$> [1..]
                        <**> [pred, succ]