support Turbo sqr and twicce; SPECIALIZE some more; renderer' refactored; convert...
[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 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 = zr * zr + zi * zi
111       !dzd2 = dzr * dzr + dzi * 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 !N !(Complex c) !(Complex c) !(Complex c) !N
131
132 priority :: Job c -> (N, N)
133 priority (Job k _ _ _ _ _ n) = (n, k)
134
135 addJob :: RealFloat c => IORef N -> Complex c -> N -> PriorityQueue IO (Job c) -> IOUArray (N,N) Bool -> N -> N -> IO ()
136 {-# SPECIALIZE addJob :: IORef N -> Complex Double -> N -> PriorityQueue IO (Job Double) -> IOUArray (N,N) Bool -> N -> N -> IO () #-}
137 {-# SPECIALIZE addJob :: IORef N -> Complex DoubleDouble -> N -> PriorityQueue IO (Job DoubleDouble) -> IOUArray (N,N) Bool -> N -> N -> IO () #-}
138 {-# SPECIALIZE addJob :: IORef N -> Complex QuadDouble -> N -> PriorityQueue IO (Job QuadDouble) -> IOUArray (N,N) Bool -> N -> N -> IO () #-}
139 addJob count 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     k <- readIORef count
144     writeIORef count $! k + 1
145     enqueue todo $ Job k i j (coords c zoom i j) 0 0 0
146
147 renderer :: (Turbo c, RealFloat c) => (N -> N -> B -> B -> B -> IO ()) -> Complex c -> N -> IO (IO ())
148 {-# SPECIALIZE renderer :: (N -> N -> B -> B -> B -> IO ()) -> Complex Double -> N -> IO (IO ()) #-}
149 {-# SPECIALIZE renderer :: (N -> N -> B -> B -> B -> IO ()) -> Complex DoubleDouble -> N -> IO (IO ()) #-}
150 {-# SPECIALIZE renderer :: (N -> N -> B -> B -> B -> IO ()) -> Complex QuadDouble -> N -> IO (IO ()) #-}
151 renderer output c zoom = do
152   workerts <- mapM (\w -> forkOnIO w $ worker c zoom output w) [ 0 .. workers - 1 ]
153   return $ do
154     mapM_ killThread workerts
155
156 coords :: RealFloat c => Complex c -> N -> N -> N -> Complex c
157 {-# SPECIALIZE coords :: Complex Double -> N -> N -> N -> Complex Double #-}
158 {-# SPECIALIZE coords :: Complex DoubleDouble -> N -> N -> N -> Complex DoubleDouble #-}
159 {-# SPECIALIZE coords :: Complex QuadDouble -> N -> N -> N -> Complex QuadDouble #-}
160 coords c zoom i j = c + ( (fromIntegral (i - width`div`2) * k)
161                         :+(fromIntegral (height`div`2 - j) * k))
162   where !k = convert (1/2^^zoom :: Double)
163
164 border :: [(N, N)]
165 border = concat
166   [ [ (j, i) | i <- [ 0 .. width  - 1 ], j <- [ 0 .. workers - 1 ] ]
167   , [ (j, i) | j <- [ 0 .. height - 1 ], i <- [ 0, width - 1 ] ]
168   , [ (j, i) | i <- [ 0 .. width  - 1 ], j <- [ height - workers .. height - 1 ] ]
169   ]
170
171 worker :: (Turbo c, RealFloat c) => Complex c -> N -> (N -> N -> B -> B -> B -> IO ()) -> N -> IO ()
172 {-# SPECIALIZE worker :: Complex Double -> N -> (N -> N -> B -> B -> B -> IO ()) -> N -> IO () #-}
173 {-# SPECIALIZE worker :: Complex DoubleDouble -> N -> (N -> N -> B -> B -> B -> IO ()) -> N -> IO () #-}
174 {-# SPECIALIZE worker :: Complex QuadDouble -> N -> (N -> N -> B -> B -> B -> IO ()) -> N -> IO () #-}
175 worker c zoom output w = do
176   sync <- newArray rng False
177   queue <- newPriorityQueue priority
178   count <- newIORef 0
179   let addJ = addJob count c zoom queue sync
180   mapM_ (uncurry $ flip addJ) . filter mine $ border
181   compute addJ output queue
182   where
183     mine (j, _) = j `mod` workers == w
184
185 compute :: (Turbo c, RealFloat c) => (N -> N -> IO ()) -> (N -> N -> B -> B -> B -> IO ()) -> PriorityQueue IO (Job c) -> IO ()
186 {-# SPECIALIZE compute :: (N -> N -> IO ()) -> (N -> N -> B -> B -> B -> IO ()) -> PriorityQueue IO (Job Double) -> IO () #-}
187 {-# SPECIALIZE compute :: (N -> N -> IO ()) -> (N -> N -> B -> B -> B -> IO ()) -> PriorityQueue IO (Job DoubleDouble) -> IO () #-}
188 {-# SPECIALIZE compute :: (N -> N -> IO ()) -> (N -> N -> B -> B -> B -> IO ()) -> PriorityQueue IO (Job QuadDouble) -> IO () #-}
189 compute addJ output queue = do
190   mjob <- dequeue queue
191   case mjob of
192     Just (Job k i j c z dz n) -> do
193       let done' !(zr:+zi) !(dzr:+dzi) !n' = do
194             let (r, g, b) = colour (convert zr :+ convert zi) (convert dzr :+ convert dzi) n'
195             output i j r g b
196             sequence_
197               [ addJ x y
198               | u <- spreadX
199               , v <- spreadY
200               , let x = i + u
201               , let y = j + v
202               , inRange rng (y, x)
203               ]
204           todo' z' dz' n' = enqueue queue $ Job k i j c z' dz' n'
205       calculate c limit z dz n done' todo'
206       compute addJ output queue
207     Nothing -> return ()
208
209 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 ()
210 {-# SPECIALIZE calculate :: Complex Double -> N -> Complex Double -> Complex Double -> N -> (Complex Double -> Complex Double -> N -> IO ()) -> (Complex Double -> Complex Double -> N -> IO ()) -> IO () #-}
211 {-# SPECIALIZE calculate :: Complex DoubleDouble -> N -> Complex DoubleDouble -> Complex DoubleDouble -> N -> (Complex DoubleDouble -> Complex DoubleDouble -> N -> IO ()) -> (Complex DoubleDouble -> Complex DoubleDouble -> N -> IO ()) -> IO () #-}
212 {-# SPECIALIZE calculate :: Complex QuadDouble -> N -> Complex QuadDouble -> Complex QuadDouble -> N -> (Complex QuadDouble -> Complex QuadDouble -> N -> IO ()) -> (Complex QuadDouble -> Complex QuadDouble -> N -> IO ()) -> IO () #-}
213 calculate !c !m0 !z0 !dz0 !n0 done todo = go m0 z0 dz0 n0
214   where
215     go !m !z@(zr:+zi) !dz !n
216       | not (sqr zr + sqr zi < er2) = done z dz n
217       | m <= 0 = todo z dz n
218       | otherwise = go (m - 1) (sqr z + c) (let !zdz = z * dz in twice zdz + 1) (n + 1)
219     !er2 = convert escapeR2
220
221 renderer' output !c !zoom
222   | zoom < 48 = renderer output (f c :: Complex Double      ) zoom
223   | zoom < 96 = renderer output (f c :: Complex DoubleDouble) zoom
224   | otherwise = renderer output (f c :: Complex QuadDouble  ) zoom
225   where f !(cr :+ ci) = convert cr :+ convert ci
226
227 main :: IO ()
228 main = do
229   _ <- unsafeInitGUIForThreadedRTS
230   _ <- initGL
231   glconfig <- glConfigNew [ GLModeRGBA, GLModeDouble ]
232   canvas <- glDrawingAreaNew glconfig
233   widgetSetSizeRequest canvas width height
234   imgdata <- mallocBytes $ width * height * 3
235   pokeArray imgdata (replicate (height * width * 3) (255 :: B))
236   let output x y r g b = do
237         let p = imgdata `plusPtr` ((y * width + x) * 3)
238         pokeByteOff p 0 r
239         pokeByteOff p 1 g
240         pokeByteOff p 2 b
241   window <- windowNew
242   eventb <- eventBoxNew
243   set window [ containerBorderWidth := 0, containerChild := eventb ]
244   set eventb [ containerBorderWidth := 0, containerChild := canvas ]
245   mapM_ (flip forkOnIO $ fpu_fix_start nullPtr) [ 0 .. numCapabilities - 1 ]
246   stop0 <- renderer' output c0 zoom0
247   sR <- newIORef (c0, zoom0, stop0)
248   _ <- eventb `on` buttonPressEvent $ {-# SCC "cb/event" #-} tryEvent $ do
249     LeftButton <- eventButton
250     (x, y) <- eventCoordinates
251     liftIO $ do
252       (c, zoom, stop) <- readIORef sR
253       stop
254       pokeArray imgdata (replicate (height * width * 3) (255 :: B))
255       let c' = c + ((convert x :+ convert (-y)) - (fromIntegral width :+ fromIntegral (-height)) * (0.5 :+ 0)) * ((1/2^^zoom) :+ 0)
256           zoom' = zoom + 1
257       stop' <- renderer' output c' zoom'
258       writeIORef sR (c', zoom', stop')
259       print (c', zoom')
260   _ <- onRealize canvas $ {-# SCC "cb/realize" #-}withGLDrawingArea canvas $ \_ -> do
261     GL.clearColor $= (GL.Color4 0.0 0.0 0.0 0.0)
262     GL.matrixMode $= GL.Projection
263     GL.loadIdentity
264     GL.ortho 0.0 1.0 0.0 1.0 (-1.0) 1.0
265     GL.drawBuffer $= GL.BackBuffers
266     [tex] <- GL.genObjectNames 1
267     GL.texture GL.Texture2D $= GL.Enabled
268     GL.textureBinding GL.Texture2D $= Just tex
269     GL.texImage2D Nothing GL.NoProxy 0 GL.RGB' (GL.TextureSize2D (fromIntegral width) (fromIntegral height)) 0 (GL.PixelData GL.RGB GL.UnsignedByte nullPtr)
270     GL.textureFilter GL.Texture2D $= ((GL.Nearest, Nothing), GL.Nearest)
271     GL.textureWrapMode GL.Texture2D GL.S $= (GL.Repeated, GL.ClampToEdge)
272     GL.textureWrapMode GL.Texture2D GL.T $= (GL.Repeated, GL.ClampToEdge)
273   _ <- onExpose canvas $ {-# SCC "cb/expose" #-} \_ -> do
274     withGLDrawingArea canvas $ \glwindow -> do
275       let v :: GLfloat -> GLfloat -> GLfloat -> GLfloat -> IO ()
276           v tx ty vx vy = GL.texCoord (GL.TexCoord2 tx ty) >> GL.vertex (GL.Vertex2 vx vy)
277           w = fromIntegral width
278           h = fromIntegral height
279       GL.clear [ GL.ColorBuffer ]
280       GL.texSubImage2D Nothing 0 (GL.TexturePosition2D 0 0) (GL.TextureSize2D w h) (GL.PixelData GL.RGB GL.UnsignedByte imgdata)
281       GL.renderPrimitive GL.Quads $ do
282         v 0 1 0 0 >> v 0 0 0 1 >> v 1 0 1 1 >> v 1 1 1 0
283       glDrawableSwapBuffers glwindow
284     return True
285   _ <- onDestroy window mainQuit
286   _ <- timeoutAdd (widgetQueueDraw canvas >> return True) 200
287   widgetShowAll window
288   mainGUI
289
290 spreadX :: [ N ]
291 spreadX = [ -1, 0, 1 ]
292
293 spreadY :: [ N ]
294 spreadY = [ -workers, 0, workers ]
295
296 workers :: N
297 workers = numCapabilities
298
299 limit :: N
300 limit = (2^(10::N)-1)
301
302 width, height :: N
303 width = 512
304 height = 512
305
306 rng :: ((N, N), (N, N))
307 rng = ((0, 0), (height - 1, width - 1))
308
309 c0 :: Complex QuadDouble
310 c0 = 0
311
312 zoom0 :: N
313 zoom0 = 6
314
315 escapeR, escapeR2 :: R
316 escapeR = 65536
317 escapeR2 = escapeR * escapeR