Pipeline.hs

{-# LANGUAGE AllowAmbiguousTypes #-}

module Mischief.Render.Material.Pipeline where

import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.State (MonadState (get), put)
import Data.Data
import Data.Foldable
import Data.Kind (Type)
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Text qualified as T
import Data.Text.Encoding qualified as T
import Foreign (Ptr, nullPtr, with)
import Foreign.C.ConstPtr
import GHC.TypeLits
import Language.Haskell.TH (Extension (AllowAmbiguousTypes))
import Mischief.Assets (Asset (..))
import Mischief.ECS
import Mischief.Math qualified as Math
import Mischief.Render.Core
import Mischief.Render.Material.Definition
import Mischief.Render.Texture
import Mischief.Render.TextureFormat
import Mischief.Shader hiding (FragmentOutput, VertexOutput)
import Mischief.Shader.Generation hiding (FragmentOutput)
import Mischief.Shader.Internal.Bindings (Bindable, ErasedBinding (..), ErasedBindingValue, createBindGroup, createBindLayout, dummyBinding)
import Mischief.Shader.Internal.Params
import Mischief.Shader.Internal.State
import Mischief.Shader.Types qualified as ST
import Mischief.WGPU
import Mischief.WGPU.Framework (loadShaderFromBytes)
import Mischief.WGPU.Opaque
import Mischief.WGPU.Types.Enums
import Mischief.WGPU.Types.General

data PipelineKey = PipelineKey
  { materialType :: TypeRep,
    format :: WGPUTextureFormat
  }
  deriving (Eq, Ord)

newtype PipelineCache = PipelineCache (Map PipelineKey RenderPipeline) deriving anyclass (Component)

qgetPipeline :: forall mat. (Material mat) => TextureFormat -> Query System RenderPipeline
qgetPipeline format' = do
  let format = toWGPUFormat format'
  device <- qres
  qlift $ do
    PipelineCache cache <- resOrInsert $ PipelineCache Map.empty

    let key =
          PipelineKey
            { materialType = typeRep $ Proxy @mat,
              format
            }

    case Map.lookup key cache of
      Just pipeline -> pure pipeline
      Nothing -> liftIO $ createPipeline @mat device format

createPipeline :: forall mat. (Material mat) => RenderDevice -> WGPUTextureFormat -> IO RenderPipeline
createPipeline (RenderDevice device) format = do
  BindLayout bindLayout <- createBindLayout (Set.toList $ collectBindings @mat) (RenderDevice device)

  pipelineLayoutDesc <- with bindLayout $ \bindGroupLayout -> do
    pure $
      WGPUPipelineLayoutDescriptor
        { nextInChain = nullPtr,
          label = WGPUStringView {_data = ConstPtr nullPtr, length = 0},
          bindGroupLayoutCount = 1,
          bindGroupLayouts = ConstPtr {unConstPtr = bindGroupLayout},
          immediateSize = 0
        }

  pipelineLayout <- with pipelineLayoutDesc $ \pipelineLayoutDesc -> do
    wgpuDeviceCreatePipelineLayout device (ConstPtr {unConstPtr = pipelineLayoutDesc})

  -- putStrLn $ "------------------------------------------\n" ++ T.unpack (genShader @mat @(VO mat) @FragmentOutput "fragment" (collectBindings @mat) (convertFrag $ fragment @mat))

  -- putStrLn $ "------------------------------------------\n" ++ T.unpack (genShader @mat @(VI mat) @(VO mat) "vertex" (Set.toList $ collectBindings @mat) (vertex @mat))

  vertShader <- loadShaderFromBytes device (T.encodeUtf8 $ genShader @mat @(VI mat) @(VO mat) "vertex" (Set.toList $ collectBindings @mat) (vertex @mat))
  fragShader <- loadShaderFromBytes device (T.encodeUtf8 $ genShader @mat @(VO mat) @FragmentOutput "fragment" (Set.toList $ collectBindings @mat) (convertFrag $ fragment @mat))

  pipeline <- withWGPUString "main" $ \vertexEntry -> do
    withWGPUString "main" $ \fragmentEntry -> do
      withWGPUString "render pipeline" $ \label -> do
        let multisample =
              WGPUMultisampleState
                { nextInChain = nullPtr,
                  count = 1,
                  mask = 0xFFFFFFFF,
                  alphaToCoverageEnabled = wgpuFalse
                }

        let primitive =
              newWGPUPrimitiveState
                { topology = wGPUPrimitiveTopology_TriangleList
                }
        with blendState $ \blend -> do
          let target =
                WGPUColorTargetState
                  { nextInChain = nullPtr,
                    format,
                    blend = ConstPtr blend,
                    writeMask = wGPUColorWriteMask_All
                  }

          let vertex =
                newVertexState
                  { entryPoint = vertexEntry,
                    _module = vertShader,
                    _WGPUVertexState = Proxy
                  }

          with target $ \target -> do
            let fragment =
                  newFragmentState
                    { entryPoint = fragmentEntry,
                      _module = fragShader,
                      _WGPUFragmentState = Proxy,
                      targetCount = 1,
                      targets = ConstPtr target
                    }

            with fragment $ \fragment -> do
              let pipelineDesc =
                    WGPURenderPipelineDescriptor
                      { nextInChain = nullPtr,
                        label,
                        layout = pipelineLayout,
                        vertex,
                        multisample,
                        primitive,
                        fragment = ConstPtr fragment,
                        depthStencil = ConstPtr nullPtr
                      }

              with pipelineDesc $ \desc -> wgpuDeviceCreateRenderPipeline device desc

  pure $ RenderPipeline pipeline