more strictness; use sqr in colour
[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 = 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 !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' :: Real c => (N -> N -> B -> B -> B -> IO ()) -> Complex c -> N -> IO (IO ())
222 renderer' output !c !zoom
223   | zoom < 48 = renderer output (f c :: Complex Double      ) zoom
224   | zoom < 96 = renderer output (f c :: Complex DoubleDouble) zoom
225   | otherwise = renderer output (f c :: Complex QuadDouble  ) zoom
226   where f !(cr :+ ci) = convert cr :+ convert ci
227
228 main :: IO ()
229 main = do
230   _ <- unsafeInitGUIForThreadedRTS
231   _ <- initGL
232   glconfig <- glConfigNew [ GLModeRGBA, GLModeDouble ]
233   canvas <- glDrawingAreaNew glconfig
234   widgetSetSizeRequest canvas width height
235   imgdata <- mallocBytes $ width * height * 3
236   pokeArray imgdata (replicate (height * width * 3) (255 :: B))
237   let output x y r g b = do
238         let p = imgdata `plusPtr` ((y * width + x) * 3)
239         pokeByteOff p 0 r
240         pokeByteOff p 1 g
241         pokeByteOff p 2 b
242   window <- windowNew
243   eventb <- eventBoxNew
244   set window [ containerBorderWidth := 0, containerChild := eventb,windowResizable := False ]
245   set eventb [ containerBorderWidth := 0, containerChild := canvas ]
246   mapM_ (flip forkOnIO $ fpu_fix_start nullPtr) [ 0 .. numCapabilities - 1 ]
247   stop0 <- renderer' output c0 zoom0
248   sR <- newIORef (c0, zoom0, stop0)
249   _ <- eventb `on` buttonPressEvent $ {-# SCC "cb/event" #-} tryEvent $ do
250     LeftButton <- eventButton
251     (x, y) <- eventCoordinates
252     liftIO $ do
253       (c, zoom, stop) <- readIORef sR
254       stop
255       pokeArray imgdata (replicate (height * width * 3) (255 :: B))
256       let c' = c + ((convert x :+ convert (-y)) - (fromIntegral width :+ fromIntegral (-height)) * (0.5 :+ 0)) * ((1/2^^zoom) :+ 0)
257           zoom' = zoom + 1
258       stop' <- renderer' output c' zoom'
259       writeIORef sR (c', zoom', stop')
260       print (c', zoom')
261   _ <- onRealize canvas $ {-# SCC "cb/realize" #-}withGLDrawingArea canvas $ \_ -> do
262     GL.clearColor $= (GL.Color4 0.0 0.0 0.0 0.0)
263     GL.matrixMode $= GL.Projection
264     GL.loadIdentity
265     GL.ortho 0.0 1.0 0.0 1.0 (-1.0) 1.0
266     GL.drawBuffer $= GL.BackBuffers
267     [tex] <- GL.genObjectNames 1
268     GL.texture GL.Texture2D $= GL.Enabled
269     GL.textureBinding GL.Texture2D $= Just tex
270     GL.texImage2D Nothing GL.NoProxy 0 GL.RGB' (GL.TextureSize2D (fromIntegral width) (fromIntegral height)) 0 (GL.PixelData GL.RGB GL.UnsignedByte nullPtr)
271     GL.textureFilter GL.Texture2D $= ((GL.Nearest, Nothing), GL.Nearest)
272     GL.textureWrapMode GL.Texture2D GL.S $= (GL.Repeated, GL.ClampToEdge)
273     GL.textureWrapMode GL.Texture2D GL.T $= (GL.Repeated, GL.ClampToEdge)
274   _ <- onExpose canvas $ {-# SCC "cb/expose" #-} \_ -> do
275     withGLDrawingArea canvas $ \glwindow -> do
276       let v :: GLfloat -> GLfloat -> GLfloat -> GLfloat -> IO ()
277           v tx ty vx vy = GL.texCoord (GL.TexCoord2 tx ty) >> GL.vertex (GL.Vertex2 vx vy)
278           w = fromIntegral width
279           h = fromIntegral height
280       GL.clear [ GL.ColorBuffer ]
281       GL.texSubImage2D Nothing 0 (GL.TexturePosition2D 0 0) (GL.TextureSize2D w h) (GL.PixelData GL.RGB GL.UnsignedByte imgdata)
282       GL.renderPrimitive GL.Quads $ do
283         v 0 1 0 0 >> v 0 0 0 1 >> v 1 0 1 1 >> v 1 1 1 0
284       glDrawableSwapBuffers glwindow
285     return True
286   _ <- onDestroy window mainQuit
287   _ <- timeoutAdd (widgetQueueDraw canvas >> return True) 200
288   widgetShowAll window
289   mainGUI
290
291 spreadX :: [ N ]
292 spreadX = [ -1, 0, 1 ]
293
294 spreadY :: [ N ]
295 spreadY = [ -workers, 0, workers ]
296
297 workers :: N
298 workers = numCapabilities
299
300 limit :: N
301 limit = (2^(10::N)-1)
302
303 width, height :: N
304 width = 512
305 height = 512
306
307 rng :: ((N, N), (N, N))
308 rng = ((0, 0), (height - 1, width - 1))
309
310 c0 :: Complex QuadDouble
311 c0 = 0
312
313 zoom0 :: N
314 zoom0 = 6
315
316 escapeR, escapeR2 :: R
317 escapeR = 65536
318 escapeR2 = escapeR * escapeR