Prime sieve

Primes.hs

{-# OPTIONS_GHC -Wall #-}

module Main where

import Control.Monad ((<=<))
import Data.Maybe (listToMaybe)
import Numeric.Natural (Natural)
import System.Environment (getArgs)
import Text.Read (readMaybe)

data Stream a = Stream a (Stream a)

unfoldStream :: (b -> (a, b)) -> b -> Stream a
unfoldStream f b =
  let (a, b') = f b
   in Stream a (unfoldStream f b')

takeStream :: Int -> Stream a -> [a]
takeStream 0 _ = []
takeStream n (Stream a as) = a : takeStream (n - 1) as

naturals :: Stream Natural
naturals = unfoldStream (\n -> (n, n + 1)) 0

divisibleBy :: Integral a => a -> a -> Bool
divisibleBy n m = 0 == mod m n

filterStream :: (a -> Bool) -> Stream a -> Stream a
filterStream p (Stream a as) =
  if p a
    then Stream a (filterStream p as)
    else filterStream p as

dropStream :: Int -> Stream a -> Stream a
dropStream 0 as = as
dropStream n (Stream _ as) = dropStream (n - 1) as

primes :: Stream Natural
primes =
  unfoldStream
    (\(Stream n ns) -> (n, filterStream (not . divisibleBy n) ns))
    (dropStream 2 naturals)

main :: IO ()
main = do
  Just n <- (readMaybe <=< listToMaybe) <$> getArgs
  print $ takeStream n primes