haskell-marcher.hs

{-# LANGUAGE FlexibleInstances #-}
import Data.Word
import Data.Bits
import Data.Kind
import Graphics.Gloss.Raster.Field
import System.Environment
import System.Exit
import Data.Char

-- Random
newtype RandM r = RandM { runRand :: Word32 -> (r, Word32) }

instance Functor RandM where
    fmap f (RandM g) = RandM $ \s -> let (v, n) = g s in (f v, n)

instance Applicative RandM where
    pure a = RandM $ \s -> (a, s)
    m1 <*> m2 = do
        f <- m1
        f <$> m2

instance Monad RandM where
    return = pure
    (RandM a) >>= f = RandM $ \s -> let (v, n) = a s in runRand (f v) n

pcgHash :: Word32 -> Word32
pcgHash input =
    (word `shiftR` 22) `xor` word
    where state = input * 747796405 + 2891336453
          word = ((state `shiftR` fromIntegral ((state `shiftR` 28) + 4)) `xor` state) * 277803737

setSeed :: Word32 -> RandM ()
setSeed s = RandM $ const ((), s)

getRand :: RandM Word32
getRand = RandM $ \s -> let n = pcgHash s in (n, n)

getRand01 :: RandM Float
getRand01 = do
    x <- getRand
    return $ fromIntegral x / 4294967295.0

-- Math utility
type Vec3 = (Float, Float, Float)

dot :: Vec3 -> Vec3 -> Float
dot (x1, y1, z1) (x2, y2, z2) = x1 * x2 + y1 * y2 + z1 * z2

mag :: Vec3 -> Float
mag (x, y, z) = sqrt $ x * x + y * y + z * z

normalize :: Vec3 -> Vec3
normalize v = v / toVec (mag v)

vecMap :: (Float -> Float) -> Vec3 -> Vec3
vecMap f (x, y, z) = (f x, f y, f z)

vecUnOp :: (Float -> Float) -> Vec3 -> Vec3
vecUnOp = fmap

vecBinOp :: (Float -> Float -> Float) -> Vec3 -> Vec3 -> Vec3
vecBinOp f (x1, y1, z1) (x2, y2, z2) = (f x1 x2, f y1 y2, f z1 z2)

fmod :: Vec3 -> Vec3 -> Vec3
fmod x y = x - y * floored
    where floored = (fromIntegral . floor) `vecMap` (x / y)

toVec :: Float -> Vec3
toVec x = (x, x, x)

toColor :: Vec3 -> Color
toColor (r, g, b) = rgb r g b

instance Num Vec3 where
    (+) = vecBinOp (+)
    (-) = vecBinOp (-)
    (*) = vecBinOp (*)
    negate = vecUnOp (0 -)
    abs = vecUnOp abs
    signum = vecUnOp signum
    fromInteger i = vecUnOp (+ fromInteger i) (0, 0, 0)

instance Fractional Vec3 where
    (/) = vecBinOp (/)
    fromRational i = vecUnOp (+ fromRational i) (0, 0, 0)

-- Math
uniformSampleSphere :: RandM Vec3
uniformSampleSphere = do
    r1 <- getRand01
    r2 <- getRand01
    let sinTheta = sqrt $ 1.0 - r1 * r1
    let phi = 2.0 * pi * r2
    let x = sinTheta * cos phi
    let z = sinTheta * sin phi
    return (x, r1, z)

uniformSampleHemisphere :: Vec3 -> RandM Vec3
uniformSampleHemisphere normal = do
    sample <- uniformSampleSphere
    if dot sample normal < 0.0
    then return $ (*2) `vecMap` sample
    else return sample

type SDF = Vec3 -> Float
sdSphere :: Float -> SDF
sdSphere radius p = mag p - radius

sdTranslate :: Vec3 -> SDF -> SDF
sdTranslate offset sdf p = sdf $ p - offset

sdScale :: Float -> SDF -> SDF
sdScale factor sdf p = (factor *) $ sdf $ (/ factor) `vecMap` p

sdRepeat :: Vec3 -> SDF -> SDF
sdRepeat period sdf p = sdf $ fmod (p + scaledPeriod) period - scaledPeriod
    where scaledPeriod = (* 0.5) `vecMap` period

sdUnion :: SDF -> SDF -> SDF
sdUnion a b p = min (a p) (b p)

sdIntersection :: SDF -> SDF -> SDF
sdIntersection a b p = max (a p) (b p)

sdDifference :: SDF -> SDF -> SDF
sdDifference a b p = max (- a p) (b p)

calcNormal :: SDF -> Vec3 -> Vec3
calcNormal sdf hit =
    normalize (x, y, z)
    where helper offset = sdf (hit + offset) - sdf (hit - offset)
          x = helper (eps, 0, 0)
          y = helper (0, eps, 0)
          z = helper (0, 0, eps)
          eps = 0.0001

maxDist = 10

march :: SDF -> Vec3 -> Vec3 -> Float
march sdf ro rd = innerMarch sdf ro 0.0
    where innerMarch sdf p t =
              let dist = sdf p in
              if dist < 0.001 || t > maxDist then t
              else
                  let t' = t + dist in
                  let p' = p + ((* dist) `vecMap` rd) in
                  innerMarch sdf p' t'

-- Main
checkHit :: SDF -> Vec3 -> Vec3 -> RandM Bool
checkHit sdf hit normal = do
    sample <- uniformSampleHemisphere normal
    return $ march scene offset sample < maxDist
    where offset = hit + vecMap (* 0.01) normal

calcAO :: SDF -> Int -> Vec3 -> RandM Color
calcAO sdf samples hit = do
    res <- calcAOInner samples 0
    let res' = res / fromIntegral samples
    return $ toColor $ toVec res'
    where normal = calcNormal sdf hit
          calcAOInner idx count =
              if idx <= 0 then return count
              else do
                  contrib <- checkHit sdf hit normal
                  let count' = if contrib then count + 1 else count
                  calcAOInner (idx - 1) count'

shade :: SDF -> Vec3 -> RandM Color
shade sdf hit = do 
    calcAO sdf 200 hit

scene :: SDF
scene = sdRepeat (toVec 2) $ sdSphere 0.5

draw :: Float -> Point -> Color
draw time (x, y) =
    let ro = (0, 0, -1) in
    let rd = normalize (x, y, 1) in
    let d = march scene ro rd in
    if d >= 10 then rgb 0 0 0
    else let hit = ro + toVec d * rd in
         fst $ runRand (shade scene hit) 1337

main :: IO ()
main = animateField (InWindow "HSMarch" (320, 320) (0, 0)) (1, 1) draw

Requires gloss-raster.