Eraosthenes's sieve with Haskell's Concurrency

primes.hs

{-# LANGUAGE TypeApplications #-}

import Control.Concurrent
import Control.Concurrent.MVar
import Control.Monad
import Control.Parallel
import Data.Vector ((!))
import Debug.Trace
import System.Exit (die)
import qualified Data.Vector as V

n :: Int
n = 1000000

-- updates the mvar v with a
updMVar :: MVar a -> a -> IO ()
updMVar v a = do
  _ <- takeMVar v
  putMVar v a

worker :: MVar Int -> V.Vector (MVar Bool) -> MVar Bool -> IO ()
worker i v d = do

  idx <- takeMVar i
  putMVar i (idx + 1) -- update current position

  print idx

  if (idx >= n) 
    then updMVar d True -- if done set it 
    else do

    forkIO $ worker i v d -- otherwise fork a worker
    
    forM_ [idx*2, idx*3  .. n - 1] $ \j -> do
      updMVar (v ! j) False

finish :: V.Vector (MVar Bool) -> IO ()
finish vec = do
  ans <- V.foldM (\s v -> do
                     ok <- takeMVar v
                     putMVar v ok
                     return $ if ok then (s + 1) else s)
         0 vec
  die $ "# primes = " ++ show ans ++ "\n"

  -- This is not actually the answer, but is the answer + 2
  -- but this mathematical inaccuracy shouldn't affect the program.
  
main = do
  cidxVar <- newMVar (2 :: Int) -- stores current index in sieve
  cvecVar <- V.replicateM n (newMVar True) -- sieve values as vector of mvars
  doneVar <- newMVar False -- is everything done?

  print $ "Okay done."

  forkIO $ worker cidxVar cvecVar doneVar
  
  print "Forked."
  
  forever $ do
    -- print "This shouldn't be needed." 
    -- ^^ Removing this line causes the program to stop in between, and hang.
    isDone <- takeMVar doneVar -- check if we're done
    if (isDone)
      then finish cvecVar      -- then finish the job
      else putMVar doneVar isDone -- else put it back and continue