Generation.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeAbstractions #-}

module Mischief.Shader.Generation where

import Control.Monad.State hiding (get)
import Data.Data hiding (cast)
import Data.Default
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Text (Text)
import Data.Text qualified as T
import GHC.Generics
import GHC.TypeLits (KnownNat, Nat, natVal)
import Mischief.Render.Core
import Mischief.Render.Texture
import Mischief.Shader.Functions
import Mischief.Shader.Internal.Bindings
import Mischief.Shader.Internal.Params
import Mischief.Shader.Internal.State
import Mischief.Shader.Types hiding ((<>))
import Unsafe.Coerce (unsafeCoerce)
import Prelude (String, map, pure, show, undefined, ($), (.), (<>))
import Prelude qualified

data Binds f = Binds
  { sampler :: Binding f Sampler,
    texture :: Binding f Texture
  }
  deriving (Generic, Bindable)

data VertexOutput f = VertexOutput
  { position :: BuiltIn f "position" (Vec 4 Float),
    uv :: Location f 0 (Vec 2 Float)
  }
  deriving (Generic, ShaderParam)

newtype FragmentOutput f = FragmentOutput (Location f 0 (Vec 4 Float))
  deriving stock (Generic)
  deriving anyclass (ShaderParam)

-- frag :: Binds -> BuiltIn "location" Vec2f -> Shader (Location 0 Vec4f)
-- frag b input = do
--   let result = sample (get b.texture) (get b.sampler) (get input)
--   pure (set result)

shaderToWGSL :: forall a m. (ShaderParam a) => Map TypeRep ErasedValue -> Shader m (a GPU) -> Text
shaderToWGSL dict (Shader s) = do
  let (a, ShaderState {ast}) = runState s ShaderState {counter = 0, ast = Empty, dict}
  let ast' = case (paramDefinition @a) of
        StructParam n _ -> do
          let Values vs = collectValues a
          let vst = map valueToWGSL vs
          Concat ast (Return (StructInit n vst))
        AnonParam _ _ -> do
          let Values [v] = collectValues a
          Concat ast (Return (unwrapValue v))
  stmtToWGSL ast'

unwrapValue :: forall a. Value -> Expr a
unwrapValue (Value (a :: Expr b)) = unsafeCoerce a

valueToWGSL :: Value -> Text
valueToWGSL (Value (a :: Expr a)) = exprToWGSL a

genFunction :: forall a m. (ShaderParam a) => Text -> Text -> Text -> Map TypeRep ErasedValue -> Shader m (a GPU) -> Text
genFunction name input output dict s = "fn " <> name <> "(" <> input <> ") -> " <> output <> " {\n" <> shaderToWGSL dict s <> "\n}"

genInput :: ParamDefinition -> Text
genInput (StructParam name _) = "input: " <> name
genInput (AnonParam n m) = n <> "input: " <> m

genOutput :: ParamDefinition -> Text
genOutput (StructParam name _) = name
genOutput (AnonParam n m) = n <> m

genShader :: forall m p a. (ShaderParam p, ShaderParam a) => Text -> [ErasedBinding] -> (p GPU -> Shader m (a GPU)) -> Text
genShader tag b s = genBindings 0 b <> genParams (Proxy @p) <> genParams (Proxy @a) <> "@" <> tag <> "\n" <> genFunction "main" (genInput $ paramDefinition @p) (genOutput $ paramDefinition @a) (genDict 0 b) (s dummyParam)

genDict :: Nat -> [ErasedBinding] -> Map TypeRep ErasedValue
genDict _ [] = Map.empty
genDict n (ErasedBinding _ m (_ :: Proxy b) : xs) =
  let (y, n') = dummyBinding @b n
   in Map.union (Map.singleton (typeRep m) (ErasedValue y)) (genDict n' xs)

genBindings :: Nat -> [ErasedBinding] -> Text
genBindings _ [] = ""
genBindings n (ErasedBinding _ _ (_ :: Proxy b) : xs) =
  let (Bindings x, n') = collectBindings n (Proxy @(b Internal))
   in T.concat (map genBinding x) <> "\n" <> genBindings n' xs

genBinding :: BindingData -> Text
genBinding NormalData {bType, index} = "@group(0) @binding(" <> T.pack (show index) <> ")\nvar b" <> T.pack (show index) <> " : " <> bType <> ";\n\n"
genBinding UniformData {structName, structFields, index} =
  "@group(0) @binding(" <> T.pack (show index) <> ")\nvar<uniform> b" <> T.pack (show index) <> " : " <> structName <> ";\n\n" <> genBuffer structName structFields

genBuffer :: Text -> [(Text, Text)] -> Text
genBuffer name fields =
  let fields' = map (\(x, y) -> "  " <> x <> ": " <> y) fields
   in "struct " <> name <> " {\n" <> T.intercalate ",\n" fields' <> "\n};\n\n"

genParams :: forall p. (ShaderParam p) => Proxy p -> Text
genParams _ =
  let definition = paramDefinition @p
   in case definition of
        StructParam name (Params params) -> "struct " <> name <> " {\n" <> T.concat (map genParam params) <> "};\n\n"
        _ -> ""

genParam :: ParamData -> Text
genParam ParamData {pType, index = BuiltInKind s} = "  @builtin(" <> s <> ") " <> s <> " : " <> pType <> ",\n"
genParam ParamData {pType, index = LocKind n} = "  @location(" <> T.pack (show n) <> ") l" <> T.pack (show n) <> " : " <> pType <> ",\n"

genShaderT :: forall m b a p. (Bindable b, ShaderParam p, ShaderParam a) => [ErasedBinding] -> (p GPU -> Shader m (a GPU)) -> String
genShaderT p = T.unpack . genShader "no_tag" p