packages feed

snowglobe-3: SnowGlobe.hs

-- SnowGlobe -- randomized fractal snowflakes demo
-- (GPL3+) 2012,2014 Claude Heiland-Allen <claude@mathr.co.uk>
-- usage: snowglobe
-- 'f' toggles full screen, 'shift-R' toggles record, any other key to quit
-- tested with: ghc-7.0.4, ghc-7.4-rc1, ghc-7.4.2, ghc-7.6.2, ghc-7.8.2

import Vector hiding (i)
import Graphics.UI.GLUT hiding (scale, FramebufferObject)
import qualified Graphics.UI.GLUT as GL
import Graphics.Rendering.OpenGL.Raw
  ( glTexImage2D, gl_TEXTURE_2D, gl_LUMINANCE32F, gl_LUMINANCE, gl_RGBA
  , gl_UNSIGNED_BYTE, gl_FALSE, glClampColor, glGenerateMipmap
  , gl_CLAMP_VERTEX_COLOR, gl_CLAMP_READ_COLOR, gl_CLAMP_FRAGMENT_COLOR
  , glGenFramebuffers, glBindFramebuffer, glFramebufferTexture2D
  , gl_FRAMEBUFFER, gl_COLOR_ATTACHMENT0, glUniformMatrix3fv, gl_ALPHA
  , glGetTexImage
  )
import Control.Exception (evaluate)
import Control.Monad (forM_, replicateM, when)
import Data.IORef (IORef, modifyIORef, newIORef, readIORef, writeIORef)
import Data.List (foldl', minimumBy)
import Data.Ord (comparing)
import Data.Map (Map)
import qualified Data.Map as M
import qualified Data.Set as S
import System.Exit (exitSuccess)
import System.IO (hPutStrLn, stderr, stdout, hPutStr, withBinaryFile, IOMode(WriteMode), hPutBuf)
import System.Random (randomRIO)
import Unsafe.Coerce (unsafeCoerce) -- for OpenGL uniform newtype
import Foreign (alloca, peek, nullPtr, withArray)
import Graphics.Rendering.OpenGL.Capture (capturePPM)
import Data.ByteString as BS (hPut, pack)
import Foreign (Ptr, allocaArray)
import Data.Word (Word8)

shader :: Maybe String -> Maybe String -> IO Program
shader mV mF = do
  p <- createProgram
  vs <- case mV of
    Nothing -> return []
    Just v -> do
      vert <- createShader VertexShader
      shaderSourceBS vert $= BS.pack (map (toEnum . fromEnum) v)
      compileShader vert
      msg <- get (shaderInfoLog vert)
      when (not (null msg)) $ hPutStrLn stderr msg
      return [vert]
  fs <- case mF of
    Nothing -> return []
    Just f -> do
      frag <- createShader FragmentShader
      shaderSourceBS frag $= BS.pack (map (toEnum . fromEnum) f)
      compileShader frag
      msg <- get (shaderInfoLog frag)
      when (not (null msg)) $ hPutStrLn stderr msg
      return [frag]
  attachedShaders p $= concat [vs, fs]
  linkProgram p
  msg <- get (programInfoLog p)
  when (not (null msg)) $ hPutStrLn stderr msg
  return p

type N = Int

sum' :: Num a => [a] -> a
sum' = foldl' (+) 0

size :: N
size = 2048

er :: R
er = 4

accuracy :: R
accuracy = 1 / fromIntegral size

tSize :: N
tSize = 1024

maxSpeed :: R
maxSpeed = 10

data SnowGlobe = SnowGlobe
  { pInitial :: Program, pInitial'er, pInitial'rho :: UniformLocation
  , pStep :: Program, pStep'er, pStep'ts, pStep'src :: UniformLocation
  , pColour :: Program, pColour'src, pColour'speed, pColour'colour :: UniformLocation
  , tPing, tPong :: TextureObject, fBuffer :: FramebufferObject
  , sFlakes :: [Flake], sTextures :: Map N TextureObject
  , wSize :: Size, wFullScreen :: Maybe Size
  , sRenderFlake :: IO Render, sNextName :: N
  , sRecord, sRecordFlakes :: Bool
  }

data Flake = Flake
  { flakeName :: !N
  , flakeTime :: !R
  , flakePosition :: !V2
  , flakeVelocity :: !V2
  }

flakeMass :: Flake -> R
flakeMass f = let s = flakeSize f in 1 + s * s

flakeSize :: Flake -> R
flakeSize f = sin (pi * flakeTime f)

flakeForce :: Flake -> V2 -> V2
flakeForce f p =
  let d@(V2 x y) = p ^-^ flakePosition f
      m = let d2 = x * x + y * y in d2 * d2
  in  (flakeMass f / m) *^ d

flakeUpdate :: R -> [Flake] -> Flake -> Flake
flakeUpdate dt fs f = f
  { flakeTime = flakeTime f + dt / 2
  , flakePosition = flakePosition f ^+^ (dt *^ flakeVelocity f)
  , flakeVelocity = mapVector ((* maxSpeed) . tanh . (/ maxSpeed)) $ 0.999 *^ (flakeVelocity f ^+^ ((dt / flakeMass f) *^ (flakeField fs (flakePosition f))))
  }

mapVector :: (R -> R) -> V2 -> V2
mapVector f (V2 x y) = V2 (f x) (f y)

flakeField :: [Flake] -> V2 -> V2
flakeField fs p = foldl' (^+^) o [ flakeForce f (p ^+^ d) | f <- fs, p /= flakePosition f, dx <- [-2,0,2], dy <- [-2,0,2], let d = V2 dx dy ]

flakeSpawn :: [Flake] -> N -> IO Flake
flakeSpawn fs name = do
  xs <- replicateM 32 $ randomRIO (-1, 1)
  ys <- replicateM 32 $ randomRIO (-1, 1)
  let ps = zipWith V2 xs ys
      p = fst . minimumBy (comparing snd) . map (\p' -> (p', norm2 (flakeField fs p'))) $ ps
  return Flake{ flakeName = name, flakeTime = 0, flakePosition = p, flakeVelocity = o }
  where
    norm2 (V2 x y) = x * x + y * y

flakesUpdate :: R -> [Flake] -> [Flake]
flakesUpdate dt fs =
  let gs = map (flakeUpdate dt fs) fs
      alive = filter ((< 1) . flakeTime) gs
  in  map flakeWrap alive

flakeWrap :: Flake -> Flake
flakeWrap f = f{ flakePosition = mapVector wrap (flakePosition f) }
  where
    wrap x =
      let y = (x + 1) / 2
          z = y - fromIntegral (floor y :: N)
      in  z * 2 - 1

-- amortized rendering over several frames
data Render = Done TextureObject | Step (IO Render)

-- initial pass
flakeRenderStart :: IORef SnowGlobe -> IO Render
flakeRenderStart sR = do
  s0 <- readIORef sR
  p <- replicateM 4 $ randomRIO (0.02, 0.98)
  let rts = flakeTransforms p
      rho = maximum (map fst rts)
      ts = map snd rts
      passes = clamp 4 256 . round . logBase rho $ accuracy
  loadIdentity
  ortho2D 0 1 0 1
  viewport $= (Position 0 0, Size (fromIntegral size) (fromIntegral size))
  currentProgram $= Just (pInitial s0)
  uniform (pInitial'er s0) $= TexCoord1 (realToFrac er :: GLfloat)
  uniform (pInitial'rho s0) $=  TexCoord1 (realToFrac rho :: GLfloat)
  bindFBO (fBuffer s0) (tPing s0)
  unitQuad
  unbindFBO
  currentProgram $= Nothing
  return $ Step (flakeRenderPass sR ts passes passes)

-- multi step passes
flakeRenderPass :: IORef SnowGlobe -> [M3] -> N -> N -> IO Render
flakeRenderPass sR _  passes 0 = flakeRenderFinish sR passes
flakeRenderPass sR ts passes n = do
  s0 <- readIORef sR
  loadIdentity
  ortho2D 0 1 0 1
  viewport $= (Position 0 0, Size (fromIntegral size) (fromIntegral size))
  currentProgram $= Just (pStep s0)
  uniform (pStep'er s0) $= TexCoord1 (realToFrac er :: GLfloat)
  uniform (pStep'src s0) $= TexCoord1 (0 :: GLint)
  withArray (map ((realToFrac :: Float -> GLfloat) . (realToFrac :: Double -> Float)) . concatMap matrixToList $ ts) $ glUniformMatrix3fv (unsafeCoerce $ pStep'ts s0) 6 1
  bindFBO (fBuffer s0) (tPong s0)
  textureBinding Texture2D $= Just (tPing s0)
  unitQuad
  textureBinding Texture2D $= Nothing
  unbindFBO
  writeIORef sR s0{ tPing = tPong s0, tPong = tPing s0 }
  currentProgram $= Nothing
  if even n
    then return $ Step (flakeRenderPass sR ts passes (n - 1))
    else                flakeRenderPass sR ts passes (n - 1)

matrixToList :: M3 -> [R]
matrixToList (M3 a b c d e f g h i) = [a,b,c,d,e,f,g,h,i]

-- colourize pass
flakeRenderFinish :: IORef SnowGlobe -> N -> IO Render
flakeRenderFinish sR passes = do
  s0 <- readIORef sR
  t <- newTexRGBA tSize
  bindFBO (fBuffer s0) t
  loadIdentity
  ortho2D (-1) 1 (-1) 1
  viewport $= (Position 0 0, Size (fromIntegral tSize) (fromIntegral tSize))
  textureBinding Texture2D $= Just (tPong s0)
  currentProgram $= Just (pColour s0)
  uniform (pColour'src s0) $= TexCoord1 (0 :: GLint)
  uniform (pColour'speed s0) $= TexCoord1 (1 / fromIntegral passes :: GLfloat)
  uniform (pColour'colour s0) $= TexCoord3 1 1 (1 :: GLfloat)
  fullQuad
  currentProgram $= Nothing
  unbindFBO
  textureBinding Texture2D $= Just t
  glGenerateMipmap gl_TEXTURE_2D
  textureBinding Texture2D $= Nothing
  return $ Done t

transformRST :: R -> R -> V2 -> M3
transformRST a l (V2 x y) = M3 c s x  (-s) c y  0 0 1
  where c = l * cos a
        s = l * sin a

flakeTransforms :: [R] -> [(R, M3)]
flakeTransforms [a,b,c,d] = [(la,inv ta),(lb,inv tb),(lc,inv tc1),(ld,inv td1),(lc,inv tc2),(ld,inv td2)]
  where
    lx = a + b
    ly = 2 * (c + d)
    la  = a / lx
    lb  = b / lx
    lc  = c / ly
    ld  = d / ly
    u   = V2 0 0
    v   = V2 0 la
    l   =     pi / 3
    ll  = 2 * pi / 3
    r   = -l
    rr  = -ll
    ta  = transformRST  0 la u
    tb  = transformRST  0 lb v
    tc1 = transformRST  l lc v
    tc2 = transformRST  r lc v
    td1 = transformRST ll ld v
    td2 = transformRST rr ld v
flakeTransforms _ = error "flakeTransforms"

main :: IO ()
main = do
  _ <- getArgsAndInitialize
  let wSize' = Size 1280 720
  initialWindowSize $= wSize'
  initialDisplayMode $= [DoubleBuffered]
  _ <- createWindow "SnowGlobe"
  pInitial'     <- shader Nothing (Just $ unlines
    [ "uniform float er;"
    , "uniform float rho;"
    , ""
    , "void main() {"
    , "  vec2 p = er * (gl_TexCoord[0].xy * 2.0 - vec2(1.0));"
    , "  float l = length(p);"
    , "  float n;"
    , "  if (l >= er) {"
    , "    n = 0.0;"
    , "  } else if (er > l && l >= rho * er) {"
    , "    n = (log(er) - log(l)) / -log(rho);"
    , "  } else {"
    , "    n = -1.0;"
    , "  }"
    , "  gl_FragData[0] = vec4(n);"
    , "}"
    ])
  pInitial'er'  <- get $ uniformLocation pInitial' "er"
  pInitial'rho' <- get $ uniformLocation pInitial' "rho"
  pStep'        <- shader Nothing (Just $ unlines
    [ "uniform float er;"
    , "uniform mat3 ts[6];"
    , ""
    , "uniform sampler2D src;"
    , ""
    , "void main() {"
    , "  vec2 p0 = er * (gl_TexCoord[0].xy * 2.0 - vec2(1.0));"
    , "  float m = -1.0;"
    , "  for (int i = 0; i < 6; ++i) {"
    , "    vec3 p = ts[i] * vec3(p0, 1.0);"
    , "    vec2 q = p.xy / p.z;"
    , "    float l = length(q);"
    , "    if (l < er) {"
    , "      m = max(m, texture2D(src, (q / er + vec2(1.0)) / 2.0).x);"
    , "    }"
    , "  }"
    , "  if (m >= 0.0) {"
    , "    m += 1.0;"
    , "  }"
    , "  m = max(m, texture2D(src, (p0 / er + vec2(1.0)) / 2.0).x);"
    , "  gl_FragData[0] = vec4(m);"
    , "}"
    ])
  pStep'er'     <- get $ uniformLocation pStep' "er"
  pStep'ts'     <- get $ uniformLocation pStep' "ts"
  pStep'src'    <- get $ uniformLocation pStep' "src"
  pColour'      <- shader Nothing (Just $ unlines
    [ "uniform sampler2D src;"
    , "uniform float speed;"
    , "uniform vec3 colour;"
    , ""
    , "void main() {"
    , "  vec2 p = gl_TexCoord[0].xy;"
    , "  float n = texture2D(src, p).x;"
    , "  p -= vec2(0.5);"
    , "  const mat2 r = mat2(0.5, 0.8660254037844386, -0.8660254037844386, 0.5);"
    , "  for (int i = 1; i < 6; ++i) {"
    , "    p = r * p;"
    , "    n = max(n, texture2D(src, p + vec2(0.5)).x);"
    , "  }"
    , "  if (n > 0.0) {"
    , "    n *= speed;"
    , "  } else {"
    , "    n = 0.0;"
    , "  }"
    , "  n *= n;"
    , "  n *= n;"
    , "  gl_FragData[0] = vec4(colour, n);"
    , "}"
    ])
  pColour'src'  <- get $ uniformLocation pColour' "src"
  pColour'speed'<- get $ uniformLocation pColour' "speed"
  pColour'colour'<- get $ uniformLocation pColour' "colour"
  tPing'        <- newTex size
  tPong'        <- newTex size
  fBuffer'      <- newFBO
  glClampColor gl_CLAMP_VERTEX_COLOR   gl_FALSE
  glClampColor gl_CLAMP_READ_COLOR     gl_FALSE
  glClampColor gl_CLAMP_FRAGMENT_COLOR gl_FALSE
  sR <- newIORef SnowGlobe
    { pInitial = pInitial', pInitial'er = pInitial'er', pInitial'rho = pInitial'rho'
    , pStep = pStep', pStep'er = pStep'er', pStep'ts = pStep'ts', pStep'src = pStep'src'
    , pColour = pColour', pColour'src = pColour'src', pColour'speed = pColour'speed', pColour'colour = pColour'colour'
    , tPing = tPing', tPong = tPong', fBuffer = fBuffer'
    , sFlakes = [], sTextures = M.empty, wSize = wSize', wFullScreen = Nothing
    , sNextName = 0, sRenderFlake = return undefined
    , sRecord = False, sRecordFlakes = False
    }
  modifyIORef sR $ \s'->s'{ sRenderFlake = flakeRenderStart sR }
  addTimerCallback 40 timer
  displayCallback $= display sR
  reshapeCallback $= Just (reshape sR)
  keyboardMouseCallback $= Just (keyboard sR)
  mainLoop

keyboard :: IORef SnowGlobe -> Key -> KeyState -> Modifiers -> Position -> IO ()
keyboard sR (SpecialKey KeyF11) Down _ _ = toggleFullScreen sR
keyboard sR (Char 'f') Down _ _ = toggleFullScreen sR
keyboard sR (Char 'R') Down _ _ = modifyIORef sR $ \s -> s{ sRecord = not (sRecord s) }
keyboard sR (Char 'S') Down _ _ = modifyIORef sR $ \s -> s{ sRecordFlakes = not (sRecordFlakes s) }
keyboard _ (Char _) Down _ _ = exitSuccess
keyboard _ _ _ _ _ = return ()

toggleFullScreen :: IORef SnowGlobe -> IO ()
toggleFullScreen sR = do
  s <- readIORef sR
  case wFullScreen s of
    Nothing -> do
      writeIORef sR s{ wFullScreen = Just (wSize s) }
      cursor $= None
      fullScreen
    Just ws -> do
      writeIORef sR s{ wFullScreen = Nothing }
      cursor $= Inherit
      windowSize $= ws

reshape :: IORef SnowGlobe -> Size -> IO ()
reshape sR sz = do
  s <- readIORef sR
  writeIORef sR s{ wSize = sz }

timer :: IO ()
timer = do
  addTimerCallback 40 timer
  postRedisplay Nothing

display' :: IORef SnowGlobe -> IO ()
display' sR = do
  update sR
  s <- readIORef sR
  _ <- evaluate (sum' . map flakeName . sFlakes $ s)
  let names = S.fromList $ map flakeName (sFlakes s)
      expired = S.filter (`S.notMember` names) (M.keysSet (sTextures s))
      sTextures' = foldr M.delete (sTextures s) (S.toList expired)
  deleteObjectNames [sTextures s M.! n | n <- S.toList expired]
  modifyIORef sR $ \s'->s'{ sTextures = sTextures' }
  r <- sRenderFlake s
  case r of
    Done t -> do
      f <- flakeSpawn (sFlakes s) (sNextName s)
      modifyIORef sR $ \s'->s'
        { sRenderFlake = flakeRenderStart sR
        , sFlakes = f : sFlakes s'
        , sTextures = M.insert (flakeName f) t (sTextures s')
        , sNextName = sNextName s' + 1
        }
      when (sRecordFlakes s) $
        saveTexture t ("snowglobe-" ++ show (flakeName f) ++ ".pgm")
    Step sRenderFlake' ->
      modifyIORef sR $ \s'->s'{ sRenderFlake = sRenderFlake' }

update :: IORef SnowGlobe -> IO ()
update sR = do
  s <- readIORef sR
  let sFlakes' = flakesUpdate (1 / 256) (sFlakes s)
  writeIORef sR (s{ sFlakes = sFlakes' })

display :: IORef SnowGlobe -> IO ()
display sR = do
  s <- readIORef sR
  loadIdentity
  let Size w h = wSize s
      r = 0.7
      (x, y)
        | h < w     = (r, r * fromIntegral h / fromIntegral w)
        | otherwise = (r * fromIntegral w / fromIntegral h, r)
  ortho2D (-x) x (-y) y
  viewport $= (Position 0 0, wSize s)
  clearColor $= Color4 0 0 0.25 1
  clear [ColorBuffer]
  texture Texture2D $= Enabled
  blend $= Enabled
  blendFunc $= (SrcAlpha, OneMinusSrcAlpha)
  forM_ (sFlakes s) (flakeDraw s)
  blend $= Disabled
  texture Texture2D $= Disabled
  swapBuffers
  when (sRecord s) $ hPut stdout =<< capturePPM
  reportErrors
  display' sR

flakeDraw :: SnowGlobe -> Flake -> IO ()
flakeDraw s f = do
  let d :: GLdouble
      d = realToFrac $ flakeSize f / 4
      a :: GLdouble
      a = realToFrac $ 360 * flakeTime f * sin (fromIntegral (flakeName f))
      p :: GLdouble -> GLdouble -> IO ()
      p u v = do
        texCoord $ TexCoord2 ((1+u)/2) ((1+v)/2)
        vertex $ Vertex2 u v
      V2 x y = flakePosition f
  case flakeName f `M.lookup` sTextures s of
    Nothing -> return ()
    t -> do
      textureBinding Texture2D $= t
      unsafePreservingMatrix $ do
        translate $ Vector3 (realToFrac x :: GLdouble) (realToFrac y :: GLdouble) 0
        rotate a (Vector3 0 0 1)
        GL.scale d d d
        renderPrimitive Quads $ p (-1) (-1) >> p 1 (-1) >> p 1 1 >> p (-1) 1
      textureBinding Texture2D $= Nothing

newTex :: N -> IO TextureObject
newTex s = do
  [t] <- genObjectNames 1
  textureBinding Texture2D $= Just t
  glTexImage2D gl_TEXTURE_2D 0 (fromIntegral gl_LUMINANCE32F) (fromIntegral s) (fromIntegral s) 0 gl_LUMINANCE gl_UNSIGNED_BYTE nullPtr
  textureFilter Texture2D $= ((Linear', Nothing), Linear')
  textureWrapMode Texture2D S $= (Repeated, ClampToEdge)
  textureWrapMode Texture2D T $= (Repeated, ClampToEdge)
  textureBinding Texture2D $= Nothing
  return t

newTexRGBA :: N -> IO TextureObject
newTexRGBA s = do
  [t] <- genObjectNames 1
  textureBinding Texture2D $= Just t
  glTexImage2D gl_TEXTURE_2D 0 (fromIntegral gl_RGBA) (fromIntegral s) (fromIntegral s) 0 gl_RGBA gl_UNSIGNED_BYTE nullPtr
  textureFilter Texture2D $= ((Linear', Just Linear'), Linear')
  textureWrapMode Texture2D S $= (Repeated, ClampToEdge)
  textureWrapMode Texture2D T $= (Repeated, ClampToEdge)
  textureBinding Texture2D $= Nothing
  return t

saveTexture :: TextureObject -> FilePath -> IO ()
saveTexture t f = withBinaryFile f WriteMode $ \h -> do
  let header = "P5\n" ++ show tSize ++ " " ++ show tSize ++ "\n255\n"
      n = tSize * tSize
  hPutStr h header
  allocaArray n $ \p -> do
    textureBinding Texture2D $= Just t
    glGetTexImage gl_TEXTURE_2D 0 gl_ALPHA gl_UNSIGNED_BYTE (p :: Ptr Word8)
    textureBinding Texture2D $= Nothing
    hPutBuf h p n

newtype FramebufferObject = FramebufferObject GLuint

newFBO :: IO FramebufferObject
newFBO = fmap FramebufferObject (alloca $ \p -> glGenFramebuffers 1 p >> peek p)

bindFBO :: FramebufferObject -> TextureObject -> IO ()
bindFBO (FramebufferObject f) (TextureObject t) = do
  glBindFramebuffer gl_FRAMEBUFFER f
  glFramebufferTexture2D gl_FRAMEBUFFER gl_COLOR_ATTACHMENT0 gl_TEXTURE_2D t 0

unbindFBO :: IO ()
unbindFBO = do
  glFramebufferTexture2D gl_FRAMEBUFFER gl_COLOR_ATTACHMENT0 gl_TEXTURE_2D 0 0
  glBindFramebuffer gl_FRAMEBUFFER 0

fullQuad :: IO ()
fullQuad = do
  renderPrimitive Quads $ do
    t (0.5-r) (0.5+r) >> v (-r2) ( r2)
    t (0.5-r) (0.5-r) >> v (-r2) (-r2)
    t (0.5+r) (0.5-r) >> v ( r2) (-r2)
    t (0.5+r) (0.5+r) >> v ( r2) ( r2)
  where
    r = sqrt 0.5 / realToFrac er
    r2 = 1
    t, v :: GLdouble -> GLdouble -> IO ()
    t x y = texCoord (TexCoord2 x y)
    v x y = vertex (Vertex2 x y)

unitQuad :: IO ()
unitQuad = renderPrimitive Quads $ do
    t 0 1 >> v 0 1
    t 0 0 >> v 0 0
    t 1 0 >> v 1 0
    t 1 1 >> v 1 1
  where
    t, v :: GLdouble -> GLdouble -> IO ()
    t x y = texCoord (TexCoord2 x y)
    v x y = vertex (Vertex2 x y)

clamp :: Ord a => a -> a -> a -> a
clamp mi ma x = mi `max` x `min` ma