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