raytracer in haskell

sphtrace.hs

{-# LANGUAGE BangPatterns #-}

import Control.Applicative
import System.IO
import Data.Maybe
import Data.Complex
import Data.Word
import qualified Data.ByteString as B

data Vec a = Vec { vecx, vecy, vecz :: !a }

instance Functor Vec where
   fmap f (Vec x y z) = Vec (f x) (f y) (f z)

instance Applicative Vec where
   pure x = Vec x x x
   (Vec f g h) <*> (Vec x y z) = Vec (f x) (g y) (h z)

instance Foldable Vec where
   foldr f a (Vec x y z) = f x $ f y $ f z a
   foldl1 f (Vec x y z) = f z $ f y x

instance Num a => Num (Vec a) where
   (+) = liftA2 (+)
   (*) = liftA2 (*)
   (-) = liftA2 (-)
   abs v = abs <$> v
   signum v = signum <$> v
   fromInteger i = pure $ fromInteger i

dot :: Num a => Vec a -> Vec a -> a
dot u v = foldl1 (+) $ u * v

cross :: Num a => Vec a -> Vec a -> Vec a
cross (Vec ux uy uz) (Vec vx vy vz) = Vec (uy * vz - uz * vy) (uz * vx - ux * vz) (ux * vy - uy * vx)

normsq :: Num a => Vec a -> a
normsq v = v `dot` v

norm :: Floating a => Vec a -> a
norm v = sqrt $ normsq v

normalize :: Floating a => Vec a -> Vec a
normalize v = fmap (/ len) v
   where len = norm v

transpose :: Vec a -> Vec a -> Vec a -> Vec (Vec a)
transpose u v w = Vec x y z
   where x = Vec (vecx u) (vecx v) (vecx w)
         y = Vec (vecy u) (vecy v) (vecy w)
         z = Vec (vecz u) (vecz v) (vecz w)

data Camera a = Cam {
   cam_size :: (Int, Int),
   cam_pos :: Vec a,
   cam_scale :: (a, a),
   cam_aspect, cam_cz :: a,
   cam_dir :: Vec (Vec a)
}

calc_camera :: Floating a => Int -> Int -> a -> Vec a -> Vec a -> Vec a -> Camera a
calc_camera w h fov pos center up = Cam (w, h) pos (scale_x, scale_y) aspect cz (transpose di dj dk)
   where (w', h') = (fromIntegral w, fromIntegral h)
         (scale_x, scale_y) = (2 / (w' - 1), 2 / (h' - 1))
         aspect = h' / w'
         cz = 1 / tan(fov * pi / 360)
         dk = normalize $ center - pos
         di = normalize $ dk `cross` up
         dj = di `cross` dk

trace_ray :: RealFloat a => (Vec a -> a) -> Int -> Int -> a -> Camera a -> Maybe (Vec a)
trace_ray dist_fn sx sy max_t cam = ray_loop 0
   where pos = cam_pos cam
         (scale_x, scale_y) = cam_scale cam
         cx = fromIntegral sx * scale_x - 1
         cy = (1 - fromIntegral sy * scale_y) * (cam_aspect cam)
         c = Vec cx cy (cam_cz cam)
         m = normalize $ dot c <$> cam_dir cam
         ray_loop t
            | t > max_t = Nothing
            | d < eps = Just $! r
            | otherwise = ray_loop (t + d)
            where r = fmap (*t) m + pos
                  d = dist_fn r
                  eps = 0.001

data Light a = Light { light_pos :: Vec a, light_val :: a }

trace_shadow :: RealFloat a => (Vec a -> a) -> Light a -> Vec a -> a -> Maybe a
trace_shadow dist_fn light pos k = ray_loop 0.02 1
   where lpos = light_pos light
         dp = lpos - pos
         dist = norm dp
         m = (/ dist) <$> dp
         ray_loop t acc
            | t > dist = Just $! s
            | d < eps = Nothing
            | otherwise = ray_loop (t + d) s
            where r = fmap (*t) m + pos
                  d = dist_fn r
                  s = min acc $ k * d / t
                  eps = 0.001

calc_diffuse :: RealFloat a => Light a -> Vec a -> Vec a -> a
calc_diffuse light pos snorm = col * il * int / dsq
   where col = light_val light
         d = light_pos light - pos
         dsq = normsq d
         int = max 0 $ snorm `dot` d
         il = 1 / sqrt dsq

gradient :: Floating a => (Vec a -> a) -> Vec a -> Vec a
gradient f p = normalize $ Vec dx dy dz
   where dx = f (p + epsx) - f (p - epsx)
         dy = f (p + epsy) - f (p - epsy)
         dz = f (p + epsz) - f (p - epsz)
         (epsx, epsy, epsz) = (Vec eps 0 0, Vec 0 eps 0, Vec 0 0 eps)
         eps = 1e-6

data Image = Image Int Int [Word8]

render :: RealFloat a => (Vec a -> a) -> Camera a -> Light a -> a -> a -> Image
render scene cam light ambient max_t = Image w h [pixel x y | y <- [0 .. h-1], x <- [0 .. w-1]]
   where (w, h) = cam_size cam
         clamp mn mx = max mn . min mx
         pixel sx sy = fromMaybe 0 val
            where val = do
                  hit_pos <- trace_ray scene sx sy max_t cam
                  shadow <- trace_shadow scene light hit_pos 8
                  let diffuse = calc_diffuse light hit_pos $ gradient scene hit_pos
                  let color = clamp 0 1 $ ambient + diffuse * shadow
                  return . truncate $ color * 255

write_ppm :: String -> Image -> IO ()
write_ppm filename (Image w h pixels) = do
   handle <- openBinaryFile filename WriteMode
   hPutStr handle $ "P5\n" ++ show w ++ " " ++ show h ++ "\n255\n"
   B.hPut handle $ B.pack pixels
   hClose handle

-- Scene definition
translate :: RealFloat a => (a, a, a) -> (Vec a -> a) -> Vec a -> a
translate (dx, dy, dz) f p = f (p - Vec dx dy dz)

mandel_dist :: RealFloat a => Int -> Complex a -> a
mandel_dist max_it c = mandel_loop 0 1 0
   where mandel_loop !z !dz it
            | it >= max_it = 0
            | zsq > 1024 = dist
            | otherwise = mandel_loop z' dz' (it + 1)
            where z' = z * z + c
                  dz' = z * dz * 2 + 1
                  zsq = sqr z'
                  dist = sqrt (zsq / sqr dz') * log zsq * 0.5
                  sqr (re :+ im) = re * re + im * im

scene :: RealFloat a => Vec a -> a
scene p = minimum [translate (2, -0.5, 0) (sphere 0.6) p,
                   translate (3.2, 0, 0.5) (cone (pi/6)) p,
                   translate (-1, 0, 0) (torus 1.5 0.3) p,
                   max (box (Vec 1 1 1) p) (-sphere 1.3 p),
                   max (translate (0, -1, 0) (cylinder_xy 0.6) p) (translate (0, -1, 2.5) (cylinder_yz 0.6) p),
                   max (translate (2.1, 0, 1.5) mandel p) (plane (Vec 0 1 0) 0.6 p),
                   plane (Vec 1 0 0) 2.1 p, plane (Vec 0 1 0) 1 p, plane (Vec 0 0 1) 1.5 p]
   where sphere r p = norm p - r
         torus r1 r2 (Vec x y z) = let q = sqrt (x*x + z*z) - r1 in sqrt (q*q + y*y) - r2
         plane n b p = p `dot` n + b
         cone t (Vec x y z) = let q = sqrt (x*x + z*z) in q * cos(t) + y * sin(t)
         box b p = let d = abs p - b in min 0 (maximum d) + norm (fmap (max 0) d)
         cylinder_xy r (Vec x y _) = sqrt (x*x + y*y) - r
         cylinder_yz r (Vec _ y z) = sqrt (y*y + z*z) - r
         mandel (Vec re _ im) = mandel_dist 256 (re :+ im)

main = write_ppm "out.ppm" $ render scene cam light 0.03 42
   where cam = calc_camera 640 480 60 (Vec 3 3 5) (Vec 0 (-1) 0) (Vec 0 1 0)
         light = Light (Vec 3 2 2) 7