mandelbrot Haskell GHC #2 program

--
-- The Computer Language Benchmarks Game
-- https://salsa.debian.org/benchmarksgame-team/benchmarksgame/
--
-- Contributed by Spencer Janssen, Trevor McCort, Christophe Poucet and Don Stewart
-- Parallelised by Tim Newsham
--
-- Must be compiled with the -fexcess-precision flag as a pragma. GHC
-- currently doesn't recognise the -fexcess-precision flag on the command
-- line (!).
--
-- The following flags are suggested when compiling:
--
--      ghc -optc-march=pentium4 -optc-mfpmath=sse -optc-msse2 -threaded --make
--
-- Run with -N6 on a quad core (more capabilities to hide latency)
--
--      $ time ./A 6400 +RTS -N6
--

import System.Environment
import System.IO
import Foreign
import Foreign.Marshal.Array
import Control.Concurrent
import Control.Concurrent.MVar
import Control.Concurrent.Chan
import Control.Monad

main = do
    -- width in pixels
    w <- getArgs >>= readIO . head
        -- width in bytes
    let n      = w `div` 8
        -- width of a pixel in the complex plane
        m  = 2 / fromIntegral w
        coords = [T 1 0 y (fromIntegral y * m - 1) | y <- [0..w-1]]
    q <- newChan
    replies <- replicateM w newEmptyMVar
    mapM_ (writeChan q) $ zip coords replies
    replicateM_ 4 . forkIO $ worker q w m n

    putStrLn ("P4\n"++show w++" "++show w)
    mapM_ (takeMVar >=> \b -> hPutBuf stdout b n) replies

-- Worker computes one line of the image and sends it to the master
-- q - work queue
-- w - width in pixels
-- m - width of a pixel in the complex plane
-- n - width in bytes
worker q w m n = forever (do
    (coord, reply) <- readChan q
    p <- mallocArray0 n
    unfold (next_x w m n) p coord
    putMVar reply p)

-- f - takes coordinates and returns Nothing if done
--     or the next byte of the bitmap otherwise.
-- ptr - buffer to write to
-- x0 - initial coordinates 
unfold :: (T -> Maybe (Word8,T)) -> Ptr Word8 -> T -> IO (Ptr Word8)
unfold !f !ptr !x0 = go ptr x0
  where
    -- p - pointer into the buffer
    -- x - coordinates
    go !p !x = case f x of
        Just (w,y)          -> poke p w >> go (p `plusPtr` 1) y
        Nothing             -> return ptr

-- T bs x y ci
--    bx - x position in bytes
--    x  - x position in pixels
--    y  - y position in pixels
--    ci - y position in complex plane
data T = T !Int !Int !Int !Double

-- w - image width in pixels
-- iw - pixel width in the complex plane
-- bw - image width in bytes
next_x !w !iw !bw (T bx x y ci)
    | bx == bw  = Nothing
    | otherwise = Just (loop_x w x 8 iw ci 0, T (bx+1) (x+8) y ci)

-- w - image width in pixels
-- x - current x coordinate in pixels
-- n - bit positition from 8 to 0
-- iw - pixel width in the complex plane
-- ci - current y coordinate in complex plane
-- b - accumulated bit value.
loop_x !w !x !n !iw !ci !b
    | x < w = if n == 0
                    then b
                    else loop_x w (x+1) (n-1) iw ci (b+b+v)
    | otherwise = b `shiftL` n
  where
    v = fractal 0 0 (fromIntegral x * iw - 1.5) ci 50

-- julia function (r :+ i) (cr :+ ci) with max iterations k.
fractal :: Double -> Double -> Double -> Double -> Int -> Word8
fractal !r !i !cr !ci !k
    | r2 + i2 > 4 = 0
    | k == 0      = 1
    | otherwise   = fractal (r2-i2+cr) ((r+r)*i+ci) cr ci (k-1)
  where
    (!r2,!i2) = (r*r,i*i)
    

NOTES:
64-bit Ubuntu quad core
The Glorious Glasgow Haskell
Compilation System,
version 9.4.8
LLVM version 17.0.2


 Thu, 07 Mar 2024 23:25:54 GMT

MAKE:
mv mandelbrot.ghc-2.ghc mandelbrot.ghc-2.hs
~/.ghcup/bin/ghc --make -optlo-passes='module(default<O2>)' -O2 -fno-llvm-tbaa -XBangPatterns -threaded -rtsopts  mandelbrot.ghc-2.hs -o mandelbrot.ghc-2.ghc_run
Loaded package environment from /home/dunham/.ghc/x86_64-linux-9.4.8/environments/default
[1 of 2] Compiling Main             ( mandelbrot.ghc-2.hs, mandelbrot.ghc-2.o )
[2 of 2] Linking mandelbrot.ghc-2.ghc_run
rm mandelbrot.ghc-2.hs

11.67s to complete and log all make actions

COMMAND LINE:
 ./mandelbrot.ghc-2.ghc_run +RTS -N4 -RTS 16000

(BINARY) PROGRAM OUTPUT NOT SHOWN