{-|
Specialization constants, normalized to stacks of 32-bit units.


@
data MySpec = MySpec { width :: Word32, height :: Word32, scale :: Float }
instance Specialization MySpec where -- ... pack the fields

withSpecialization sp \\mSpec ->
  -- build a 'Vulkan.Core10.PipelineShaderStageCreateInfo' with
  --   specializationInfo = mSpec
  ...
@
-}
module Vulkan.Utils.Pipeline.Specialization
  ( withSpecialization
  , allocateSpecialization
  , Specialization (..)
  , SpecializationConst (..)
  ) where

import Control.Monad.IO.Class (liftIO)
import Control.Monad.IO.Unlift (MonadUnliftIO, withRunInIO)
import Control.Monad.Trans.Resource (MonadResource, allocate)
import Data.Bool (bool)
import Data.Int (Int32)
import Data.Vector (Vector)
import qualified Data.Vector as Vector
import qualified Data.Vector.Storable as Storable
import Data.Word (Word32)
import Foreign.Marshal.Alloc (free)
import Foreign.Marshal.Array (mallocArray, pokeArray)
import Foreign.Ptr (castPtr)
import GHC.Float (castFloatToWord32)
import qualified Vulkan.Core10 as Vk

{- | Provide a 'Vk.SpecializationInfo' describing @spec@ to the callback.

The info (and the buffer its @data'@ pointer references) is valid only for the
duration of the callback, which is exactly the window in which it needs to live:
pipeline creation copies the constant values out. Build and create the pipeline
inside the continuation.

An empty specialization (e.g. @()@ or an empty list) yields 'Nothing', so the
shader stage's @specializationInfo@ stays unset.
-}
withSpecialization
  :: (Specialization spec, MonadUnliftIO m)
  => spec
  -> (Maybe Vk.SpecializationInfo -> m a)
  -> m a
withSpecialization :: forall spec (m :: * -> *) a.
(Specialization spec, MonadUnliftIO m) =>
spec -> (Maybe SpecializationInfo -> m a) -> m a
withSpecialization spec
spec Maybe SpecializationInfo -> m a
action =
  if Vector Word32 -> Bool
forall a. Storable a => Vector a -> Bool
Storable.null Vector Word32
specData
    then
      Maybe SpecializationInfo -> m a
action Maybe SpecializationInfo
forall a. Maybe a
Nothing
    else ((forall a. m a -> IO a) -> IO a) -> m a
forall b. ((forall a. m a -> IO a) -> IO b) -> m b
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO (((forall a. m a -> IO a) -> IO a) -> m a)
-> ((forall a. m a -> IO a) -> IO a) -> m a
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run ->
      Vector Word32 -> (Ptr Word32 -> IO a) -> IO a
forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
Storable.unsafeWith Vector Word32
specData ((Ptr Word32 -> IO a) -> IO a) -> (Ptr Word32 -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Ptr Word32
specPtr ->
        m a -> IO a
forall a. m a -> IO a
run (m a -> IO a)
-> (Maybe SpecializationInfo -> m a)
-> Maybe SpecializationInfo
-> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe SpecializationInfo -> m a
action (Maybe SpecializationInfo -> IO a)
-> Maybe SpecializationInfo -> IO a
forall a b. (a -> b) -> a -> b
$
          SpecializationInfo -> Maybe SpecializationInfo
forall a. a -> Maybe a
Just
            Vk.SpecializationInfo
              { mapEntries :: Vector SpecializationMapEntry
Vk.mapEntries = Vector SpecializationMapEntry
mapEntries
              , dataSize :: Word64
Vk.dataSize = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word64) -> Int -> Word64
forall a b. (a -> b) -> a -> b
$ Vector Word32 -> Int
forall a. Storable a => Vector a -> Int
Storable.length Vector Word32
specData Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4
              , data' :: Ptr ()
Vk.data' = Ptr Word32 -> Ptr ()
forall a b. Ptr a -> Ptr b
castPtr Ptr Word32
specPtr
              }
  where
    specData :: Storable.Vector Word32
    specData :: Vector Word32
specData = [Word32] -> Vector Word32
forall a. Storable a => [a] -> Vector a
Storable.fromList (spec -> [Word32]
forall a. Specialization a => a -> [Word32]
specializationData spec
spec)

    mapEntries :: Vector Vk.SpecializationMapEntry
    mapEntries :: Vector SpecializationMapEntry
mapEntries = Int -> Vector SpecializationMapEntry
specializationMapEntries (Vector Word32 -> Int
forall a. Storable a => Vector a -> Int
Storable.length Vector Word32
specData)

{- | Pack a specialization into a buffer tied to the current resource scope,
yielding the 'Vk.SpecializationInfo' to embed in a shader stage's
@specializationInfo@.

Unlike 'withSpecialization' this is not continuation-scoped: the backing buffer
lives until the surrounding 'Control.Monad.Trans.Resource.ResourceT' block ends,
so it survives a later pipeline creation. An empty specialization (e.g. @()@)
yields 'Nothing'.
-}
allocateSpecialization
  :: (Specialization spec, MonadResource m)
  => spec
  -> m (Maybe Vk.SpecializationInfo)
allocateSpecialization :: forall spec (m :: * -> *).
(Specialization spec, MonadResource m) =>
spec -> m (Maybe SpecializationInfo)
allocateSpecialization spec
spec =
  case spec -> [Word32]
forall a. Specialization a => a -> [Word32]
specializationData spec
spec of
    [] ->
      Maybe SpecializationInfo -> m (Maybe SpecializationInfo)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe SpecializationInfo
forall a. Maybe a
Nothing
    [Word32]
ws -> do
      let n :: Int
n = [Word32] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Word32]
ws
      -- A C-malloc'd buffer is pointer-stable and freed at scope end; the
      -- pointer must stay valid until pipeline creation copies the values.
      (_key, ptr) <- IO (Ptr Word32)
-> (Ptr Word32 -> IO ()) -> m (ReleaseKey, Ptr Word32)
forall (m :: * -> *) a.
MonadResource m =>
IO a -> (a -> IO ()) -> m (ReleaseKey, a)
allocate (Int -> IO (Ptr Word32)
forall a. Storable a => Int -> IO (Ptr a)
mallocArray Int
n) Ptr Word32 -> IO ()
forall a. Ptr a -> IO ()
free
      liftIO $ pokeArray ptr ws
      pure $
        Just
          Vk.SpecializationInfo
            { Vk.mapEntries = specializationMapEntries n
            , Vk.dataSize = fromIntegral (n * 4)
            , Vk.data' = castPtr ptr
            }

-- | One 32-bit @constantID = offset/4@ entry per slot, counting from zero.
specializationMapEntries :: Int -> Vector Vk.SpecializationMapEntry
specializationMapEntries :: Int -> Vector SpecializationMapEntry
specializationMapEntries Int
n =
  Int
-> (Int -> SpecializationMapEntry) -> Vector SpecializationMapEntry
forall a. Int -> (Int -> a) -> Vector a
Vector.generate Int
n ((Int -> SpecializationMapEntry) -> Vector SpecializationMapEntry)
-> (Int -> SpecializationMapEntry) -> Vector SpecializationMapEntry
forall a b. (a -> b) -> a -> b
$ \Int
ix ->
    Vk.SpecializationMapEntry
      { constantID :: Word32
Vk.constantID = Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
ix
      , offset :: Word32
Vk.offset = Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4)
      , size :: Word64
Vk.size = Word64
4
      }

{- | A value that flattens to a stack of 32-bit specialization constants, in
@constant_id@ order starting from zero.
-}
class Specialization a where
  specializationData :: a -> [Word32]

instance Specialization () where
  specializationData :: () -> [Word32]
specializationData ()
_ = []

-- | Pre-packed constants, used as-is.
instance Specialization [Word32] where
  specializationData :: [Word32] -> [Word32]
specializationData = [Word32] -> [Word32]
forall a. a -> a
id

instance Specialization Word32 where
  specializationData :: Word32 -> [Word32]
specializationData Word32
x = [Word32 -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData Word32
x]

instance Specialization Int32 where
  specializationData :: Int32 -> [Word32]
specializationData Int32
x = [Int32 -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData Int32
x]

instance Specialization Float where
  specializationData :: Float -> [Word32]
specializationData Float
x = [Float -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData Float
x]

instance Specialization Bool where
  specializationData :: Bool -> [Word32]
specializationData Bool
x = [Bool -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData Bool
x]

{- | A single scalar specialization constant, reinterpreted into its 32-bit
representation.

Per the @GL_KHR_vulkan_glsl@ spec a @constant_id@ may only decorate a scalar
@int@, @float@ or @bool@; @uint@ works in practice too. All of these are 32 bits
wide.
-}
class SpecializationConst a where
  packConstData :: a -> Word32

instance SpecializationConst Word32 where
  packConstData :: Word32 -> Word32
packConstData = Word32 -> Word32
forall a. a -> a
id

-- | Two's-complement bit pattern, preserved by 'fromIntegral' at the same width.
instance SpecializationConst Int32 where
  packConstData :: Int32 -> Word32
packConstData = Int32 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral

instance SpecializationConst Float where
  packConstData :: Float -> Word32
packConstData = Float -> Word32
castFloatToWord32

instance SpecializationConst Bool where
  packConstData :: Bool -> Word32
packConstData = Word32 -> Word32 -> Bool -> Word32
forall a. a -> a -> Bool -> a
bool Word32
0 Word32
1

instance
  ( SpecializationConst a
  , SpecializationConst b
  )
  => Specialization (a, b)
  where
  specializationData :: (a, b) -> [Word32]
specializationData (a
a, b
b) =
    [ a -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData a
a
    , b -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData b
b
    ]

instance
  ( SpecializationConst a
  , SpecializationConst b
  , SpecializationConst c
  )
  => Specialization (a, b, c)
  where
  specializationData :: (a, b, c) -> [Word32]
specializationData (a
a, b
b, c
c) =
    [ a -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData a
a
    , b -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData b
b
    , c -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData c
c
    ]

instance
  ( SpecializationConst a
  , SpecializationConst b
  , SpecializationConst c
  , SpecializationConst d
  )
  => Specialization (a, b, c, d)
  where
  specializationData :: (a, b, c, d) -> [Word32]
specializationData (a
a, b
b, c
c, d
d) =
    [ a -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData a
a
    , b -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData b
b
    , c -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData c
c
    , d -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData d
d
    ]

instance
  ( SpecializationConst a
  , SpecializationConst b
  , SpecializationConst c
  , SpecializationConst d
  , SpecializationConst e
  )
  => Specialization (a, b, c, d, e)
  where
  specializationData :: (a, b, c, d, e) -> [Word32]
specializationData (a
a, b
b, c
c, d
d, e
e) =
    [ a -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData a
a
    , b -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData b
b
    , c -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData c
c
    , d -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData d
d
    , e -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData e
e
    ]

instance
  ( SpecializationConst a
  , SpecializationConst b
  , SpecializationConst c
  , SpecializationConst d
  , SpecializationConst e
  , SpecializationConst f
  )
  => Specialization (a, b, c, d, e, f)
  where
  specializationData :: (a, b, c, d, e, f) -> [Word32]
specializationData (a
a, b
b, c
c, d
d, e
e, f
f) =
    [ a -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData a
a
    , b -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData b
b
    , c -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData c
c
    , d -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData d
d
    , e -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData e
e
    , f -> Word32
forall a. SpecializationConst a => a -> Word32
packConstData f
f
    ]