{-# OPTIONS_GHC -fplugin Foreign.Storable.Generic.Plugin #-}

{-# LANGUAGE OverloadedLists #-}

module Render.DescSets.Sun
  ( Sun(..)
  , createSet0Ds
  , set0

  , pattern MAX_VIEWS

  , Buffer

  , SunInput(..)
  , initialSunInput

  , Process
  , spawn1
  , mkSun

  , Observer
  , newObserver1
  , observe1
  ) where

import RIO

import Control.Monad.Trans.Resource (MonadResource, ResourceT)
import Data.Tagged (Tagged(..))
import Data.Vector qualified as Vector
import Data.Vector.Storable qualified as VectorS
import Foreign.Storable.Generic (GStorable)
import Geomancy (Vec3, Vec4, vec3, vec4)
import Geomancy.Transform (Transform)
import Geomancy.Transform qualified as Transform
import Geomancy.Vec4 qualified as Vec4
import Vulkan.Core10 qualified as Vk
import Vulkan.CStruct.Extends (SomeStruct(..))
import Vulkan.NamedType ((:::))
import Vulkan.Zero (Zero(..))

import Engine.Camera qualified as Camera
import Engine.Types (StageRIO)
import Engine.Vulkan.DescSets ()
import Engine.Vulkan.Types (DsLayoutBindings, HasVulkan(..))
import Engine.Worker qualified as Worker
import Resource.Buffer qualified as Buffer
import Resource.Region qualified as Region
import Resource.Vulkan.DescriptorPool qualified as DescriptorPool

-- * Set0 data for light projection

-- | Maximum "guaranteed" amount for multiview passes
pattern MAX_VIEWS :: Int
pattern $mMAX_VIEWS :: forall {r}. Int -> ((# #) -> r) -> ((# #) -> r) -> r
$bMAX_VIEWS :: Int
MAX_VIEWS = 6

data Sun = Sun
  { Sun -> Transform
sunViewProjection :: Transform
  , Sun -> Vec4
sunShadow         :: Vec4 -- offsetx, offsety, index, size -- XXX: only index is used
  , Sun -> Vec4
sunPosition       :: Vec4 -- XXX: alpha available for stuff
  , Sun -> Vec4
sunDirection      :: Vec4 -- XXX: alpha available for stuff
  , Sun -> Vec4
sunColor          :: Vec4 -- XXX: RGB premultiplied, alpha is available for stuff
  }
  deriving (Int -> Sun -> ShowS
[Sun] -> ShowS
Sun -> String
(Int -> Sun -> ShowS)
-> (Sun -> String) -> ([Sun] -> ShowS) -> Show Sun
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Sun -> ShowS
showsPrec :: Int -> Sun -> ShowS
$cshow :: Sun -> String
show :: Sun -> String
$cshowList :: [Sun] -> ShowS
showList :: [Sun] -> ShowS
Show, (forall x. Sun -> Rep Sun x)
-> (forall x. Rep Sun x -> Sun) -> Generic Sun
forall x. Rep Sun x -> Sun
forall x. Sun -> Rep Sun x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. Sun -> Rep Sun x
from :: forall x. Sun -> Rep Sun x
$cto :: forall x. Rep Sun x -> Sun
to :: forall x. Rep Sun x -> Sun
Generic)

instance GStorable Sun

instance Zero Sun where
  zero :: Sun
zero = Sun
    { $sel:sunViewProjection:Sun :: Transform
sunViewProjection = Transform
forall a. Monoid a => a
mempty
    , $sel:sunShadow:Sun :: Vec4
sunShadow         = Vec4
0
    , $sel:sunPosition:Sun :: Vec4
sunPosition       = Vec4
0
    , $sel:sunDirection:Sun :: Vec4
sunDirection      = Float -> Float -> Float -> Float -> Vec4
vec4 Float
0 Float
1 Float
0 Float
0
    , $sel:sunColor:Sun :: Vec4
sunColor          = Vec4
0
    }

-- * Shadow casting descriptor set

set0
  :: Tagged Sun DsLayoutBindings
set0 :: Tagged Sun DsLayoutBindings
set0 = DsLayoutBindings -> Tagged Sun DsLayoutBindings
forall {k} (s :: k) b. b -> Tagged s b
Tagged
  [ (DescriptorSetLayoutBinding
set0bind0, DescriptorBindingFlags
forall a. Zero a => a
zero)
  ]

set0bind0 :: Vk.DescriptorSetLayoutBinding
set0bind0 :: DescriptorSetLayoutBinding
set0bind0 = Vk.DescriptorSetLayoutBinding
  { $sel:binding:DescriptorSetLayoutBinding :: Word32
binding           = Word32
0
  , $sel:descriptorType:DescriptorSetLayoutBinding :: DescriptorType
descriptorType    = DescriptorType
Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
  , $sel:descriptorCount:DescriptorSetLayoutBinding :: Word32
descriptorCount   = Word32
1
  , $sel:stageFlags:DescriptorSetLayoutBinding :: ShaderStageFlags
stageFlags        = ShaderStageFlags
Vk.SHADER_STAGE_VERTEX_BIT
  , $sel:immutableSamplers:DescriptorSetLayoutBinding :: Vector Sampler
immutableSamplers = Vector Sampler
forall a. Monoid a => a
mempty
  }

-- * Setup

type Buffer = Buffer.Allocated 'Buffer.Coherent Sun

createSet0Ds
  :: Tagged '[Sun] Vk.DescriptorSetLayout
  -> ResourceT (StageRIO st)
      ( Tagged '[Sun] (Vector Vk.DescriptorSet)
      , Buffer
      )
createSet0Ds :: forall st.
Tagged '[Sun] DescriptorSetLayout
-> ResourceT
     (StageRIO st) (Tagged '[Sun] (Vector DescriptorSet), Buffer)
createSet0Ds (Tagged DescriptorSetLayout
set0layout) = do
  DescriptorPool
descPool <- StageRIO st (ReleaseKey, DescriptorPool)
-> ResourceT (StageRIO st) DescriptorPool
forall (m :: * -> *) a.
MonadResource m =>
m (ReleaseKey, a) -> ResourceT m a
Region.local (StageRIO st (ReleaseKey, DescriptorPool)
 -> ResourceT (StageRIO st) DescriptorPool)
-> StageRIO st (ReleaseKey, DescriptorPool)
-> ResourceT (StageRIO st) DescriptorPool
forall a b. (a -> b) -> a -> b
$
    Maybe Text
-> Word32
-> [(DescriptorType, Word32)]
-> StageRIO st (ReleaseKey, DescriptorPool)
forall env (m :: * -> *).
(MonadVulkan env m, MonadResource m) =>
Maybe Text
-> Word32
-> [(DescriptorType, Word32)]
-> m (ReleaseKey, DescriptorPool)
DescriptorPool.allocate (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"Basic.Sun") Word32
1
      [ ( DescriptorType
Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
        , Word32
1 Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
1
        )
      ]

  Vector DescriptorSet
descSets <- DescriptorPool
-> Maybe Text
-> Vector DescriptorSetLayout
-> ResourceT (StageRIO st) (Vector DescriptorSet)
forall env (m :: * -> *).
MonadVulkan env m =>
DescriptorPool
-> Maybe Text
-> Vector DescriptorSetLayout
-> m (Vector DescriptorSet)
DescriptorPool.allocateSetsFrom DescriptorPool
descPool (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"Basic.Sun") [Item (Vector DescriptorSetLayout)
DescriptorSetLayout
set0layout]

  Buffer
sunData <- StageRIO st (ReleaseKey, Buffer) -> ResourceT (StageRIO st) Buffer
forall (m :: * -> *) a.
MonadResource m =>
m (ReleaseKey, a) -> ResourceT m a
Region.local (StageRIO st (ReleaseKey, Buffer)
 -> ResourceT (StageRIO st) Buffer)
-> StageRIO st (ReleaseKey, Buffer)
-> ResourceT (StageRIO st) Buffer
forall a b. (a -> b) -> a -> b
$
    Maybe Text
-> BufferUsageFlagBits
-> Int
-> Vector Sun
-> StageRIO st (ReleaseKey, Buffer)
forall env (m :: * -> *) a.
(MonadVulkan env m, MonadResource m, Storable a) =>
Maybe Text
-> BufferUsageFlagBits
-> Int
-> Vector a
-> m (ReleaseKey, Allocated 'Coherent a)
Buffer.allocateCoherent
      (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"Basic.Sun.Data")
      BufferUsageFlagBits
Vk.BUFFER_USAGE_UNIFORM_BUFFER_BIT Int
MAX_VIEWS
      (Int -> Sun -> Vector Sun
forall a. Storable a => Int -> a -> Vector a
VectorS.replicate Int
MAX_VIEWS Sun
forall a. Zero a => a
zero)

  Tagged '[Sun] (Vector DescriptorSet)
-> Buffer -> ResourceT (StageRIO st) ()
forall st.
Tagged '[Sun] (Vector DescriptorSet)
-> Buffer -> ResourceT (StageRIO st) ()
updateSet0Ds (Vector DescriptorSet -> Tagged '[Sun] (Vector DescriptorSet)
forall {k} (s :: k) b. b -> Tagged s b
Tagged Vector DescriptorSet
descSets) Buffer
sunData

  pure (Vector DescriptorSet -> Tagged '[Sun] (Vector DescriptorSet)
forall {k} (s :: k) b. b -> Tagged s b
Tagged Vector DescriptorSet
descSets, Buffer
sunData)

updateSet0Ds
  :: Tagged '[Sun] (Vector Vk.DescriptorSet)
  -> Buffer.Allocated 'Buffer.Coherent Sun
  -> ResourceT (StageRIO st) ()
updateSet0Ds :: forall st.
Tagged '[Sun] (Vector DescriptorSet)
-> Buffer -> ResourceT (StageRIO st) ()
updateSet0Ds (Tagged Vector DescriptorSet
ds) Buffer.Allocated{Buffer
aBuffer :: Buffer
$sel:aBuffer:Allocated :: forall {k} (s :: Store) (a :: k). Allocated s a -> Buffer
aBuffer} = do
  App GlobalHandles st
context <- (App GlobalHandles st -> App GlobalHandles st)
-> ResourceT (StageRIO st) (App GlobalHandles st)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks App GlobalHandles st -> App GlobalHandles st
forall a. a -> a
id
  Device
-> Vector (SomeStruct WriteDescriptorSet)
-> Vector CopyDescriptorSet
-> ResourceT (StageRIO st) ()
forall (io :: * -> *).
MonadIO io =>
Device
-> Vector (SomeStruct WriteDescriptorSet)
-> Vector CopyDescriptorSet
-> io ()
Vk.updateDescriptorSets (App GlobalHandles st -> Device
forall a. HasVulkan a => a -> Device
getDevice App GlobalHandles st
context) [Item (Vector (SomeStruct WriteDescriptorSet))
SomeStruct WriteDescriptorSet
writeSet0b0] Vector CopyDescriptorSet
forall a. Monoid a => a
mempty

  where
    destSet0 :: DescriptorSet
destSet0 = case Vector DescriptorSet -> Maybe DescriptorSet
forall (m :: * -> *) a. Monad m => Vector a -> m a
Vector.headM Vector DescriptorSet
ds of
      Maybe DescriptorSet
Nothing ->
        String -> DescriptorSet
forall a. HasCallStack => String -> a
error String
"assert: descriptor sets promised to contain [Sun]"
      Just DescriptorSet
one ->
        DescriptorSet
one

    writeSet0b0 :: SomeStruct WriteDescriptorSet
writeSet0b0 = WriteDescriptorSet '[] -> SomeStruct WriteDescriptorSet
forall (a :: [*] -> *) (es :: [*]).
(Extendss a es, PokeChain es, Show (Chain es)) =>
a es -> SomeStruct a
SomeStruct WriteDescriptorSet '[]
forall a. Zero a => a
zero
      { $sel:dstSet:WriteDescriptorSet :: DescriptorSet
Vk.dstSet          = DescriptorSet
destSet0
      , $sel:dstBinding:WriteDescriptorSet :: Word32
Vk.dstBinding      = Word32
0
      , $sel:dstArrayElement:WriteDescriptorSet :: Word32
Vk.dstArrayElement = Word32
0
      , $sel:descriptorCount:WriteDescriptorSet :: Word32
Vk.descriptorCount = Word32
1
      , $sel:descriptorType:WriteDescriptorSet :: DescriptorType
Vk.descriptorType  = DescriptorType
Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
      , $sel:bufferInfo:WriteDescriptorSet :: Vector DescriptorBufferInfo
Vk.bufferInfo      = [Item (Vector DescriptorBufferInfo)
DescriptorBufferInfo
set0bind0I]
      }
      where
        set0bind0I :: DescriptorBufferInfo
set0bind0I = Vk.DescriptorBufferInfo
          { $sel:buffer:DescriptorBufferInfo :: Buffer
Vk.buffer = Buffer
aBuffer
          , $sel:offset:DescriptorBufferInfo :: DeviceSize
Vk.offset = DeviceSize
0
          , $sel:range:DescriptorBufferInfo :: DeviceSize
Vk.range  = DeviceSize
Vk.WHOLE_SIZE
          }

data SunInput = SunInput
  { SunInput -> Vec4
siColor :: Vec4

  , SunInput -> Float
siInclination :: Float
  , SunInput -> Float
siAzimuth     :: Float
  , SunInput -> Float
siRadius      :: Float
  , SunInput -> Vec3
siTarget      :: Vec3

  , SunInput -> Float
siDepthRange :: Float
  , SunInput -> Float
siSize       :: Float
  , SunInput -> Float
siShadowIx   :: Float
  }

initialSunInput :: SunInput
initialSunInput :: SunInput
initialSunInput = SunInput
  { $sel:siColor:SunInput :: Vec4
siColor = Float -> Float -> Float -> Float -> Vec4
vec4 Float
1 Float
1 Float
1 Float
1

  , $sel:siInclination:SunInput :: Float
siInclination = Float
τFloat -> Float -> Float
forall a. Fractional a => a -> a -> a
/Float
8
  , $sel:siAzimuth:SunInput :: Float
siAzimuth     = -Float
τFloat -> Float -> Float
forall a. Fractional a => a -> a -> a
/Float
8
  , $sel:siRadius:SunInput :: Float
siRadius      = Float
forall a. (Eq a, Num a) => a
Camera.PROJECTION_FAR Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
2
  , $sel:siTarget:SunInput :: Vec3
siTarget      = Vec3
0

  , $sel:siDepthRange:SunInput :: Float
siDepthRange = Float
forall a. (Eq a, Num a) => a
Camera.PROJECTION_FAR
  , $sel:siSize:SunInput :: Float
siSize       = Float
512
  , $sel:siShadowIx:SunInput :: Float
siShadowIx   = -Float
1
  }

type Process = Worker.Cell SunInput ("bounding box" ::: Transform, Sun)

spawn1
  :: ( MonadResource m
     , MonadUnliftIO m
     )
  => SunInput
  -> m Process
spawn1 :: forall (m :: * -> *).
(MonadResource m, MonadUnliftIO m) =>
SunInput -> m Process
spawn1 = (SunInput -> (Transform, Sun)) -> SunInput -> m Process
forall (m :: * -> *) input output.
(MonadUnliftIO m, MonadResource m) =>
(input -> output) -> input -> m (Cell input output)
Worker.spawnCell SunInput -> (Transform, Sun)
mkSun

mkSun :: SunInput -> ("bounding box" ::: Transform, Sun)
mkSun :: SunInput -> (Transform, Sun)
mkSun SunInput{Float
Vec3
Vec4
$sel:siColor:SunInput :: SunInput -> Vec4
$sel:siInclination:SunInput :: SunInput -> Float
$sel:siAzimuth:SunInput :: SunInput -> Float
$sel:siRadius:SunInput :: SunInput -> Float
$sel:siTarget:SunInput :: SunInput -> Vec3
$sel:siDepthRange:SunInput :: SunInput -> Float
$sel:siSize:SunInput :: SunInput -> Float
$sel:siShadowIx:SunInput :: SunInput -> Float
siColor :: Vec4
siInclination :: Float
siAzimuth :: Float
siRadius :: Float
siTarget :: Vec3
siDepthRange :: Float
siSize :: Float
siShadowIx :: Float
..} =
  ( Transform
bbTransform
  , Sun
      { $sel:sunViewProjection:Sun :: Transform
sunViewProjection = [Transform] -> Transform
forall a. Monoid a => [a] -> a
mconcat [Transform]
vp
      , $sel:sunShadow:Sun :: Vec4
sunShadow         = Float -> Float -> Float -> Float -> Vec4
vec4 Float
0 Float
0 Float
siShadowIx Float
siSize
      , $sel:sunPosition:Sun :: Vec4
sunPosition       = Vec3 -> Float -> Vec4
forall a. Coercible a Vec3 => a -> Float -> Vec4
Vec4.fromVec3 Vec3
position Float
0
      , $sel:sunDirection:Sun :: Vec4
sunDirection      = Vec3 -> Float -> Vec4
forall a. Coercible a Vec3 => a -> Float -> Vec4
Vec4.fromVec3 Vec3
direction Float
0
      , $sel:sunColor:Sun :: Vec4
sunColor          = Vec4
siColor
      }
  )
  where
    vp :: [Transform]
vp =
      [ Float -> Transform
Transform.rotateY (-Float
siAzimuth)
      , Float -> Transform
Transform.rotateX (-Float
siInclination)

      , Float -> Float -> Float -> Transform
Transform.translate Float
0 Float
0 Float
siRadius

      -- XXX: some area beyond the near plane receives light, but not shadows
      , Float -> Float -> Float -> Transform
Transform.scale3
          (Float
1 Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
siSize)
          (Float
1 Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
siSize)
          (Float
1 Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
siDepthRange)
      ]

    position :: Vec3
position = Vec3 -> Transform -> Vec3
Transform.apply (Float -> Float -> Float -> Vec3
vec3 Float
0 Float
0 Float
siRadius) Transform
rotation

    direction :: Vec3
direction = Vec3 -> Transform -> Vec3
Transform.apply (Float -> Float -> Float -> Vec3
vec3 Float
0 Float
0 (Float -> Vec3) -> Float -> Vec3
forall a b. (a -> b) -> a -> b
$ -Float
1) Transform
rotation

    bbTransform :: Transform
bbTransform = [Transform] -> Transform
forall a. Monoid a => [a] -> a
mconcat
      [ -- XXX: orient wire box "green/near -> far/red"
        Float -> Transform
Transform.rotateX (Float
τFloat -> Float -> Float
forall a. Fractional a => a -> a -> a
/Float
4)
        -- XXX: the rest must be matched with VP flipped
      , Float -> Float -> Float -> Transform
Transform.translate Float
0 Float
0 Float
0.5                 -- XXX: shift origin to the near face

        -- XXX: reverse light transform
      , Float -> Float -> Float -> Transform
Transform.scale3 Float
siSize Float
siSize Float
siDepthRange -- XXX: size to projection volume
      , Float -> Float -> Float -> Transform
Transform.translate Float
0 Float
0 (-Float
siRadius)         -- XXX: translate near face to radius
      , Item [Transform]
Transform
rotation                                    -- XXX: apply sphere coords
      ]

    rotation :: Transform
rotation = [Transform] -> Transform
forall a. Monoid a => [a] -> a
mconcat
      [ Float -> Transform
Transform.rotateX Float
siInclination
      , Float -> Transform
Transform.rotateY Float
siAzimuth
      ]

type Observer = Worker.ObserverIO (VectorS.Vector ("bounding box" ::: Transform))

newObserver1 :: MonadIO m => m Observer
newObserver1 :: forall (m :: * -> *). MonadIO m => m Observer
newObserver1 = Vector Transform -> m Observer
forall (m :: * -> *) a. MonadIO m => a -> m (ObserverIO a)
Worker.newObserverIO Vector Transform
forall a. Monoid a => a
mempty

observe1 :: MonadUnliftIO m => Process -> Observer -> Buffer -> m ()
observe1 :: forall (m :: * -> *).
MonadUnliftIO m =>
Process -> Observer -> Buffer -> m ()
observe1 Process
sunP Observer
sunOut Buffer
sunData =
  Process
-> Observer
-> (Vector Transform -> GetOutput Process -> m (Vector Transform))
-> m ()
forall (m :: * -> *) output a.
(MonadUnliftIO m, HasOutput output) =>
output -> ObserverIO a -> (a -> GetOutput output -> m a) -> m ()
Worker.observeIO_ Process
sunP Observer
sunOut \Vector Transform
_oldBB (Transform
bb, Sun
sun) -> do
    -- XXX: must stay the same or descsets must be updated with a new buffer
    Buffer
_same <- Vector Sun -> Buffer -> m Buffer
forall (io :: * -> *) a.
(MonadUnliftIO io, Storable a) =>
Vector a -> Allocated 'Coherent a -> io (Allocated 'Coherent a)
Buffer.updateCoherent (Sun -> Vector Sun
forall a. Storable a => a -> Vector a
VectorS.singleton Sun
sun) Buffer
sunData
    pure $ Transform -> Vector Transform
forall a. Storable a => a -> Vector a
VectorS.singleton Transform
bb

τ :: Float
τ :: Float
τ = Float
2 Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
forall a. Floating a => a
pi