mandelbrot.hs

import Control.Applicative
import Control.Concurrent
import Control.Monad
import Data.Complex
import Data.List (find)
import Graphics.EasyRaster.GTK
import System.Environment
import System.IO

-- | Type representing points in an image
type Point = (Coord, Coord)

-- | Type representing colors of pixels
type Color = (Channel, Channel, Channel)

-- | Communicate results between threads.
data Result a = Result { getResult :: a }  -- | A result
              | Done                      -- | Process is done
              deriving (Eq)

infinity = 1.0/0.0

-- | Some nice-looking colors.  (Actually they are hideous, but I am
-- lazy...)
colors :: [Color]
colors = [(0,0,255),(0,255,0),(0,255,255),(255,0,0),
          (255,0,255),(255,255,0),(0,0,128),(0,128,0),(0,128,128),(128,0,0),
          (128,0,128),(128,128,0),(128,128,255),(128,255,128),
          (128,255,255),(255,128,128),
          (255,128,255),(255,255,128),(128,128,128)]

-- | Calculate the sequence for a given c.
iterateMandelbrot :: Complex Double -> [Complex Double]
iterateMandelbrot c = iterate (\z -> z ** 2 + c) c

-- | Calculate whether a given complex number "c" is in the Mandelbrot
-- set.  Return 0 if the point is in the set, number of iterations
-- until divergence otherwise.
inMandelbrotSet :: Complex Double -> Int
inMandelbrotSet c = case find (\(i,z) -> magnitude z == infinity) (take 50 zs) of
  Just (i,_) -> i
  Nothing -> 0
  where zs = zip [1..] (iterateMandelbrot c)

-- | Convert a pixel position in the image buffer to a coordinate on
-- the complex plane.  (One axis.)
toPlaneCoord :: Double -> Double -> Int -> Int -> Double
toPlaneCoord min max x xmax =
  (s * ((fromIntegral x) / (fromIntegral xmax))) + min
    where s = max - min

-- | Given the maximum x and y bounds of the image, return a function
-- that calculates the color for any point on the image.
mandelbrotSet :: Int -> Int -> Point -> Color
mandelbrotSet xmax ymax (x,y) = case inMandelbrotSet c of
  0 -> (0,0,0)
  x -> colors !! (x `mod` 19)
 where c = toPlaneCoord (-2) 0.5 x xmax :+ toPlaneCoord (-1) 1 y ymax

-- | Called when a thread is finished; kills the thread and notifies
-- the coordinator of its completion.
imDone mvar = do
  putMVar mvar True
  myThreadId >>= killThread

-- | Given a pair of coordinates representing the size of the image,
-- and a function from coordinates to colors, run the function for
-- each point in the image and return the generated image.
doIb :: Point -> (Point -> Color) -> IO ImageBuffer
doIb pmax f = do
  ib <- ib_new pmax
  todo <- newChan
  results <- newChan
  done <- newEmptyMVar

  -- consume colors to write, write them to image
  forkIO $ forever $ do
    result <- readChan results
    when (result == Done) $ imDone done

    let (p,c) = getResult result
    let (rowNum,colNum) = p
    ( ib_write_pixel ib p ) $! c
    when (colNum `mod` 20 == 0 && rowNum == 0) $ print colNum

  -- consume list of pixels, produce colors to write
  let numThreads = 6
  replicateM_ numThreads $ forkIO $ forever $ do
    ps <- readChan todo
    when (null ps) $ imDone done
    forM_ ps (\p -> do
                 let fp = f $! p
                 writeChan results $! (fp `seq` Result (p,fp)))

  -- produce list of coordinates
  forM (ib_coords2 ib) $ \ps -> do
    writeChan todo ps

  -- ensure workers kill themselves when they get to the end of the chan
  replicateM_ numThreads $ (writeChan todo [])
  putStrLn "workers started"

  -- wait for workers to finish
  replicateM_ numThreads (takeMVar done)
  putStrLn "workers done"

  -- wait for image writer to finish
  writeChan results Done
  takeMVar done

  return ib

-- | Main function; generate the image of size xmax, ymax and save it
-- to my homedir as a PNG.
mainMandelbrot xmax ymax = do
  ib <- doIb (xmax,ymax) (mandelbrotSet xmax ymax)
  putStrLn "Writing image..."
  ib_save ib IFT_PNG "/home/jon/out.png"
  putStrLn "done"
  return ()

-- | Entry point.  Accepts two args on the command-line, xmax and
-- ymax.
main = do
  init_system
  xmax:ymax:[] <- getArgs
  mainMandelbrot (read xmax) (read ymax)
  return ()