set number of worker threads from RTS option; use forkOnIO to distribute workers...
[maximus:gmndl.git] / gmndl.hs
1 {-# LANGUAGE BangPatterns #-}
2 module Main (main) where
3
4 import Control.Concurrent (killThread)
5 import Control.Monad (when)
6 import Control.Monad.Trans (liftIO)
7 import Data.Array.IO (IOUArray, newArray, readArray, writeArray, inRange)
8 import Data.IORef (IORef, newIORef, readIORef, writeIORef)
9 import Data.PriorityQueue (PriorityQueue, newPriorityQueue, enqueue, dequeue)
10 import Foreign (mallocBytes, nullPtr, plusPtr, pokeArray, pokeByteOff, Word8)
11 import GHC.Conc (forkOnIO, numCapabilities)
12 import Graphics.UI.Gtk
13 import Graphics.UI.Gtk.OpenGL
14 import qualified Graphics.Rendering.OpenGL as GL
15 import Graphics.Rendering.OpenGL (($=), GLfloat)
16 import Numeric.QD.DoubleDouble (DoubleDouble(), fromDouble, toDouble)
17 import Numeric.QD.FPU.Raw (fpu_fix_start)
18
19 type B = Word8
20 type N = Int
21 type R = Double
22 type C = Cx DoubleDouble
23
24 data Cx a = Cx {-# UNPACK #-} !a {-# UNPACK #-} !a deriving (Read, Show, Eq)
25 instance Num a => Num (Cx a) where
26   {-# SPECIALIZE instance Num (Cx DoubleDouble) #-}
27   (!(Cx a b)) + (!(Cx c d)) = Cx (a + c) (b + d)
28   (!(Cx a b)) - (!(Cx c d)) = Cx (a - c) (b - d)
29   (!(Cx a b)) * (!(Cx c d)) = Cx (a * c - b * d) (a * d + b * c)
30   negate !(Cx a b) = Cx (-a) (-b)
31   abs x = error $ "Cx.abs: " ++ show x
32   signum x = error $ "Cx.signum: " ++ show x
33   fromInteger !x = Cx (fromInteger x) 0
34
35 hsv2rgb :: R -> R -> R -> (R, R, R)
36 hsv2rgb !h !s !v
37   | s == 0 = (v, v, v)
38   | h == 1 = hsv2rgb 0 s v
39   | otherwise =
40       let !i = floor (h * 6) `mod` 6 :: N
41           !f = (h * 6) - fromIntegral i
42           !p = v * (1 - s)
43           !q = v * (1 - s * f)
44           !t = v * (1 - s * (1 - f))
45       in  case i of
46             0 -> (v, t, p)
47             1 -> (q, v, p)
48             2 -> (p, v, t)
49             3 -> (p, q, v)
50             4 -> (t, p, v)
51             5 -> (v, p, q)
52             _ -> (0, 0, 0)
53
54 colour :: C -> C -> N -> (B, B, B)
55 colour !(Cx zr2 zi2) !(Cx dzr2 dzi2) !n =
56   let !zr = toDouble zr2
57       !zi = toDouble zi2
58       !dzr = toDouble dzr2
59       !dzi = toDouble dzi2
60       !il2 = 1 / log 2
61       !zd2 = zr * zr + zi * zi
62       !dzd2 = dzr * dzr + dzi * dzi
63       !d = (fromIntegral n :: R) - log (log zd2 / log escapeR2) * il2
64       !dwell = fromIntegral (floor d :: N)
65       !finala = atan2 zi zr
66       !de = (log zd2 * il2) * sqrt zd2 / sqrt dzd2
67       !dscale = log de * il2 + 32
68       !hue = log d * il2 / 3
69       !saturation = 0 `max` (log d * il2 / 8) `min` 1
70       !value = 0 `max` (1 - dscale / 48) `min` 1
71       !h = hue - fromIntegral (floor hue :: N)
72       !k = dwell / 2
73       !satf = if k - fromIntegral (floor k :: N) >= (0.5 :: R) then 0.9 else 1
74       !valf = if finala < 0 then 0.9 else 1
75       (!r, !g, !b) = hsv2rgb h (satf * saturation) (valf * value)
76       !rr = floor $ 0 `max` 255 * r `min` 255
77       !gg = floor $ 0 `max` 255 * g `min` 255
78       !bb = floor $ 0 `max` 255 * b `min` 255
79   in  (rr, gg, bb)
80
81 data Job = Job !N !N !N !C !C !C !N
82
83 priority :: Job -> (N, N)
84 priority (Job k _ _ _ _ _ n) = (n, k)
85
86 addJob :: IORef N -> C -> N -> PriorityQueue IO Job -> IOUArray (N,N) Bool -> N -> N -> IO ()
87 addJob count c zoom todo sync i j = do
88   already <- readArray sync (j, i)
89   when (not already) $ do
90     writeArray sync (j, i) True
91     k <- readIORef count
92     writeIORef count $! k + 1
93     enqueue todo $ Job k i j (coords c zoom i j) 0 0 0
94
95 renderer :: (N -> N -> B -> B -> B -> IO ()) -> C -> N -> IO (IO ())
96 renderer output c zoom = do
97   workerts <- mapM (\w -> forkOnIO w $ worker c zoom output w) [ 0 .. workers - 1 ]
98   return $ do
99     mapM_ killThread workerts
100
101 coords :: C -> N -> N -> N -> C
102 coords c zoom i j = c + Cx (fromIntegral (i - width`div`2) * k)
103                            (fromIntegral (height`div`2 - j) * k)
104   where !k = fromDouble (1/2^^zoom)
105
106 border :: [(N, N)]
107 border = concat
108   [ [ (j, i) | i <- [ 0 .. width  - 1 ], j <- [ 0 .. workers - 1 ] ]
109   , [ (j, i) | j <- [ 0 .. height - 1 ], i <- [ 0, width - 1 ] ]
110   , [ (j, i) | i <- [ 0 .. width  - 1 ], j <- [ height - workers .. height - 1 ] ]
111   ]
112
113 worker :: C -> N -> (N -> N -> B -> B -> B -> IO ()) -> N -> IO ()
114 worker c zoom output w = do
115   sync <- newArray rng False
116   queue <- newPriorityQueue priority
117   count <- newIORef 0
118   let addJ = addJob count c zoom queue sync
119   mapM_ (uncurry $ flip addJ) . filter mine $ border
120   compute addJ output queue
121   where
122     mine (j, _) = j `mod` workers == w
123
124 compute :: (N -> N -> IO ()) -> (N -> N -> B -> B -> B -> IO ()) -> PriorityQueue IO Job -> IO ()
125 compute addJ output queue = do
126   mjob <- dequeue queue
127   case mjob of
128     Just (Job k i j c z dz n) -> do
129       let done' z' dz' n' = do
130             let (r, g, b) = colour z' dz' n'
131             output i j r g b
132             sequence_
133               [ addJ x y
134               | u <- spreadX
135               , v <- spreadY
136               , let x = i + u
137               , let y = j + v
138               , inRange rng (y, x)
139               ]
140           todo' z' dz' n' = enqueue queue $ Job k i j c z' dz' n'
141       calculate c limit z dz n done' todo'
142       compute addJ output queue
143     Nothing -> return ()
144
145 calculate :: C -> N -> C -> C -> N -> (C -> C -> N -> IO ()) -> (C -> C -> N -> IO ()) -> IO ()
146 calculate !c !m0 !z0 !dz0 !n0 done todo = go m0 z0 dz0 n0
147   where
148     go !m !z@(Cx zr zi) !dz !n
149       | not (zr * zr + zi * zi < er2) = done z dz n
150       | m <= 0 = todo z dz n
151       | otherwise = go (m - 1) (z * z + c) (let !zdz = z * dz in zdz + zdz + 1) (n + 1)
152     !er2 = fromDouble escapeR2
153
154 main :: IO ()
155 main = do
156   _ <- unsafeInitGUIForThreadedRTS
157   _ <- initGL
158   glconfig <- glConfigNew [ GLModeRGBA, GLModeDouble ]
159   canvas <- glDrawingAreaNew glconfig
160   widgetSetSizeRequest canvas width height
161   imgdata <- mallocBytes $ width * height * 3
162   pokeArray imgdata (replicate (height * width * 3) (255 :: B))
163   let output x y r g b = do
164         let p = imgdata `plusPtr` ((y * width + x) * 3)
165         pokeByteOff p 0 r
166         pokeByteOff p 1 g
167         pokeByteOff p 2 b
168   window <- windowNew
169   eventb <- eventBoxNew
170   set window [ containerBorderWidth := 0, containerChild := eventb ]
171   set eventb [ containerBorderWidth := 0, containerChild := canvas ]
172   mapM_ (flip forkOnIO $ fpu_fix_start nullPtr) [ 0 .. numCapabilities - 1 ]
173   stop0 <- renderer output c0 zoom0
174   sR <- newIORef (c0, zoom0, stop0)
175   _ <- eventb `on` buttonPressEvent $ {-# SCC "cb/event" #-} tryEvent $ do
176     LeftButton <- eventButton
177     (x, y) <- eventCoordinates
178     liftIO $ do
179       (c, zoom, stop) <- readIORef sR
180       stop
181       pokeArray imgdata (replicate (height * width * 3) (255 :: B))
182       let c' = c + ((Cx (fromDouble x) (fromDouble (-y))) - (fromIntegral width `Cx` fromIntegral (-height))*Cx (fromDouble 0.5) 0) * Cx (fromDouble (1/2^^zoom)) 0
183           zoom' = zoom + 1
184       stop' <- renderer output c' zoom'
185       writeIORef sR (c', zoom', stop')
186       print (c', zoom')
187   _ <- onRealize canvas $ {-# SCC "cb/realize" #-}withGLDrawingArea canvas $ \_ -> do
188     GL.clearColor $= (GL.Color4 0.0 0.0 0.0 0.0)
189     GL.matrixMode $= GL.Projection
190     GL.loadIdentity
191     GL.ortho 0.0 1.0 0.0 1.0 (-1.0) 1.0
192     GL.drawBuffer $= GL.BackBuffers
193     [tex] <- GL.genObjectNames 1
194     GL.texture GL.Texture2D $= GL.Enabled
195     GL.textureBinding GL.Texture2D $= Just tex
196     GL.texImage2D Nothing GL.NoProxy 0 GL.RGB' (GL.TextureSize2D (fromIntegral width) (fromIntegral height)) 0 (GL.PixelData GL.RGB GL.UnsignedByte nullPtr)
197     GL.textureFilter GL.Texture2D $= ((GL.Nearest, Nothing), GL.Nearest)
198     GL.textureWrapMode GL.Texture2D GL.S $= (GL.Repeated, GL.ClampToEdge)
199     GL.textureWrapMode GL.Texture2D GL.T $= (GL.Repeated, GL.ClampToEdge)
200   _ <- onExpose canvas $ {-# SCC "cb/expose" #-} \_ -> do
201     withGLDrawingArea canvas $ \glwindow -> do
202       let v :: GLfloat -> GLfloat -> GLfloat -> GLfloat -> IO ()
203           v tx ty vx vy = GL.texCoord (GL.TexCoord2 tx ty) >> GL.vertex (GL.Vertex2 vx vy)
204           w = fromIntegral width
205           h = fromIntegral height
206       GL.clear [ GL.ColorBuffer ]
207       GL.texSubImage2D Nothing 0 (GL.TexturePosition2D 0 0) (GL.TextureSize2D w h) (GL.PixelData GL.RGB GL.UnsignedByte imgdata)
208       GL.renderPrimitive GL.Quads $ do
209         v 0 1 0 0 >> v 0 0 0 1 >> v 1 0 1 1 >> v 1 1 1 0
210       glDrawableSwapBuffers glwindow
211     return True
212   _ <- onDestroy window mainQuit
213   _ <- timeoutAdd (widgetQueueDraw canvas >> return True) 200
214   widgetShowAll window
215   mainGUI
216
217 spreadX :: [ N ]
218 spreadX = [ -1, 0, 1 ]
219
220 spreadY :: [ N ]
221 spreadY = [ -workers, 0, workers ]
222
223 workers :: N
224 workers = numCapabilities
225
226 limit :: N
227 limit = (2^(12::N)-1)
228
229 width, height :: N
230 width = 512
231 height = 512
232
233 rng :: ((N, N), (N, N))
234 rng = ((0, 0), (height - 1, width - 1))
235
236 c0 :: C
237 c0 = 0
238
239 zoom0 :: N
240 zoom0 = 6
241
242 escapeR, escapeR2 :: R
243 escapeR = 65536
244 escapeR2 = escapeR * escapeR