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