Sun.hs

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

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 (ResourceT)
import Control.Monad.Trans.Resource qualified as 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.Utils.Debug qualified as Debug
import Vulkan.Zero (Zero(..))

import Engine.Camera qualified as Camera
import Engine.Types (StageRIO)
import Engine.Vulkan.DescSets ()
import Engine.Vulkan.Types (DsBindings, HasVulkan(..))
import Engine.Worker qualified as Worker
import Resource.Buffer qualified as Buffer
import Resource.DescriptorSet qualified as DescriptorSet

-- * Set0 data for light projection

-- | Maximum "guaranteed" amount for multiview passes
pattern MAX_VIEWS :: Int
pattern MAX_VIEWS = 6

data Sun = Sun
  { sunViewProjection :: Transform
  , sunShadow         :: Vec4 -- offsetx, offsety, index, size -- XXX: only index is used
  , sunPosition       :: Vec4 -- XXX: alpha available for stuff
  , sunDirection      :: Vec4 -- XXX: alpha available for stuff
  , sunColor          :: Vec4 -- XXX: RGB premultiplied, alpha is available for stuff
  }
  deriving (Show, Generic)

instance GStorable Sun

instance Zero Sun where
  zero = Sun
    { sunViewProjection = mempty
    , sunShadow         = 0
    , sunPosition       = 0
    , sunDirection      = vec4 0 1 0 0
    , sunColor          = 0
    }

-- * Shadow casting descriptor set

set0
  :: Tagged Sun DsBindings
set0 = Tagged
  [ (set0bind0, zero)
  ]

set0bind0 :: Vk.DescriptorSetLayoutBinding
set0bind0 = Vk.DescriptorSetLayoutBinding
  { binding           = 0
  , descriptorType    = Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
  , descriptorCount   = 1
  , stageFlags        = Vk.SHADER_STAGE_VERTEX_BIT
  , immutableSamplers = mempty
  }

-- * Setup

type Buffer = Buffer.Allocated 'Buffer.Coherent Sun

createSet0Ds
  :: Tagged '[Sun] Vk.DescriptorSetLayout
  -> ResourceT (StageRIO st)
      ( Tagged '[Sun] (Vector Vk.DescriptorSet)
      , Buffer
      )
createSet0Ds (Tagged set0layout) = do
  context <- asks id

  (_dpKey, descPool) <- DescriptorSet.allocatePool 1 dpSizes

  let
    set0dsCI = zero
      { Vk.descriptorPool = descPool
      , Vk.setLayouts     = Vector.singleton set0layout
      }
  descSets <- fmap (Tagged @'[Sun]) $
    Vk.allocateDescriptorSets (getDevice context) set0dsCI

  let
    initialSuns = VectorS.replicate MAX_VIEWS zero
  (_, sunData) <- ResourceT.allocate
    (Buffer.createCoherent context Vk.BUFFER_USAGE_UNIFORM_BUFFER_BIT MAX_VIEWS initialSuns)
    (Buffer.destroyAll context . Just)

  updateSet0Ds descSets sunData

  let device = getDevice context
  Debug.nameObject device descPool "Sun.Pool"
  for_ (unTagged descSets) \ds ->
    Debug.nameObject device ds "Sun.DS"
  Debug.nameObject device (Buffer.aBuffer sunData) "Sun.Data"

  pure (descSets, sunData)

dpSizes :: DescriptorSet.TypeMap Word32
dpSizes =
  [ ( Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
    , uniformBuffers
    )
  -- XXX: may be required to fetch textures for shadows from texture-masked models
  -- , ( Vk.DESCRIPTOR_TYPE_SAMPLED_IMAGE
  --   , sampledImages
  --   )
  -- , ( Vk.DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER
  --   , sampledImages + shadowMaps
  --   )
  -- , ( Vk.DESCRIPTOR_TYPE_SAMPLER
  --   , staticSamplers
  --   )
  ]
  where
    uniformBuffers = 2    -- 1 scene + 1 light array
    -- sampledImages  = 128  -- max dynamic textures and cubemaps
    -- staticSamplers = 8    -- immutable samplers
    -- shadowMaps     = 2    -- max shadowmaps

updateSet0Ds
  :: Tagged '[Sun] (Vector Vk.DescriptorSet)
  -> Buffer.Allocated 'Buffer.Coherent Sun
  -> ResourceT (StageRIO st) ()
updateSet0Ds (Tagged ds) Buffer.Allocated{aBuffer} = do
  context <- asks id
  Vk.updateDescriptorSets (getDevice context) writeSets mempty

  where
    destSet0 = case Vector.headM ds of
      Nothing ->
        error "assert: descriptor sets promised to contain [Sun]"
      Just one ->
        one

    writeSet0b0 = SomeStruct zero
      { Vk.dstSet          = destSet0
      , Vk.dstBinding      = 0
      , Vk.dstArrayElement = 0
      , Vk.descriptorCount = 1
      , Vk.descriptorType  = Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
      , Vk.bufferInfo      = Vector.singleton set0bind0I
      }
      where
        set0bind0I = Vk.DescriptorBufferInfo
          { Vk.buffer = aBuffer
          , Vk.offset = 0
          , Vk.range  = Vk.WHOLE_SIZE
          }

    writeSets =
      Vector.singleton writeSet0b0

data SunInput = SunInput
  { siColor :: Vec4

  , siInclination :: Float
  , siAzimuth     :: Float
  , siRadius      :: Float
  , siTarget      :: Vec3

  , siDepthRange :: Float
  , siSize       :: Float
  , siShadowIx   :: Float
  }

initialSunInput :: SunInput
initialSunInput = SunInput
  { siColor = vec4 1 1 1 1

  , siInclination = τ/8
  , siAzimuth     = -τ/8
  , siRadius      = Camera.PROJECTION_FAR / 2
  , siTarget      = 0

  , siDepthRange = Camera.PROJECTION_FAR
  , siSize       = 512
  , siShadowIx   = -1
  }

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

spawn1 :: MonadUnliftIO m => SunInput -> m Process
spawn1 = Worker.spawnCell mkSun

mkSun :: SunInput -> ("bounding box" ::: Transform, Sun)
mkSun SunInput{..} =
  ( bbTransform
  , Sun
      { sunViewProjection = mconcat vp
      , sunShadow         = vec4 0 0 siShadowIx siSize
      , sunPosition       = Vec4.fromVec3 position 0
      , sunDirection      = Vec4.fromVec3 direction 0
      , sunColor          = siColor
      }
  )
  where
    vp =
      [ Transform.rotateY (-siAzimuth)
      , Transform.rotateX (-siInclination)

      , Transform.translate 0 0 siRadius

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

    position = Transform.apply (vec3 0 0 siRadius) rotation

    direction = Transform.apply (vec3 0 0 $ -1) rotation

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

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

    rotation = mconcat
      [ Transform.rotateX siInclination
      , Transform.rotateY siAzimuth
      ]

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

newObserver1 :: MonadIO m => m Observer
newObserver1 = Worker.newObserverIO mempty

observe1 :: MonadUnliftIO m => Process -> Observer -> Buffer -> m ()
observe1 sunP sunOut sunData =
  Worker.observeIO_ sunP sunOut \_oldBB (bb, sun) -> do
    -- XXX: must stay the same or descsets must be updated with a new buffer
    _same <- Buffer.updateCoherent (VectorS.singleton sun) sunData
    pure $ VectorS.singleton bb

τ :: Float
τ = 2 * pi