never executed always true always false
    1 {-|
    2 Module: Flaw.Graphics.Program
    3 Description: Shader program support.
    4 License: MIT
    5 -}
    6 
    7 {-# LANGUAGE TemplateHaskell #-}
    8 
    9 module Flaw.Graphics.Program
   10   ( Program
   11   , OfScalarType(..)
   12   , OfAttributeType(..)
   13   , AttributeFormat(..)
   14   , AttributeType(..)
   15   , Normalization(..)
   16   , Node
   17   , cnst
   18   , constf, const2f, const3f, const4f
   19   , cvec11, cvec111, cvec12, cvec21, cvec1111, cvec112, cvec121, cvec211, cvec22, cvec13, cvec31
   20   , cast
   21   , attribute
   22   , attributeWithType
   23   , UniformBufferSlot
   24   , UniformStorage
   25   , uniformBufferSlot
   26   , uniform
   27   , uniformArray
   28   , createUniformStorage
   29   , setUniform
   30   , renderUniform
   31   , renderIndexedUniform
   32   , renderUploadUniformStorage
   33   , renderUniformStorage
   34   , sampler
   35   , sampler1Df, sampler1D2f, sampler1D3f, sampler1D4f
   36   , sampler2Df, sampler2D2f, sampler2D3f, sampler2D4f
   37   , sampler3Df, sampler3D2f, sampler3D3f, sampler3D4f
   38   , samplerCubef, samplerCube2f, samplerCube3f, samplerCube4f
   39   , sample, sampleOffset, sampleLod, sampleLodOffset, sampleBias, sampleBiasOffset, sampleGrad, sampleGradOffset
   40   , temp
   41   , rasterize
   42   , colorTarget
   43   , dualColorTarget
   44   , depthTarget
   45   , (!)
   46   , min_, max_
   47   , clamp, lerp
   48   , equal_ , less_, lessEqual_, greater_, greaterEqual_, if_
   49   , ddx, ddy
   50   , floor_
   51   , instanceId
   52   , invSqrt
   53   , screenToTexture
   54   , normalizeSampledDepth
   55   , fragCoord
   56   ) where
   57 
   58 import Control.Monad.Reader
   59 import qualified Data.ByteString.Unsafe as B
   60 import Data.Char
   61 import Data.Int
   62 import Data.IORef
   63 import Data.Word
   64 import Foreign.ForeignPtr
   65 import Foreign.ForeignPtr.Unsafe
   66 import Foreign.Ptr
   67 import Foreign.Storable
   68 import Language.Haskell.TH
   69 
   70 import Flaw.Graphics
   71 import Flaw.Graphics.Program.Internal
   72 import Flaw.Math
   73 
   74 cnst :: OfValueType a => a -> Node a
   75 cnst value = ConstNode (valueType value) value
   76 
   77 -- | Helper method to know return value.
   78 withUndefined :: (a -> Node a) -> Node a
   79 withUndefined q = q undefined
   80 
   81 -- | Helper method to know return value in monad.
   82 withUndefinedM :: (a -> m (Node a)) -> m (Node a)
   83 withUndefinedM q = q undefined
   84 
   85 -- | Create vector as a combination of scalars/vectors.
   86 fmap concat $ forM
   87   [ [1, 1]
   88   , [1, 1, 1], [1, 2], [2, 1]
   89   , [1, 1, 1, 1], [1, 1, 2], [1, 2, 1], [2, 1, 1], [2, 2], [1, 3], [3, 1]]
   90   $ \cs -> do
   91   ps <- forM (zip cs [1..]) $ \(_c, i) -> newName ['p', intToDigit i]
   92   tvA <- newName "a"
   93   u <- newName "u"
   94   let
   95     funName = mkName $ "cvec" ++ map intToDigit cs
   96     vecType n = appT (conT $ mkName $ "Vec" ++ [intToDigit n]) (varT tvA)
   97     argType n = [t| Node $(if n > 1 then vecType n else varT tvA) |]
   98     funType = forallT [PlainTV tvA] (sequence [ [t| OfScalarType $(varT tvA) |], [t| Vectorized $(varT tvA) |] ]) $
   99       foldr (\a b -> [t| $a -> $b |]) [t| Node $(vecType $ sum cs) |] $ map argType cs
  100     construction = foldl appE (conE $ mkName $ "Combine" ++ [intToDigit $ length cs] ++ "VecNode") $
  101       (map (\a -> [| nodeValueType $(varE a) |]) ps) ++ [ [| valueType $(varE u) |] ] ++ (map varE ps)
  102   sequence
  103     [ sigD funName funType
  104     , funD funName [clause (map varP ps) (normalB [| withUndefined $ \ $(varP u) -> $construction |]) []]
  105     ]
  106 
  107 -- | Cast value to other type.
  108 cast :: (OfValueType a, OfValueType b) => Node a -> Node b
  109 cast a = withUndefined $ \u -> CastNode (nodeValueType a) (valueType u) a
  110 
  111 -- | Define vertex attribute using typed format.
  112 attribute :: OfAttributeType a
  113   => Int -- ^ Slot.
  114   -> Int -- ^ Offset.
  115   -> Int -- ^ Divisor.
  116   -> AttributeFormat a -- ^ Format.
  117   -> Program (Node a)
  118 attribute slot offset divisor format = attributeWithType slot offset divisor $ attributeFormatToType format
  119 
  120 -- | Define vertex attribute using untyped 'AttributeType'.
  121 attributeWithType :: OfAttributeType a
  122   => Int -- ^ Slot.
  123   -> Int -- ^ Offset.
  124   -> Int -- ^ Divisor.
  125   -> AttributeType -- ^ Attribute type.
  126   -> Program (Node a)
  127 attributeWithType slot offset divisor at = withUndefinedM $ \u -> withState $ \state@State
  128   { stateStage = stage
  129   } -> do
  130   if stage /= VertexStage then fail "attribute can only be defined in vertex program"
  131   else return ()
  132   tempInternal (AttributeNode Attribute
  133     { attributeSlot = slot
  134     , attributeOffset = offset
  135     , attributeDivisor = divisor
  136     , attributeType = at
  137     , attributeValueType = valueType u
  138     }) state
  139 
  140 data UniformBufferSlot = UniformBufferSlot
  141   { uniformBufferSlotIndex :: Int
  142   , uniformBufferSlotSizeRef :: IORef Int
  143   }
  144 
  145 -- | Helper object for uniform buffer.
  146 data UniformStorage d = UniformStorage
  147   { uniformStorageSlot :: Int
  148   , uniformStorageBufferId :: UniformBufferId d
  149   , uniformStorageBytes :: ForeignPtr ()
  150   , uniformStorageSize :: Int
  151   }
  152 
  153 uniformBufferSlot :: Int -> IO UniformBufferSlot
  154 uniformBufferSlot slot = do
  155   sizeRef <- newIORef 0
  156   return UniformBufferSlot
  157     { uniformBufferSlotIndex = slot
  158     , uniformBufferSlotSizeRef = sizeRef
  159     }
  160 
  161 -- | Modified alignment calculation for shaders.
  162 -- Standard math types return alignment only for their inner components, so:
  163 -- > alignment (undefined :: Float3) = alignment (undefined :: Float)
  164 -- This function calculates more realistic alignment for shader programs.
  165 shaderAlignment :: Storable a => a -> Int
  166 shaderAlignment u = f (sizeOf u) (sizeOf (undefined :: Float)) where
  167   f n s
  168     | n <= s || s >= (sizeOf (undefined :: Float4)) = s
  169     | otherwise = f n (s * 2)
  170 
  171 uniform :: (OfValueType a, Storable a) => UniformBufferSlot -> IO (Node a)
  172 uniform UniformBufferSlot
  173   { uniformBufferSlotIndex = slot
  174   , uniformBufferSlotSizeRef = sizeRef
  175   } = withUndefinedM $ \u -> do
  176   bufferSize <- readIORef sizeRef
  177   let align = shaderAlignment u
  178   let alignedBufferSize = ((bufferSize + align - 1) `quot` align) * align
  179   writeIORef sizeRef $ alignedBufferSize + sizeOf u
  180   return $ UniformNode Uniform
  181     { uniformSlot = slot
  182     , uniformOffset = alignedBufferSize
  183     , uniformSize = 0
  184     , uniformType = valueType u
  185     }
  186 
  187 uniformArray
  188   :: (OfValueType a, Storable a)
  189   => Int -- ^ size
  190   -> UniformBufferSlot -- ^ slot
  191   -> IO (Node [a])
  192 uniformArray size UniformBufferSlot
  193   { uniformBufferSlotIndex = slot
  194   , uniformBufferSlotSizeRef = sizeRef
  195   } = wu func where
  196   wu :: (a -> IO (Node [a])) -> IO (Node [a])
  197   wu f = f undefined
  198   func u = do
  199     bufferSize <- readIORef sizeRef
  200     let align = shaderAlignment u
  201     let alignedBufferSize = ((bufferSize + align - 1) `quot` align) * align
  202     writeIORef sizeRef $ alignedBufferSize + (sizeOf u) * size
  203     return $ UniformNode Uniform
  204       { uniformSlot = slot
  205       , uniformOffset = alignedBufferSize
  206       , uniformSize = size
  207       , uniformType = valueType u
  208       }
  209 
  210 createUniformStorage :: Device d => d -> UniformBufferSlot -> IO (UniformStorage d, IO ())
  211 createUniformStorage device UniformBufferSlot
  212   { uniformBufferSlotIndex = slot
  213   , uniformBufferSlotSizeRef = sizeRef
  214   } = do
  215   size <- readIORef sizeRef
  216   -- align just in case
  217   let alignedSize = ((size + 15) `quot` 16) * 16
  218   (uniformBuffer, release) <- createUniformBuffer device alignedSize
  219   bytes <- mallocForeignPtrBytes alignedSize
  220   return (UniformStorage
  221     { uniformStorageSlot = slot
  222     , uniformStorageBufferId = uniformBuffer
  223     , uniformStorageBytes = bytes
  224     , uniformStorageSize = alignedSize
  225     }, release)
  226 
  227 setUniform :: (OfValueType a, Storable a) => UniformStorage d -> Node a -> a -> IO ()
  228 setUniform UniformStorage
  229   { uniformStorageBytes = bytes
  230   } (UniformNode Uniform
  231   { uniformOffset = offset
  232   }) value = do
  233   withForeignPtr bytes $ \ptr -> do
  234     pokeByteOff ptr offset value
  235 setUniform _ _ _ = undefined
  236 
  237 renderUniform :: (OfValueType a, Storable a) => UniformStorage d -> Node a -> a -> Render c ()
  238 renderUniform uniformStorage node value = liftIO $ setUniform uniformStorage node value
  239 
  240 setIndexedUniform :: (OfValueType a, Storable a) => UniformStorage d -> Node [a] -> Int -> a -> IO ()
  241 setIndexedUniform UniformStorage
  242   { uniformStorageBytes = bytes
  243   } (UniformNode Uniform
  244   { uniformOffset = offset
  245   }) i value = do
  246   withForeignPtr bytes $ \ptr -> do
  247     pokeElemOff (ptr `plusPtr` offset) i value
  248 setIndexedUniform _ _ _ _ = undefined
  249 
  250 renderIndexedUniform :: (OfValueType a, Storable a) => UniformStorage d -> Node [a] -> Int -> a -> Render c ()
  251 renderIndexedUniform uniformStorage node i value = liftIO $ setIndexedUniform uniformStorage node i value
  252 
  253 renderUploadUniformStorage :: Context c d => UniformStorage d -> Render c ()
  254 renderUploadUniformStorage UniformStorage
  255   { uniformStorageBufferId = uniformBuffer
  256   , uniformStorageBytes = bytes
  257   , uniformStorageSize = size
  258   } = do
  259   bs <- liftIO $ B.unsafePackCStringLen (castPtr $ unsafeForeignPtrToPtr bytes, size)
  260   renderUploadUniformBuffer uniformBuffer bs
  261   liftIO $ touchForeignPtr bytes
  262 
  263 renderUniformStorage :: Context c d => UniformStorage d -> Render c ()
  264 renderUniformStorage UniformStorage
  265   { uniformStorageSlot = slot
  266   , uniformStorageBufferId = uniformBuffer
  267   } = renderUniformBuffer slot uniformBuffer
  268 
  269 sampler :: (OfValueType s, OfValueType c) => Int -> SamplerDimension -> SamplerNode s c
  270 sampler slot dimension = withUndefined2 f where
  271   f s c = SamplerNode Sampler
  272     { samplerSlot = slot
  273     , samplerDimension = dimension
  274     , samplerSampleType = valueType s
  275     , samplerCoordsType = valueType c
  276     }
  277   withUndefined2 :: (s -> c -> SamplerNode s c) -> SamplerNode s c
  278   withUndefined2 q = q undefined undefined
  279 
  280 sample :: (OfVectorType (v c), OfVectorType (v Int32)) => SamplerNode s (v c) -> Node (v c) -> Node s
  281 sample s c = SampleNode
  282   { sampleNodeSamplerNode = s
  283   , sampleNodeCoordsNode = c
  284   , sampleNodeOffsetNode = Nothing
  285   , sampleNodeLod = SampleNodeAutoLod
  286   }
  287 
  288 sampleOffset :: (OfVectorType (v c), OfVectorType (v Int32)) => SamplerNode s (v c) -> Node (v c) -> Node (v Int32) -> Node s
  289 sampleOffset s c o = SampleNode
  290   { sampleNodeSamplerNode = s
  291   , sampleNodeCoordsNode = c
  292   , sampleNodeOffsetNode = Just o
  293   , sampleNodeLod = SampleNodeAutoLod
  294   }
  295 
  296 sampleLod :: (OfVectorType (v c), OfVectorType (v Int32)) => SamplerNode s (v c) -> Node (v c) -> Node c -> Node s
  297 sampleLod s c l = SampleNode
  298   { sampleNodeSamplerNode = s
  299   , sampleNodeCoordsNode = c
  300   , sampleNodeOffsetNode = Nothing
  301   , sampleNodeLod = SampleNodeLod l
  302   }
  303 
  304 sampleLodOffset :: (OfVectorType (v c), OfVectorType (v Int32)) => SamplerNode s (v c) -> Node (v c) -> Node c -> Node (v Int32) -> Node s
  305 sampleLodOffset s c l o = SampleNode
  306   { sampleNodeSamplerNode = s
  307   , sampleNodeCoordsNode = c
  308   , sampleNodeOffsetNode = Just o
  309   , sampleNodeLod = SampleNodeLod l
  310   }
  311 
  312 sampleBias :: (OfVectorType (v c), OfVectorType (v Int32)) => SamplerNode s (v c) -> Node (v c) -> Node c -> Node s
  313 sampleBias s c b = SampleNode
  314   { sampleNodeSamplerNode = s
  315   , sampleNodeCoordsNode = c
  316   , sampleNodeOffsetNode = Nothing
  317   , sampleNodeLod = SampleNodeBiasLod b
  318   }
  319 
  320 sampleBiasOffset :: (OfVectorType (v c), OfVectorType (v Int32)) => SamplerNode s (v c) -> Node (v c) -> Node c -> Node (v Int32) -> Node s
  321 sampleBiasOffset s c b o = SampleNode
  322   { sampleNodeSamplerNode = s
  323   , sampleNodeCoordsNode = c
  324   , sampleNodeOffsetNode = Just o
  325   , sampleNodeLod = SampleNodeBiasLod b
  326   }
  327 
  328 sampleGrad :: (OfVectorType (v c), OfVectorType (v Int32)) => SamplerNode s (v c) -> Node (v c) -> Node (v c) -> Node (v c) -> Node s
  329 sampleGrad s c gx gy = SampleNode
  330   { sampleNodeSamplerNode = s
  331   , sampleNodeCoordsNode = c
  332   , sampleNodeOffsetNode = Nothing
  333   , sampleNodeLod = SampleNodeGradLod gx gy
  334   }
  335 
  336 sampleGradOffset :: (OfVectorType (v c), OfVectorType (v Int32)) => SamplerNode s (v c) -> Node (v c) -> Node (v c) -> Node (v c) -> Node (v Int32) -> Node s
  337 sampleGradOffset s c gx gy o = SampleNode
  338   { sampleNodeSamplerNode = s
  339   , sampleNodeCoordsNode = c
  340   , sampleNodeOffsetNode = Just o
  341   , sampleNodeLod = SampleNodeGradLod gx gy
  342   }
  343 
  344 withState :: (State -> IO (State, a)) -> Program a
  345 withState f = do
  346   stateVar <- ask
  347   liftIO $ do
  348     state <- readIORef stateVar
  349     (newState, result) <- f state
  350     writeIORef stateVar newState
  351     return result
  352 
  353 temp :: OfValueType a => Node a -> Program (Node a)
  354 temp = withState . tempInternal
  355 
  356 tempInternal :: OfValueType a => Node a -> State -> IO (State, Node a)
  357 tempInternal node state@State
  358   { stateStage = stage
  359   , stateTemps = temps
  360   , stateTempsCount = tempsCount
  361   } = do
  362   if stage == EndStage then fail "failed to add temp after end of the program"
  363   else return ()
  364   return (state
  365     { stateTemps = (Temp
  366       { tempIndex = tempsCount
  367       , tempNode = node
  368       , tempStage = stage
  369       , tempType = nodeValueType node
  370       }) : temps
  371     , stateTempsCount = tempsCount + 1
  372     }, TempNode tempsCount)
  373 
  374 rasterize :: Node Float4 -> Program () -> Program ()
  375 rasterize positionNode pixelProgram = withState $ \state@State
  376   { stateStage = stage
  377   , stateTargets = targets
  378   } -> do
  379   if stage /= VertexStage then fail $ show ("wrong stage to add pixel program", stage)
  380   else return ()
  381   let positionTarget = PositionTarget positionNode
  382   pixelStateVar <- newIORef state
  383     { stateStage = PixelStage
  384     , stateTargets = positionTarget : targets
  385     }
  386   runReaderT pixelProgram pixelStateVar
  387   pixelState <- readIORef pixelStateVar
  388   return (pixelState
  389     { stateStage = EndStage
  390     }, ())
  391 
  392 colorTarget :: Int -> Node Float4 -> Program ()
  393 colorTarget i colorNode = withState $ \state@State
  394   { stateStage = stage
  395   , stateTargets = targets
  396   } -> do
  397   if stage /= PixelStage then fail "colorTarget can be used only in pixel program"
  398   else return ()
  399   let target = ColorTarget i colorNode
  400   return (state
  401     { stateTargets = targets ++ [target]
  402     }, ())
  403 
  404 dualColorTarget :: Node Float4 -> Node Float4 -> Program ()
  405 dualColorTarget colorNode1 colorNode2 = withState $ \state@State
  406   { stateStage = stage
  407   , stateTargets = targets
  408   } -> do
  409   if stage /= PixelStage then fail "dualColorTarget can be used only in pixel program"
  410   else return ()
  411   let target = DualColorTarget colorNode1 colorNode2
  412   return (state
  413     { stateTargets = target : targets
  414     }, ())
  415 
  416 depthTarget :: Node Float -> Program ()
  417 depthTarget depthNode = withState $ \state@State
  418   { stateStage = stage
  419   , stateTargets = targets
  420   } -> do
  421   if stage /= PixelStage then fail "depthTarget can be used only in pixel program"
  422   else return ()
  423   let target = DepthTarget depthNode
  424   return (state
  425     { stateTargets = target : targets
  426     }, ())
  427 
  428 constf :: Float -> Node Float
  429 constf = cnst
  430 const2f :: Float2 -> Node Float2
  431 const2f = cnst
  432 const3f :: Float3 -> Node Float3
  433 const3f = cnst
  434 const4f :: Float4 -> Node Float4
  435 const4f = cnst
  436 
  437 sampler1Df :: Int -> SamplerNode Float Float
  438 sampler1Df slot = sampler slot Sampler1D
  439 sampler1D2f :: Int -> SamplerNode Float2 Float
  440 sampler1D2f slot = sampler slot Sampler1D
  441 sampler1D3f :: Int -> SamplerNode Float3 Float
  442 sampler1D3f slot = sampler slot Sampler1D
  443 sampler1D4f :: Int -> SamplerNode Float4 Float
  444 sampler1D4f slot = sampler slot Sampler1D
  445 
  446 sampler2Df :: Int -> SamplerNode Float Float2
  447 sampler2Df slot = sampler slot Sampler2D
  448 sampler2D2f :: Int -> SamplerNode Float2 Float2
  449 sampler2D2f slot = sampler slot Sampler2D
  450 sampler2D3f :: Int -> SamplerNode Float3 Float2
  451 sampler2D3f slot = sampler slot Sampler2D
  452 sampler2D4f :: Int -> SamplerNode Float4 Float2
  453 sampler2D4f slot = sampler slot Sampler2D
  454 
  455 sampler3Df :: Int -> SamplerNode Float Float3
  456 sampler3Df slot = sampler slot Sampler3D
  457 sampler3D2f :: Int -> SamplerNode Float2 Float3
  458 sampler3D2f slot = sampler slot Sampler3D
  459 sampler3D3f :: Int -> SamplerNode Float3 Float3
  460 sampler3D3f slot = sampler slot Sampler3D
  461 sampler3D4f :: Int -> SamplerNode Float4 Float3
  462 sampler3D4f slot = sampler slot Sampler3D
  463 
  464 samplerCubef :: Int -> SamplerNode Float Float3
  465 samplerCubef slot = sampler slot SamplerCube
  466 samplerCube2f :: Int -> SamplerNode Float2 Float3
  467 samplerCube2f slot = sampler slot SamplerCube
  468 samplerCube3f :: Int -> SamplerNode Float3 Float3
  469 samplerCube3f slot = sampler slot SamplerCube
  470 samplerCube4f :: Int -> SamplerNode Float4 Float3
  471 samplerCube4f slot = sampler slot SamplerCube
  472 
  473 (!) :: (OfValueType a, OfValueType b, Integral b) => Node [a] -> Node b -> Node a
  474 a ! b = IndexNode (nodeArrayValueType a) (nodeValueType b) a b
  475 
  476 min_ :: OfValueType a => Node a -> Node a -> Node a
  477 min_ a b = MinNode (nodeValueType a) a b
  478 
  479 max_ :: OfValueType a => Node a -> Node a -> Node a
  480 max_ a b = MaxNode (nodeValueType a) a b
  481 
  482 clamp :: OfValueType a => Node a -> Node a -> Node a -> Node a
  483 clamp a b c = ClampNode (nodeValueType a) a b c
  484 
  485 lerp :: OfValueType a => Node a -> Node a -> Node a -> Node a
  486 lerp a b c = LerpNode (nodeValueType a) a b c
  487 
  488 equal_ :: OfValueType a => Node a -> Node a -> Node Bool
  489 equal_ a b = EqualNode (nodeValueType a) a b
  490 
  491 less_ :: OfValueType a => Node a -> Node a -> Node Bool
  492 less_ a b = LessNode (nodeValueType a) a b
  493 
  494 lessEqual_ :: OfValueType a => Node a -> Node a -> Node Bool
  495 lessEqual_ a b = LessEqualNode (nodeValueType a) a b
  496 
  497 greater_ :: OfValueType a => Node a -> Node a -> Node Bool
  498 greater_ = flip less_
  499 
  500 greaterEqual_ :: OfValueType a => Node a -> Node a -> Node Bool
  501 greaterEqual_ = flip lessEqual_
  502 
  503 if_ :: OfValueType a => Node Bool -> Node a -> Node a -> Node a
  504 if_ c a b = IfNode (nodeValueType a) c a b
  505 
  506 ddx :: OfValueType a => Node a -> Node a
  507 ddx a = DdxNode (nodeValueType a) a
  508 
  509 ddy :: OfValueType a => Node a -> Node a
  510 ddy a = DdyNode (nodeValueType a) a
  511 
  512 floor_ :: OfValueType a => Node a -> Node a
  513 floor_ a = FloorNode (nodeValueType a) a
  514 
  515 instanceId :: Node Word32
  516 instanceId = InstanceIdNode
  517 
  518 invSqrt :: (OfValueType a, Floating a) => Node a -> Node a
  519 invSqrt a = InvSqrtNode (nodeValueType a) a
  520 
  521 screenToTexture :: OfValueType a => Node a -> Node a
  522 screenToTexture a = ScreenToTextureNode (nodeValueType a) a
  523 
  524 normalizeSampledDepth :: Node Float -> Node Float
  525 normalizeSampledDepth = NormalizeSampledDepthNode
  526 
  527 fragCoord :: Node Float4
  528 fragCoord = FragCoordNode