mischief-render-0.1.0.0: src/Mischief/Shader/Internal/State.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeAbstractions #-}
{-# OPTIONS_GHC -Wno-redundant-constraints #-}

module Mischief.Shader.Internal.State where

import Control.Monad.State
import Data.Data hiding (cast)
import Data.Map (Map)
import Data.Singletons (SingI, demote)
import Data.Text (Text)
import Data.Text qualified as T
import GHC.TypeLits (Nat, type (+))
import Mischief.Shader.Internal.Singletons
import Unsafe.Coerce (unsafeCoerce)
import Prelude hiding (cos, sin)

data ErasedValue where
  ErasedValue :: a -> ErasedValue

data ShaderState = ShaderState
  { counter :: VarIndex,
    ast :: Stmt,
    dict :: Map TypeRep ErasedValue
  }

newtype VarIndex = VarIndex Integer
  deriving newtype (Show, Num)

nameIndex :: VarIndex -> Text
nameIndex i = "v" <> T.pack (show i)

data Stmt where
  Concat :: Stmt -> Stmt -> Stmt
  Let :: VarIndex -> Expr a -> Stmt
  Return :: Expr a -> Stmt
  Debug :: Text -> Stmt
  Empty :: Stmt

data CustomType a = CustomType

instance Eq (CustomType a) where
  _ == _ = False

data Expr (a :: Types) where
  Function :: Text -> [Param] -> Expr a
  Add :: (SingI a) => Expr a -> Expr a -> Expr a
  Sub :: (SingI a) => Expr a -> Expr a -> Expr a
  Mult :: (SingI a) => Expr a -> Expr a -> Expr a
  Div :: (SingI a) => Expr a -> Expr a -> Expr a
  Neg :: (SingI a) => Expr a -> Expr a
  Abs :: (SingI a) => Expr a -> Expr a
  Signum :: (SingI a) => Expr a -> Expr a
  Cast :: (SingI a, SingI b) => Proxy b -> Expr a -> Expr b
  CastVec :: (SingI (VectorType n a), SingI (VectorType n b)) => Proxy (VectorType n b) -> Expr (VectorType n a) -> Expr (VectorType n b)
  ConstantInt :: Int -> Expr (Primitive TInt)
  ConstantFloat :: Float -> Expr (Primitive TFloat)
  ConstantUInt :: Int -> Expr (Primitive TUInt)
  ConstantArray :: forall (a' :: Types). (SingI a') => (Proxy a') -> [Expr a'] -> Expr (ArrayType a')
  ConstantVec :: forall (n :: Nat) (a' :: PrimitiveTypes). (PrintVecElements n a', SingI (VectorType n a')) => Proxy (VectorType n a') -> VecInit n a' -> Expr (VectorType n a')
  ConstantMat :: forall (n :: Nat) (m :: Nat). (PrintMatElements n m, SingI (MatrixType n m)) => Proxy (MatrixType n m) -> MatInit n m -> Expr (MatrixType n m)
  AccessField :: (SingI a) => Expr a -> Text -> Expr b
  Index :: (SingI (ArrayType a)) => Expr (ArrayType a) -> Expr (Primitive TUInt) -> Expr a
  Var :: (SingI a) => VarIndex -> Expr a
  BindingVar :: (SingI a) => Integer -> Expr a
  BuiltInVar :: (SingI a) => Text -> Expr a
  StructInit :: Text -> [Text] -> Expr a
  LocationVar :: (SingI a) => Integer -> Expr a
  VarCustom :: Text -> Expr a
  CustomOperator :: Text -> Expr a -> Expr b -> Expr c
  Void :: (SingI a) => Expr a
  AppendVecs :: forall (n :: Nat) (m :: Nat) a. (SingI a, PrintVecElements (n + m) a, SingI (VectorType (n + m) a)) => Proxy (VectorType (n + m) a) -> Expr (VectorType n a) -> Expr (VectorType m a) -> Expr (VectorType (n + m) a)
  ExtendVec :: forall (n :: Nat) a. (SingI a, PrintVecElements (n + 1) a, SingI (VectorType (n + 1) a)) => Proxy (VectorType (n + 1) a) -> Expr (VectorType n a) -> Expr (Primitive a) -> Expr (VectorType (n + 1) a)
  ConstantBool :: Bool -> Expr TBool
  If :: Expr TBool -> Expr a -> Expr a -> Expr a

data Param where
  Param :: (SingI a) => Expr a -> Param

type family VectorOrSingleOf a where
  VectorOrSingleOf (VectorType n a) = a
  VectorOrSingleOf (Primitive a) = a

instance (SingI a) => Num (Expr (Primitive a)) where
  (+) = Add
  (*) = Mult
  abs = Abs
  signum = Signum
  fromInteger i =
    let t = demote @a
     in case t of
          TInt -> unsafeCoerce $ ConstantInt (fromInteger i)
          TFloat -> unsafeCoerce $ ConstantFloat (fromInteger i)
          TUInt -> unsafeCoerce $ ConstantUInt (fromInteger i)
  negate = Neg

instance Fractional (Expr (Primitive TFloat)) where
  fromRational = ConstantFloat . fromRational
  recip = Div 1
  (/) = Div

type family VecInit (n :: Nat) (a :: PrimitiveTypes) where
  VecInit 2 a = (Expr (Primitive a), Expr (Primitive a))
  VecInit 3 a = (Expr (Primitive a), Expr (Primitive a), Expr (Primitive a))
  VecInit 4 a = (Expr (Primitive a), Expr (Primitive a), Expr (Primitive a), Expr (Primitive a))

class PrintVecElements (n :: Nat) (a :: PrimitiveTypes) where
  printVecElements :: VecInit n a -> Text

instance PrintVecElements 2 a where
  printVecElements (a, b) = exprToWGSL a <> ", " <> exprToWGSL b

instance PrintVecElements 3 a where
  printVecElements (a, b, c) = exprToWGSL a <> ", " <> exprToWGSL b <> ", " <> exprToWGSL c

instance PrintVecElements 4 a where
  printVecElements (a, b, c, d) = exprToWGSL a <> ", " <> exprToWGSL b <> ", " <> exprToWGSL c <> ", " <> exprToWGSL d

type family ArrayInit (n :: Nat) a where
  ArrayInit 2 a = (Expr a, Expr a)
  ArrayInit 3 a = (Expr a, Expr a, Expr a)
  ArrayInit 4 a = (Expr a, Expr a, Expr a, Expr a)
  ArrayInit 5 a = (Expr a, Expr a, Expr a, Expr a, Expr a)
  ArrayInit 6 a = (Expr a, Expr a, Expr a, Expr a, Expr a, Expr a)

class PrintArrayElements (n :: Nat) (a :: Types) where
  printArrayElements :: ArrayInit n a -> Text

instance PrintArrayElements 2 a where
  printArrayElements (a, b) = exprToWGSL a <> ", " <> exprToWGSL b

instance PrintArrayElements 3 a where
  printArrayElements (a, b, c) = exprToWGSL a <> ", " <> exprToWGSL b <> ", " <> exprToWGSL c

instance PrintArrayElements 4 a where
  printArrayElements (a, b, c, d) = exprToWGSL a <> ", " <> exprToWGSL b <> ", " <> exprToWGSL c <> ", " <> exprToWGSL d

instance PrintArrayElements 5 a where
  printArrayElements (a, b, c, d, e) = exprToWGSL a <> ", " <> exprToWGSL b <> ", " <> exprToWGSL c <> ", " <> exprToWGSL d <> ", " <> exprToWGSL e

instance PrintArrayElements 6 a where
  printArrayElements (a, b, c, d, e, f) = exprToWGSL a <> ", " <> exprToWGSL b <> ", " <> exprToWGSL c <> ", " <> exprToWGSL d <> ", " <> exprToWGSL e <> ", " <> exprToWGSL f

type family MatInit (n :: Nat) (m :: Nat) where
  MatInit 2 a = (Expr (VectorType a TFloat), Expr (VectorType a TFloat))
  MatInit 3 a = (Expr (VectorType a TFloat), Expr (VectorType a TFloat), Expr (VectorType a TFloat))
  MatInit 4 a = (Expr (VectorType a TFloat), Expr (VectorType a TFloat), Expr (VectorType a TFloat), Expr (VectorType a TFloat))

class PrintMatElements (n :: Nat) (m :: Nat) where
  printMatElements :: MatInit n m -> Text

instance PrintMatElements 2 n where
  printMatElements (a, b) =
    exprToWGSL a <> ", " <> exprToWGSL b

instance PrintMatElements 3 n where
  printMatElements (a, b, c) =
    exprToWGSL a <> ", " <> exprToWGSL b <> ", " <> exprToWGSL c

instance PrintMatElements 4 n where
  printMatElements (a, b, c, d) =
    exprToWGSL a <> ", " <> exprToWGSL b <> ", " <> exprToWGSL c <> ", " <> exprToWGSL d

showArgs :: [Param] -> Text
showArgs [] = ""
showArgs [Param x] = exprToWGSL x
showArgs ((Param x) : xs) = exprToWGSL x <> ", " <> showArgs xs

exprToWGSL :: Expr a -> Text
exprToWGSL (Function t x) = t <> "(" <> showArgs x <> ")"
exprToWGSL (Add a b) = "(" <> exprToWGSL a <> " + " <> exprToWGSL b <> ")"
exprToWGSL (Mult a b) = "(" <> exprToWGSL a <> " * " <> exprToWGSL b <> ")"
exprToWGSL (Sub a b) = "(" <> exprToWGSL a <> " - " <> exprToWGSL b <> ")"
exprToWGSL (Div a b) = "(" <> exprToWGSL a <> " / " <> exprToWGSL b <> ")"
exprToWGSL (Neg a) = "(-" <> exprToWGSL a <> ")"
exprToWGSL (Abs a) = exprToWGSL (Function "abs" [Param a])
exprToWGSL (Signum a) = exprToWGSL (Function "sign" [Param a])
exprToWGSL (Cast (_ :: Proxy b) (a :: Expr a)) =
  let aType = demote @a
      bType = demote @b
   in if aType == bType
        then exprToWGSL a
        else case bType of
          Primitive TInt -> "i32(" <> exprToWGSL a <> ")"
          Primitive TFloat -> "f32(" <> exprToWGSL a <> ")"
          Primitive TUInt -> "u32(" <> exprToWGSL a <> ")"
          _ -> undefined
exprToWGSL (CastVec (_ :: Proxy b) (a :: Expr a)) =
  let aType = demote @a
      bType = demote @b
   in if aType == bType
        then exprToWGSL a
        else case bType of
          VectorType 2 TInt -> "vec2<i32>(" <> exprToWGSL a <> ")"
          VectorType 2 TFloat -> "vec2<f32>(" <> exprToWGSL a <> ")"
          VectorType 2 TUInt -> "vec2<u32>(" <> exprToWGSL a <> ")"
          VectorType 3 TInt -> "vec3<i32>(" <> exprToWGSL a <> ")"
          VectorType 3 TFloat -> "vec3<f32>(" <> exprToWGSL a <> ")"
          VectorType 3 TUInt -> "vec3<u32>(" <> exprToWGSL a <> ")"
          VectorType 4 TInt -> "vec4<i32>(" <> exprToWGSL a <> ")"
          VectorType 4 TFloat -> "vec4<f32>(" <> exprToWGSL a <> ")"
          VectorType 4 TUInt -> "vec4<u32>(" <> exprToWGSL a <> ")"
          _ -> undefined
exprToWGSL (ConstantArray (_ :: Proxy a) x) =
  let element = typeToWGSL' (demote @a)
   in "array<" <> element <> ", " <> T.show (length x) <> ">(" <> T.intercalate ", " (map exprToWGSL x) <> ")"
exprToWGSL (ConstantVec @n @a' (_ :: Proxy a) x) =
  let elements = printVecElements @n @a' x
   in case demote @a of
        VectorType 2 TInt -> "vec2<i32>(" <> elements <> ")"
        VectorType 2 TFloat -> "vec2<f32>(" <> elements <> ")"
        VectorType 2 TUInt -> "vec2<u32>(" <> elements <> ")"
        VectorType 3 TInt -> "vec3<i32>(" <> elements <> ")"
        VectorType 3 TFloat -> "vec3<f32>(" <> elements <> ")"
        VectorType 3 TUInt -> "vec3<u32>(" <> elements <> ")"
        VectorType 4 TInt -> "vec4<i32>(" <> elements <> ")"
        VectorType 4 TFloat -> "vec4<f32>(" <> elements <> ")"
        VectorType 4 TUInt -> "vec4<u32>(" <> elements <> ")"
        _ -> undefined
exprToWGSL (ConstantMat @n @m (_ :: Proxy a) x) =
  let elements = printMatElements @n @m x
   in case demote @a of
        MatrixType 2 2 -> "mat2x2<f32>(" <> elements <> ")"
        MatrixType 3 2 -> "mat2x3<f32>(" <> elements <> ")"
        MatrixType 4 2 -> "mat2x4<f32>(" <> elements <> ")"
        MatrixType 2 3 -> "mat3x2<f32>(" <> elements <> ")"
        MatrixType 3 3 -> "mat3x3<f32>(" <> elements <> ")"
        MatrixType 4 3 -> "mat3x4<f32>(" <> elements <> ")"
        MatrixType 2 4 -> "mat4x2<f32>(" <> elements <> ")"
        MatrixType 3 4 -> "mat4x3<f32>(" <> elements <> ")"
        MatrixType 4 4 -> "mat4x4<f32>(" <> elements <> ")"
        _ -> undefined
exprToWGSL (ConstantInt x) = T.pack $ show x <> "i"
exprToWGSL (ConstantFloat x) = T.pack $ show x <> "f"
exprToWGSL (ConstantUInt x) = T.pack $ show x <> "u"
exprToWGSL (Var x) = nameIndex x
exprToWGSL (BindingVar x) = "b" <> T.pack (show x)
exprToWGSL Void = ""
exprToWGSL (AccessField x f) = exprToWGSL x <> "." <> f
exprToWGSL (Index a i) = exprToWGSL a <> "[" <> exprToWGSL i <> "]"
exprToWGSL (BuiltInVar x) = "input." <> x
exprToWGSL (LocationVar x) = "input.l" <> T.pack (show x)
exprToWGSL (StructInit "" [p]) = p
exprToWGSL (StructInit name params) = name <> "(" <> T.intercalate "," params <> ")"
exprToWGSL (CustomOperator op a b) = "(" <> exprToWGSL a <> " " <> op <> " " <> exprToWGSL b <> ")"
exprToWGSL (VarCustom t) = t
exprToWGSL (AppendVecs (_ :: Proxy v) a b) = case demote @v of
  VectorType 2 TInt -> "vec2<i32>(" <> exprToWGSL a <> "," <> exprToWGSL b <> ")"
  VectorType 2 TFloat -> "vec2<f32>(" <> exprToWGSL a <> "," <> exprToWGSL b <> ")"
  VectorType 2 TUInt -> "vec2<u32>(" <> exprToWGSL a <> "," <> exprToWGSL b <> ")"
  VectorType 3 TInt -> "vec3<i32>(" <> exprToWGSL a <> "," <> exprToWGSL b <> ")"
  VectorType 3 TFloat -> "vec3<f32>(" <> exprToWGSL a <> "," <> exprToWGSL b <> ")"
  VectorType 3 TUInt -> "vec3<u32>(" <> exprToWGSL a <> "," <> exprToWGSL b <> ")"
  VectorType 4 TInt -> "vec4<i32>(" <> exprToWGSL a <> "," <> exprToWGSL b <> ")"
  VectorType 4 TFloat -> "vec4<f32>(" <> exprToWGSL a <> "," <> exprToWGSL b <> ")"
  VectorType 4 TUInt -> "vec4<u32>(" <> exprToWGSL a <> "," <> exprToWGSL b <> ")"
  _ -> undefined
exprToWGSL (ExtendVec (_ :: Proxy v) a b) = case demote @v of
  VectorType 2 TInt -> "vec2<i32>(" <> exprToWGSL a <> "," <> exprToWGSL b <> ")"
  VectorType 2 TFloat -> "vec2<f32>(" <> exprToWGSL a <> "," <> exprToWGSL b <> ")"
  VectorType 2 TUInt -> "vec2<u32>(" <> exprToWGSL a <> "," <> exprToWGSL b <> ")"
  VectorType 3 TInt -> "vec3<i32>(" <> exprToWGSL a <> "," <> exprToWGSL b <> ")"
  VectorType 3 TFloat -> "vec3<f32>(" <> exprToWGSL a <> "," <> exprToWGSL b <> ")"
  VectorType 3 TUInt -> "vec3<u32>(" <> exprToWGSL a <> "," <> exprToWGSL b <> ")"
  VectorType 4 TInt -> "vec4<i32>(" <> exprToWGSL a <> "," <> exprToWGSL b <> ")"
  VectorType 4 TFloat -> "vec4<f32>(" <> exprToWGSL a <> "," <> exprToWGSL b <> ")"
  VectorType 4 TUInt -> "vec4<u32>(" <> exprToWGSL a <> "," <> exprToWGSL b <> ")"
  _ -> undefined
exprToWGSL (ConstantBool b) = if b then "true" else "false"
exprToWGSL (If bool a b) = "if " <> exprToWGSL bool <> " {" <> exprToWGSL a <> "} else {" <> exprToWGSL b <> "}"

stmtToWGSL :: Stmt -> Text
stmtToWGSL Empty = ""
stmtToWGSL (Concat Empty a) = stmtToWGSL a
stmtToWGSL (Concat a b) = stmtToWGSL a <> "\n" <> stmtToWGSL b
stmtToWGSL (Let a b) = "  let " <> nameIndex a <> " = " <> exprToWGSL b <> ";"
stmtToWGSL (Return a) = "  return " <> exprToWGSL a <> ";"
stmtToWGSL (Debug t) = "  //" <> t

-- stmtToWGSL (VLet a b) = "  var " <> nameIndex a <> " = " <> exprToWGSL b <> ";"

newtype Shader m a = Shader (State ShaderState a) deriving newtype (Functor, Applicative, Monad, MonadState ShaderState)

typeToWGSL :: forall (a :: Types). (SingI a) => Expr a -> Text
typeToWGSL _ = typeToWGSL' (demote @a)

typeToWGSL' :: Types -> Text
typeToWGSL' a = case a of
  Primitive TInt -> "i32"
  Primitive TFloat -> "f32"
  Primitive TUInt -> "u32"
  -- ArrayType n t -> "array<" <> typeToWGSL' t <> ", " <> T.pack (show n) <> ">"
  VectorType 2 TFloat -> "vec2<f32>"
  VectorType 2 TInt -> "vec2<i32>"
  VectorType 2 TUInt -> "vec2<u32>"
  VectorType 3 TFloat -> "vec3<f32>"
  VectorType 3 TInt -> "vec3<i32>"
  VectorType 3 TUInt -> "vec3<u32>"
  VectorType 4 TFloat -> "vec4<f32>"
  VectorType 4 TInt -> "vec4<i32>"
  VectorType 4 TUInt -> "vec4<u32>"
  MatrixType 2 2 -> "mat2x2<f32>"
  MatrixType 2 3 -> "mat2x3<f32>"
  MatrixType 2 4 -> "mat2x4<f32>"
  MatrixType 3 2 -> "mat3x2<f32>"
  MatrixType 3 3 -> "mat3x3<f32>"
  MatrixType 3 4 -> "mat3x4<f32>"
  MatrixType 4 2 -> "mat4x2<f32>"
  MatrixType 4 3 -> "mat4x3<f32>"
  MatrixType 4 4 -> "mat4x4<f32>"
  _ -> undefined

addStmt :: Stmt -> Shader m ()
addStmt stmt = do
  state <- get
  put state {ast = Concat state.ast stmt}

newVarIndex :: Shader m VarIndex
newVarIndex = do
  ShaderState {counter, ast, dict} <- get
  put (ShaderState {counter = counter + 1, ast, dict})
  pure counter

var :: (SingI a) => Expr a -> Shader m (Expr a)
var expr = do
  index <- newVarIndex
  addStmt (Let index expr)
  pure $ Var index

comment :: Text -> Shader m ()
comment t = addStmt (Debug t)

data CPU

data GPU

data Internal