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