{-# 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
{ forall f. Binds f -> Binding f Sampler
sampler :: Binding f Sampler,
forall f. Binds f -> Binding f Texture
texture :: Binding f Texture
}
deriving ((forall x. Binds f -> Rep (Binds f) x)
-> (forall x. Rep (Binds f) x -> Binds f) -> Generic (Binds f)
forall x. Rep (Binds f) x -> Binds f
forall x. Binds f -> Rep (Binds f) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall f x. Rep (Binds f) x -> Binds f
forall f x. Binds f -> Rep (Binds f) x
$cfrom :: forall f x. Binds f -> Rep (Binds f) x
from :: forall x. Binds f -> Rep (Binds f) x
$cto :: forall f x. Rep (Binds f) x -> Binds f
to :: forall x. Rep (Binds f) x -> Binds f
Generic, Typeable Binds
Typeable Binds =>
(Nat -> Proxy (Binds Internal) -> (Bindings, Nat))
-> (Nat
-> Proxy (Binds Internal) -> ([WGPUBindGroupLayoutEntry], Nat))
-> (Nat -> Binds CPU -> ([WGPUBindGroupEntry], Nat))
-> (Nat -> (Binds GPU, Nat))
-> Bindable Binds
Nat -> (Binds GPU, Nat)
Nat -> Proxy (Binds Internal) -> ([WGPUBindGroupLayoutEntry], Nat)
Nat -> Proxy (Binds Internal) -> (Bindings, Nat)
Nat -> Binds CPU -> ([WGPUBindGroupEntry], Nat)
forall (a :: * -> *).
Typeable a =>
(Nat -> Proxy (a Internal) -> (Bindings, Nat))
-> (Nat -> Proxy (a Internal) -> ([WGPUBindGroupLayoutEntry], Nat))
-> (Nat -> a CPU -> ([WGPUBindGroupEntry], Nat))
-> (Nat -> (a GPU, Nat))
-> Bindable a
$ccollectBindings :: Nat -> Proxy (Binds Internal) -> (Bindings, Nat)
collectBindings :: Nat -> Proxy (Binds Internal) -> (Bindings, Nat)
$ccollectLayouts :: Nat -> Proxy (Binds Internal) -> ([WGPUBindGroupLayoutEntry], Nat)
collectLayouts :: Nat -> Proxy (Binds Internal) -> ([WGPUBindGroupLayoutEntry], Nat)
$ccollectEntries :: Nat -> Binds CPU -> ([WGPUBindGroupEntry], Nat)
collectEntries :: Nat -> Binds CPU -> ([WGPUBindGroupEntry], Nat)
$cdummyBinding :: Nat -> (Binds GPU, Nat)
dummyBinding :: Nat -> (Binds GPU, Nat)
Bindable)
data VertexOutput f = VertexOutput
{ forall f. VertexOutput f -> BuiltIn f "position" (Vec 4 Float)
position :: BuiltIn f "position" (Vec 4 Float),
forall f. VertexOutput f -> Location f 0 (Vec 2 Float)
uv :: Location f 0 (Vec 2 Float)
}
deriving ((forall x. VertexOutput f -> Rep (VertexOutput f) x)
-> (forall x. Rep (VertexOutput f) x -> VertexOutput f)
-> Generic (VertexOutput f)
forall x. Rep (VertexOutput f) x -> VertexOutput f
forall x. VertexOutput f -> Rep (VertexOutput f) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall f x. Rep (VertexOutput f) x -> VertexOutput f
forall f x. VertexOutput f -> Rep (VertexOutput f) x
$cfrom :: forall f x. VertexOutput f -> Rep (VertexOutput f) x
from :: forall x. VertexOutput f -> Rep (VertexOutput f) x
$cto :: forall f x. Rep (VertexOutput f) x -> VertexOutput f
to :: forall x. Rep (VertexOutput f) x -> VertexOutput f
Generic, Typeable (VertexOutput GPU)
Params
ParamDefinition
VertexOutput GPU
Typeable (VertexOutput GPU) =>
Params
-> ParamDefinition
-> VertexOutput GPU
-> (VertexOutput GPU -> Values)
-> ShaderParam VertexOutput
VertexOutput GPU -> Values
forall (a :: * -> *).
Typeable (a GPU) =>
Params
-> ParamDefinition -> a GPU -> (a GPU -> Values) -> ShaderParam a
$ccollectParams :: Params
collectParams :: Params
$cparamDefinition :: ParamDefinition
paramDefinition :: ParamDefinition
$cdummyParam :: VertexOutput GPU
dummyParam :: VertexOutput GPU
$ccollectValues :: VertexOutput GPU -> Values
collectValues :: VertexOutput GPU -> Values
ShaderParam)
newtype FragmentOutput f = FragmentOutput (Location f 0 (Vec 4 Float))
deriving stock ((forall x. FragmentOutput f -> Rep (FragmentOutput f) x)
-> (forall x. Rep (FragmentOutput f) x -> FragmentOutput f)
-> Generic (FragmentOutput f)
forall x. Rep (FragmentOutput f) x -> FragmentOutput f
forall x. FragmentOutput f -> Rep (FragmentOutput f) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall f x. Rep (FragmentOutput f) x -> FragmentOutput f
forall f x. FragmentOutput f -> Rep (FragmentOutput f) x
$cfrom :: forall f x. FragmentOutput f -> Rep (FragmentOutput f) x
from :: forall x. FragmentOutput f -> Rep (FragmentOutput f) x
$cto :: forall f x. Rep (FragmentOutput f) x -> FragmentOutput f
to :: forall x. Rep (FragmentOutput f) x -> FragmentOutput f
Generic)
deriving anyclass (Typeable (FragmentOutput GPU)
Params
ParamDefinition
FragmentOutput GPU
Typeable (FragmentOutput GPU) =>
Params
-> ParamDefinition
-> FragmentOutput GPU
-> (FragmentOutput GPU -> Values)
-> ShaderParam FragmentOutput
FragmentOutput GPU -> Values
forall (a :: * -> *).
Typeable (a GPU) =>
Params
-> ParamDefinition -> a GPU -> (a GPU -> Values) -> ShaderParam a
$ccollectParams :: Params
collectParams :: Params
$cparamDefinition :: ParamDefinition
paramDefinition :: ParamDefinition
$cdummyParam :: FragmentOutput GPU
dummyParam :: FragmentOutput GPU
$ccollectValues :: FragmentOutput GPU -> Values
collectValues :: FragmentOutput GPU -> Values
ShaderParam)
shaderToWGSL :: forall a m. (ShaderParam a) => Map TypeRep ErasedValue -> Shader m (a GPU) -> Text
shaderToWGSL :: forall {k} (a :: * -> *) (m :: k).
ShaderParam a =>
Map TypeRep ErasedValue -> Shader m (a GPU) -> Text
shaderToWGSL Map TypeRep ErasedValue
dict (Shader State ShaderState (a GPU)
s) = do
let (a GPU
a, ShaderState {Stmt
ast :: Stmt
ast :: ShaderState -> Stmt
ast}) = State ShaderState (a GPU) -> ShaderState -> (a GPU, ShaderState)
forall s a. State s a -> s -> (a, s)
runState State ShaderState (a GPU)
s ShaderState {counter :: VarIndex
counter = VarIndex
0, ast :: Stmt
ast = Stmt
Empty, Map TypeRep ErasedValue
dict :: Map TypeRep ErasedValue
dict :: Map TypeRep ErasedValue
dict}
let ast' :: Stmt
ast' = case (forall (a :: * -> *). ShaderParam a => ParamDefinition
paramDefinition @a) of
StructParam Text
n Params
_ -> do
let Values [Value]
vs = a GPU -> Values
forall (a :: * -> *). ShaderParam a => a GPU -> Values
collectValues a GPU
a
let vst :: [Text]
vst = (Value -> Text) -> [Value] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Value -> Text
valueToWGSL [Value]
vs
Stmt -> Stmt -> Stmt
Concat Stmt
ast (Expr (ZonkAny 0) -> Stmt
forall (a :: Types). Expr a -> Stmt
Return (Text -> [Text] -> Expr (ZonkAny 0)
forall (a :: Types). Text -> [Text] -> Expr a
StructInit Text
n [Text]
vst))
AnonParam Text
_ Text
_ -> do
let Values [Value
v] = a GPU -> Values
forall (a :: * -> *). ShaderParam a => a GPU -> Values
collectValues a GPU
a
Stmt -> Stmt -> Stmt
Concat Stmt
ast (Expr (ZonkAny 1) -> Stmt
forall (a :: Types). Expr a -> Stmt
Return (Value -> Expr (ZonkAny 1)
forall (a :: Types). Value -> Expr a
unwrapValue Value
v))
Stmt -> Text
stmtToWGSL Stmt
ast'
unwrapValue :: forall a. Value -> Expr a
unwrapValue :: forall (a :: Types). Value -> Expr a
unwrapValue (Value (Expr a
a :: Expr b)) = Expr a -> Expr a
forall a b. a -> b
unsafeCoerce Expr a
a
valueToWGSL :: Value -> Text
valueToWGSL :: Value -> Text
valueToWGSL (Value (Expr a
a :: Expr a)) = Expr a -> Text
forall (a :: Types). Expr a -> Text
exprToWGSL Expr a
a
genFunction :: forall a m. (ShaderParam a) => Text -> Text -> Text -> Map TypeRep ErasedValue -> Shader m (a GPU) -> Text
genFunction :: forall {k} (a :: * -> *) (m :: k).
ShaderParam a =>
Text
-> Text
-> Text
-> Map TypeRep ErasedValue
-> Shader m (a GPU)
-> Text
genFunction Text
name Text
input Text
output Map TypeRep ErasedValue
dict Shader m (a GPU)
s = Text
"fn " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
name Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"(" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
input Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
") -> " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
output Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" {\n" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Map TypeRep ErasedValue -> Shader m (a GPU) -> Text
forall {k} (a :: * -> *) (m :: k).
ShaderParam a =>
Map TypeRep ErasedValue -> Shader m (a GPU) -> Text
shaderToWGSL Map TypeRep ErasedValue
dict Shader m (a GPU)
s Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\n}"
genInput :: ParamDefinition -> Text
genInput :: ParamDefinition -> Text
genInput (StructParam Text
name Params
_) = Text
"input: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
name
genInput (AnonParam Text
n Text
m) = Text
n Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"input: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
m
genOutput :: ParamDefinition -> Text
genOutput :: ParamDefinition -> Text
genOutput (StructParam Text
name Params
_) = Text
name
genOutput (AnonParam Text
n Text
m) = Text
n Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
m
genShader :: forall m p a. (ShaderParam p, ShaderParam a) => Text -> [ErasedBinding] -> (p GPU -> Shader m (a GPU)) -> Text
genShader :: forall {k} (m :: k) (p :: * -> *) (a :: * -> *).
(ShaderParam p, ShaderParam a) =>
Text -> [ErasedBinding] -> (p GPU -> Shader m (a GPU)) -> Text
genShader Text
tag [ErasedBinding]
b p GPU -> Shader m (a GPU)
s = Nat -> [ErasedBinding] -> Text
genBindings Nat
0 [ErasedBinding]
b Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Proxy p -> Text
forall (p :: * -> *). ShaderParam p => Proxy p -> Text
genParams (forall {k} (t :: k). Proxy t
forall (t :: * -> *). Proxy t
Proxy @p) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Proxy a -> Text
forall (p :: * -> *). ShaderParam p => Proxy p -> Text
genParams (forall {k} (t :: k). Proxy t
forall (t :: * -> *). Proxy t
Proxy @a) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"@" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
tag Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\n" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
-> Text
-> Text
-> Map TypeRep ErasedValue
-> Shader m (a GPU)
-> Text
forall {k} (a :: * -> *) (m :: k).
ShaderParam a =>
Text
-> Text
-> Text
-> Map TypeRep ErasedValue
-> Shader m (a GPU)
-> Text
genFunction Text
"main" (ParamDefinition -> Text
genInput (ParamDefinition -> Text) -> ParamDefinition -> Text
forall a b. (a -> b) -> a -> b
$ forall (a :: * -> *). ShaderParam a => ParamDefinition
paramDefinition @p) (ParamDefinition -> Text
genOutput (ParamDefinition -> Text) -> ParamDefinition -> Text
forall a b. (a -> b) -> a -> b
$ forall (a :: * -> *). ShaderParam a => ParamDefinition
paramDefinition @a) (Nat -> [ErasedBinding] -> Map TypeRep ErasedValue
genDict Nat
0 [ErasedBinding]
b) (p GPU -> Shader m (a GPU)
s p GPU
forall (a :: * -> *). ShaderParam a => a GPU
dummyParam)
genDict :: Nat -> [ErasedBinding] -> Map TypeRep ErasedValue
genDict :: Nat -> [ErasedBinding] -> Map TypeRep ErasedValue
genDict Nat
_ [] = Map TypeRep ErasedValue
forall k a. Map k a
Map.empty
genDict Nat
n (ErasedBinding m
_ Proxy a
m (Proxy b
_ :: Proxy b) : [ErasedBinding]
xs) =
let (b GPU
y, Nat
n') = forall (a :: * -> *). Bindable a => Nat -> (a GPU, Nat)
dummyBinding @b Nat
n
in Map TypeRep ErasedValue
-> Map TypeRep ErasedValue -> Map TypeRep ErasedValue
forall k a. Ord k => Map k a -> Map k a -> Map k a
Map.union (TypeRep -> ErasedValue -> Map TypeRep ErasedValue
forall k a. k -> a -> Map k a
Map.singleton (Proxy a -> TypeRep
forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep Proxy a
m) (b GPU -> ErasedValue
forall a. a -> ErasedValue
ErasedValue b GPU
y)) (Nat -> [ErasedBinding] -> Map TypeRep ErasedValue
genDict Nat
n' [ErasedBinding]
xs)
genBindings :: Nat -> [ErasedBinding] -> Text
genBindings :: Nat -> [ErasedBinding] -> Text
genBindings Nat
_ [] = Text
""
genBindings Nat
n (ErasedBinding m
_ Proxy a
_ (Proxy b
_ :: Proxy b) : [ErasedBinding]
xs) =
let (Bindings [BindingData]
x, Nat
n') = Nat -> Proxy (b Internal) -> (Bindings, Nat)
forall (a :: * -> *).
Bindable a =>
Nat -> Proxy (a Internal) -> (Bindings, Nat)
collectBindings Nat
n (forall t. Proxy t
forall {k} (t :: k). Proxy t
Proxy @(b Internal))
in [Text] -> Text
T.concat ((BindingData -> Text) -> [BindingData] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map BindingData -> Text
genBinding [BindingData]
x) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\n" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Nat -> [ErasedBinding] -> Text
genBindings Nat
n' [ErasedBinding]
xs
genBinding :: BindingData -> Text
genBinding :: BindingData -> Text
genBinding NormalData {Text
bType :: Text
bType :: BindingData -> Text
bType, Integer
index :: Integer
index :: BindingData -> Integer
index} = Text
"@group(0) @binding(" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (Integer -> String
forall a. Show a => a -> String
show Integer
index) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
")\nvar b" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (Integer -> String
forall a. Show a => a -> String
show Integer
index) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" : " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
bType Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
";\n\n"
genBinding UniformData {Text
structName :: Text
structName :: BindingData -> Text
structName, [(Text, Text)]
structFields :: [(Text, Text)]
structFields :: BindingData -> [(Text, Text)]
structFields, Integer
index :: BindingData -> Integer
index :: Integer
index} =
Text
"@group(0) @binding(" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (Integer -> String
forall a. Show a => a -> String
show Integer
index) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
")\nvar<uniform> b" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (Integer -> String
forall a. Show a => a -> String
show Integer
index) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" : " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
structName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
";\n\n" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> [(Text, Text)] -> Text
genBuffer Text
structName [(Text, Text)]
structFields
genBuffer :: Text -> [(Text, Text)] -> Text
genBuffer :: Text -> [(Text, Text)] -> Text
genBuffer Text
name [(Text, Text)]
fields =
let fields' :: [Text]
fields' = ((Text, Text) -> Text) -> [(Text, Text)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (\(Text
x, Text
y) -> Text
" " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
x Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
": " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
y) [(Text, Text)]
fields
in Text
"struct " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
name Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" {\n" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> [Text] -> Text
T.intercalate Text
",\n" [Text]
fields' Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\n};\n\n"
genParams :: forall p. (ShaderParam p) => Proxy p -> Text
genParams :: forall (p :: * -> *). ShaderParam p => Proxy p -> Text
genParams Proxy p
_ =
let definition :: ParamDefinition
definition = forall (a :: * -> *). ShaderParam a => ParamDefinition
paramDefinition @p
in case ParamDefinition
definition of
StructParam Text
name (Params [ParamData]
params) -> Text
"struct " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
name Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" {\n" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Text] -> Text
T.concat ((ParamData -> Text) -> [ParamData] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map ParamData -> Text
genParam [ParamData]
params) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"};\n\n"
ParamDefinition
_ -> Text
""
genParam :: ParamData -> Text
genParam :: ParamData -> Text
genParam ParamData {Text
pType :: Text
pType :: ParamData -> Text
pType, index :: ParamData -> ParamKind
index = BuiltInKind Text
s} = Text
" @builtin(" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
s Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
") " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
s Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" : " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
pType Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
",\n"
genParam ParamData {Text
pType :: ParamData -> Text
pType :: Text
pType, index :: ParamData -> ParamKind
index = LocKind Integer
n} = Text
" @location(" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (Integer -> String
forall a. Show a => a -> String
show Integer
n) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
") l" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (Integer -> String
forall a. Show a => a -> String
show Integer
n) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" : " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
pType Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
",\n"
genShaderT :: forall m b a p. (Bindable b, ShaderParam p, ShaderParam a) => [ErasedBinding] -> (p GPU -> Shader m (a GPU)) -> String
genShaderT :: forall {k} (m :: k) (b :: * -> *) (a :: * -> *) (p :: * -> *).
(Bindable b, ShaderParam p, ShaderParam a) =>
[ErasedBinding] -> (p GPU -> Shader m (a GPU)) -> String
genShaderT [ErasedBinding]
p = Text -> String
T.unpack (Text -> String)
-> ((p GPU -> Shader m (a GPU)) -> Text)
-> (p GPU -> Shader m (a GPU))
-> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> [ErasedBinding] -> (p GPU -> Shader m (a GPU)) -> Text
forall {k} (m :: k) (p :: * -> *) (a :: * -> *).
(ShaderParam p, ShaderParam a) =>
Text -> [ErasedBinding] -> (p GPU -> Shader m (a GPU)) -> Text
genShader Text
"no_tag" [ErasedBinding]
p