{-# LANGUAGE AllowAmbiguousTypes #-}

module Mischief.Render.Buffer where

import Control.Monad.IO.Class
import Data.Data (Typeable)
import Foreign (Ptr, allocaBytes, castPtr, nullPtr, with, (.|.))
import Foreign.C.ConstPtr (ConstPtr (..))
import Language.Haskell.TH (Extension (AllowAmbiguousTypes))
import Mischief.ECS
import Mischief.Render.Core
import Mischief.Shader.Internal.Buffers
import Mischief.Shader.Internal.State (CPU)
import Mischief.WGPU (wgpuDeviceCreateBuffer, wgpuQueueWriteBuffer, withWGPUString)
import Mischief.WGPU.Opaque (WGPUBuffer)
import Mischief.WGPU.Types.Enums
import Mischief.WGPU.Types.General

data Buffer a = Buffer {forall {k} (a :: k). Buffer a -> Ptr WGPUBuffer
buffer :: Ptr WGPUBuffer, forall {k} (a :: k). Buffer a -> BufferDescriptor
desc :: BufferDescriptor} deriving (Typeable (Buffer a)
[HookRel (Buffer a)]
[Hook (Buffer a)]
Set DefaultComponentType
IsExclusive (IsExclusiveRel (Buffer a))
(Typeable (Buffer a), IsExclusive (IsExclusiveRel (Buffer a))) =>
Set DefaultComponentType
-> [Hook (Buffer a)]
-> [Hook (Buffer a)]
-> [Hook (Buffer a)]
-> [HookRel (Buffer a)]
-> [HookRel (Buffer a)]
-> [HookRel (Buffer a)]
-> Component (Buffer a)
forall c.
(Typeable c, IsExclusive (IsExclusiveRel c)) =>
Set DefaultComponentType
-> [Hook c]
-> [Hook c]
-> [Hook c]
-> [HookRel c]
-> [HookRel c]
-> [HookRel c]
-> Component c
forall k (a :: k). (Typeable a, Typeable k) => Typeable (Buffer a)
forall k (a :: k). (Typeable a, Typeable k) => [HookRel (Buffer a)]
forall k (a :: k). (Typeable a, Typeable k) => [Hook (Buffer a)]
forall k (a :: k).
(Typeable a, Typeable k) =>
Set DefaultComponentType
forall k (a :: k).
(Typeable a, Typeable k) =>
IsExclusive (IsExclusiveRel (Buffer a))
$crequired :: forall k (a :: k).
(Typeable a, Typeable k) =>
Set DefaultComponentType
required :: Set DefaultComponentType
$conAdd :: forall k (a :: k). (Typeable a, Typeable k) => [Hook (Buffer a)]
onAdd :: [Hook (Buffer a)]
$conSet :: forall k (a :: k). (Typeable a, Typeable k) => [Hook (Buffer a)]
onSet :: [Hook (Buffer a)]
$conRemove :: forall k (a :: k). (Typeable a, Typeable k) => [Hook (Buffer a)]
onRemove :: [Hook (Buffer a)]
$conAddRel :: forall k (a :: k). (Typeable a, Typeable k) => [HookRel (Buffer a)]
onAddRel :: [HookRel (Buffer a)]
$conSetRel :: forall k (a :: k). (Typeable a, Typeable k) => [HookRel (Buffer a)]
onSetRel :: [HookRel (Buffer a)]
$conRemoveRel :: forall k (a :: k). (Typeable a, Typeable k) => [HookRel (Buffer a)]
onRemoveRel :: [HookRel (Buffer a)]
Component)

data BufferDescriptor = BufferDescriptor
  { BufferDescriptor -> String
label :: String,
    BufferDescriptor -> Int
size :: Int,
    BufferDescriptor -> [BufferUsage]
usages :: [BufferUsage],
    BufferDescriptor -> Bool
mappedAtCreation :: Bool
  }

data BufferUsage
  = BufferUsageNone
  | BufferUsageMapRead
  | BufferUsageMapWrite
  | BufferUsageCopySrc
  | BufferUsageCopyDst
  | BufferUsageIndex
  | BufferUsageVertex
  | BufferUsageUniform
  | BufferUsageStorage
  | BufferUsageIndirect
  | BufferUsageQueryResolve

usageBits :: BufferUsage -> WGPUBufferUsage
usageBits :: BufferUsage -> WGPUBufferUsage
usageBits = \case
  BufferUsage
BufferUsageNone -> WGPUBufferUsage
wGPUBufferUsage_None
  BufferUsage
BufferUsageMapRead -> WGPUBufferUsage
wGPUBufferUsage_MapRead
  BufferUsage
BufferUsageMapWrite -> WGPUBufferUsage
wGPUBufferUsage_MapWrite
  BufferUsage
BufferUsageCopySrc -> WGPUBufferUsage
wGPUBufferUsage_CopySrc
  BufferUsage
BufferUsageCopyDst -> WGPUBufferUsage
wGPUBufferUsage_CopyDst
  BufferUsage
BufferUsageIndex -> WGPUBufferUsage
wGPUBufferUsage_Index
  BufferUsage
BufferUsageVertex -> WGPUBufferUsage
wGPUBufferUsage_Vertex
  BufferUsage
BufferUsageUniform -> WGPUBufferUsage
wGPUBufferUsage_Uniform
  BufferUsage
BufferUsageStorage -> WGPUBufferUsage
wGPUBufferUsage_Storage
  BufferUsage
BufferUsageIndirect -> WGPUBufferUsage
wGPUBufferUsage_Indirect
  BufferUsage
BufferUsageQueryResolve -> WGPUBufferUsage
wGPUBufferUsage_QueryResolve

processUsages :: [BufferUsage] -> WGPUBufferUsage
processUsages :: [BufferUsage] -> WGPUBufferUsage
processUsages = (BufferUsage -> WGPUBufferUsage -> WGPUBufferUsage)
-> WGPUBufferUsage -> [BufferUsage] -> WGPUBufferUsage
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (WGPUBufferUsage -> WGPUBufferUsage -> WGPUBufferUsage
forall a. Bits a => a -> a -> a
(.|.) (WGPUBufferUsage -> WGPUBufferUsage -> WGPUBufferUsage)
-> (BufferUsage -> WGPUBufferUsage)
-> BufferUsage
-> WGPUBufferUsage
-> WGPUBufferUsage
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BufferUsage -> WGPUBufferUsage
usageBits) (WGPUFlags -> WGPUBufferUsage
WGPUBufferUsage (Word64 -> WGPUFlags
WGPUFlags Word64
0))

createBuffer' :: RenderDevice -> BufferDescriptor -> System (Buffer a)
createBuffer' :: forall {k} (a :: k).
RenderDevice -> BufferDescriptor -> System (Buffer a)
createBuffer' (RenderDevice Ptr WGPUDevice
device) BufferDescriptor
descriptor = IO (Buffer a) -> System (Buffer a)
forall a. IO a -> System a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Buffer a) -> System (Buffer a))
-> IO (Buffer a) -> System (Buffer a)
forall a b. (a -> b) -> a -> b
$ do
  let BufferDescriptor {String
label :: BufferDescriptor -> String
label :: String
label, Int
size :: BufferDescriptor -> Int
size :: Int
size, [BufferUsage]
usages :: BufferDescriptor -> [BufferUsage]
usages :: [BufferUsage]
usages, Bool
mappedAtCreation :: BufferDescriptor -> Bool
mappedAtCreation :: Bool
mappedAtCreation} = BufferDescriptor
descriptor

  String -> (WGPUStringView -> IO (Buffer a)) -> IO (Buffer a)
forall a. String -> (WGPUStringView -> IO a) -> IO a
withWGPUString String
label ((WGPUStringView -> IO (Buffer a)) -> IO (Buffer a))
-> (WGPUStringView -> IO (Buffer a)) -> IO (Buffer a)
forall a b. (a -> b) -> a -> b
$ \WGPUStringView
label -> do
    let desc :: WGPUBufferDescriptor
desc =
          WGPUBufferDescriptor
            { nextInChain :: Ptr WGPUChainedStruct
nextInChain = Ptr WGPUChainedStruct
forall a. Ptr a
nullPtr,
              WGPUStringView
label :: WGPUStringView
label :: WGPUStringView
label,
              usage :: WGPUBufferUsage
usage = [BufferUsage] -> WGPUBufferUsage
processUsages [BufferUsage]
usages,
              size :: Word64
size = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
size,
              mappedAtCreation :: WGPUBool
mappedAtCreation = if Bool
mappedAtCreation then WGPUBool
wgpuTrue else WGPUBool
wgpuFalse
            }

    buffer <- WGPUBufferDescriptor
-> (Ptr WGPUBufferDescriptor -> IO (Ptr WGPUBuffer))
-> IO (Ptr WGPUBuffer)
forall a b. Storable a => a -> (Ptr a -> IO b) -> IO b
with WGPUBufferDescriptor
desc ((Ptr WGPUBufferDescriptor -> IO (Ptr WGPUBuffer))
 -> IO (Ptr WGPUBuffer))
-> (Ptr WGPUBufferDescriptor -> IO (Ptr WGPUBuffer))
-> IO (Ptr WGPUBuffer)
forall a b. (a -> b) -> a -> b
$ Ptr WGPUDevice
-> ConstPtr WGPUBufferDescriptor -> IO (Ptr WGPUBuffer)
wgpuDeviceCreateBuffer Ptr WGPUDevice
device (ConstPtr WGPUBufferDescriptor -> IO (Ptr WGPUBuffer))
-> (Ptr WGPUBufferDescriptor -> ConstPtr WGPUBufferDescriptor)
-> Ptr WGPUBufferDescriptor
-> IO (Ptr WGPUBuffer)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr WGPUBufferDescriptor -> ConstPtr WGPUBufferDescriptor
forall a. Ptr a -> ConstPtr a
ConstPtr
    pure $ Buffer {buffer, desc = descriptor}

createBuffer :: forall a. (Bufferable a) => RenderDevice -> System (Buffer a)
createBuffer :: forall (a :: * -> *).
Bufferable a =>
RenderDevice -> System (Buffer a)
createBuffer RenderDevice
device = do
  RenderDevice -> BufferDescriptor -> System (Buffer a)
forall {k} (a :: k).
RenderDevice -> BufferDescriptor -> System (Buffer a)
createBuffer'
    RenderDevice
device
    BufferDescriptor
      { label :: String
label = String
"buffer",
        size :: Int
size = forall (a :: * -> *). Bufferable a => Int
getSize @a,
        usages :: [BufferUsage]
usages = [BufferUsage
BufferUsageCopyDst, BufferUsage
BufferUsageUniform],
        mappedAtCreation :: Bool
mappedAtCreation = Bool
False
      }

uploadBuffer :: forall a. (Bufferable a) => RenderQueue -> Buffer a -> a CPU -> System ()
uploadBuffer :: forall (a :: * -> *).
Bufferable a =>
RenderQueue -> Buffer a -> a CPU -> System ()
uploadBuffer (RenderQueue Ptr WGPUQueue
queue) Buffer {Ptr WGPUBuffer
buffer :: forall {k} (a :: k). Buffer a -> Ptr WGPUBuffer
buffer :: Ptr WGPUBuffer
buffer, desc :: forall {k} (a :: k). Buffer a -> BufferDescriptor
desc = BufferDescriptor {Int
size :: BufferDescriptor -> Int
size :: Int
size}} a CPU
a = IO () -> System ()
forall a. IO a -> System a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> System ()) -> IO () -> System ()
forall a b. (a -> b) -> a -> b
$ do
  Int -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. Int -> (Ptr a -> IO b) -> IO b
allocaBytes Int
size ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr -> do
    _ <- a CPU -> Ptr Word8 -> IO Int
forall (a :: * -> *). Bufferable a => a CPU -> Ptr Word8 -> IO Int
putBytes a CPU
a Ptr Word8
ptr
    wgpuQueueWriteBuffer queue buffer 0 (ConstPtr $ castPtr ptr) (fromIntegral size)

maintainBuffer :: forall a. (Typeable a, Component (a CPU), Bufferable a) => System ()
maintainBuffer :: forall (a :: * -> *).
(Typeable a, Component (a CPU), Bufferable a) =>
System ()
maintainBuffer = Query System () -> System ()
forall w (m :: * -> *) out. MonadSystem w m => Query m out -> m ()
query_ (Query System () -> System ()) -> Query System () -> System ()
forall a b. (a -> b) -> a -> b
$ do
  device <- forall c (m :: * -> *) w.
(MonadSystem w m, Component c) =>
Query m c
qres @RenderDevice
  queue <- qres @RenderQueue

  (entity, buffer, dt) <- [q|Entity, Maybe (Buffer a), (a CPU) / Changed (a CPU) || Without (Buffer a)|]

  buffer <- case buffer of
    Just Buffer a
x -> Buffer a -> Query System (Buffer a)
forall a. a -> Query System a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Buffer a
x
    Maybe (Buffer a)
Nothing -> System (Buffer a) -> Query System (Buffer a)
forall w (m :: * -> *) a. MonadSystem w m => m a -> Query m a
qlift (System (Buffer a) -> Query System (Buffer a))
-> System (Buffer a) -> Query System (Buffer a)
forall a b. (a -> b) -> a -> b
$ do
      buffer <- forall (a :: * -> *).
Bufferable a =>
RenderDevice -> System (Buffer a)
createBuffer @a RenderDevice
device
      insert buffer entity
      pure buffer

  qlift $ uploadBuffer queue buffer dt