{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_GHC -Wno-orphans #-}
{-# OPTIONS_GHC -Wno-redundant-constraints #-}

module Mischief.Shader.Types where

import Data.Data
import Data.List
import Data.Singletons (SingI)
import Data.Text qualified as T
import Data.Type.Ord (type (>=))
import GHC.Records (HasField (getField))
import GHC.TypeLits
import Mischief.Shader.Internal.Singletons
import Mischief.Shader.Internal.State
import Prelude (Bool (..), Char, Maybe (Just, Nothing), Show (show), String, undefined, ($), type (~))
import Prelude qualified

float :: Prelude.Float -> Float
float :: Float -> Float
float = Float -> Float
ConstantFloat

int :: Prelude.Int -> Int
int :: Int -> Int
int = Int -> Int
ConstantInt

uint :: Prelude.Int -> Uint
uint :: Int -> Uint
uint = Int -> Uint
ConstantUInt

f32 :: Prelude.Float -> Float
f32 :: Float -> Float
f32 = Float -> Float
float

i32 :: Prelude.Int -> Int
i32 :: Int -> Int
i32 = Int -> Int
int

u32 :: Prelude.Int -> Uint
u32 :: Int -> Uint
u32 = Int -> Uint
uint

vec2 :: forall a. (Expr (Primitive (DesugarPrimitiveExpr a)) ~ a, SingI (DesugarPrimitiveExpr a), PrintVecElements 2 (DesugarPrimitiveExpr a)) => Expr (Primitive (DesugarPrimitiveExpr a)) -> Expr (Primitive (DesugarPrimitiveExpr a)) -> Vec 2 a
vec2 :: forall a.
(Expr ('Primitive (DesugarPrimitiveExpr a)) ~ a,
 SingI (DesugarPrimitiveExpr a),
 PrintVecElements 2 (DesugarPrimitiveExpr a)) =>
Expr ('Primitive (DesugarPrimitiveExpr a))
-> Expr ('Primitive (DesugarPrimitiveExpr a)) -> Vec 2 a
vec2 Expr ('Primitive (DesugarPrimitiveExpr a))
a Expr ('Primitive (DesugarPrimitiveExpr a))
b = Proxy ('VectorType 2 (DesugarPrimitiveExpr a))
-> VecInit 2 (DesugarPrimitiveExpr a)
-> Expr ('VectorType 2 (DesugarPrimitiveExpr a))
forall (n :: Natural) (a' :: PrimitiveTypes).
(PrintVecElements n a', SingI ('VectorType n a')) =>
Proxy ('VectorType n a') -> VecInit n a' -> Expr ('VectorType n a')
ConstantVec Proxy ('VectorType 2 (DesugarPrimitiveExpr a))
forall {k} (t :: k). Proxy t
Proxy (Expr ('Primitive (DesugarPrimitiveExpr a))
a, Expr ('Primitive (DesugarPrimitiveExpr a))
b)

vec3 :: forall a. (Expr (Primitive (DesugarPrimitiveExpr a)) ~ a, SingI (DesugarPrimitiveExpr a), PrintVecElements 3 (DesugarPrimitiveExpr a)) => Expr (Primitive (DesugarPrimitiveExpr a)) -> Expr (Primitive (DesugarPrimitiveExpr a)) -> Expr (Primitive (DesugarPrimitiveExpr a)) -> Vec 3 a
vec3 :: forall a.
(Expr ('Primitive (DesugarPrimitiveExpr a)) ~ a,
 SingI (DesugarPrimitiveExpr a),
 PrintVecElements 3 (DesugarPrimitiveExpr a)) =>
Expr ('Primitive (DesugarPrimitiveExpr a))
-> Expr ('Primitive (DesugarPrimitiveExpr a))
-> Expr ('Primitive (DesugarPrimitiveExpr a))
-> Vec 3 a
vec3 Expr ('Primitive (DesugarPrimitiveExpr a))
a Expr ('Primitive (DesugarPrimitiveExpr a))
b Expr ('Primitive (DesugarPrimitiveExpr a))
c = Proxy ('VectorType 3 (DesugarPrimitiveExpr a))
-> VecInit 3 (DesugarPrimitiveExpr a)
-> Expr ('VectorType 3 (DesugarPrimitiveExpr a))
forall (n :: Natural) (a' :: PrimitiveTypes).
(PrintVecElements n a', SingI ('VectorType n a')) =>
Proxy ('VectorType n a') -> VecInit n a' -> Expr ('VectorType n a')
ConstantVec Proxy ('VectorType 3 (DesugarPrimitiveExpr a))
forall {k} (t :: k). Proxy t
Proxy (Expr ('Primitive (DesugarPrimitiveExpr a))
a, Expr ('Primitive (DesugarPrimitiveExpr a))
b, Expr ('Primitive (DesugarPrimitiveExpr a))
c)

vec4 :: forall a. (Expr (Primitive (DesugarPrimitiveExpr a)) ~ a, SingI (DesugarPrimitiveExpr a), PrintVecElements 4 (DesugarPrimitiveExpr a)) => Expr (Primitive (DesugarPrimitiveExpr a)) -> Expr (Primitive (DesugarPrimitiveExpr a)) -> Expr (Primitive (DesugarPrimitiveExpr a)) -> Expr (Primitive (DesugarPrimitiveExpr a)) -> Vec 4 a
vec4 :: forall a.
(Expr ('Primitive (DesugarPrimitiveExpr a)) ~ a,
 SingI (DesugarPrimitiveExpr a),
 PrintVecElements 4 (DesugarPrimitiveExpr a)) =>
Expr ('Primitive (DesugarPrimitiveExpr a))
-> Expr ('Primitive (DesugarPrimitiveExpr a))
-> Expr ('Primitive (DesugarPrimitiveExpr a))
-> Expr ('Primitive (DesugarPrimitiveExpr a))
-> Vec 4 a
vec4 Expr ('Primitive (DesugarPrimitiveExpr a))
a Expr ('Primitive (DesugarPrimitiveExpr a))
b Expr ('Primitive (DesugarPrimitiveExpr a))
c Expr ('Primitive (DesugarPrimitiveExpr a))
d = Proxy ('VectorType 4 (DesugarPrimitiveExpr a))
-> VecInit 4 (DesugarPrimitiveExpr a)
-> Expr ('VectorType 4 (DesugarPrimitiveExpr a))
forall (n :: Natural) (a' :: PrimitiveTypes).
(PrintVecElements n a', SingI ('VectorType n a')) =>
Proxy ('VectorType n a') -> VecInit n a' -> Expr ('VectorType n a')
ConstantVec Proxy ('VectorType 4 (DesugarPrimitiveExpr a))
forall {k} (t :: k). Proxy t
Proxy (Expr ('Primitive (DesugarPrimitiveExpr a))
a, Expr ('Primitive (DesugarPrimitiveExpr a))
b, Expr ('Primitive (DesugarPrimitiveExpr a))
c, Expr ('Primitive (DesugarPrimitiveExpr a))
d)

(+:) :: (SingI a, PrintVecElements (n + 1) a, KnownNat (n + 1)) => Expr (VectorType n a) -> Expr (Primitive a) -> Expr (VectorType (n + 1) a)
+: :: forall (a :: PrimitiveTypes) (n :: Natural).
(SingI a, PrintVecElements (n + 1) a, KnownNat (n + 1)) =>
Expr ('VectorType n a)
-> Expr ('Primitive a) -> Expr ('VectorType (n + 1) a)
(+:) = Proxy ('VectorType (n + 1) a)
-> Expr ('VectorType n a)
-> Expr ('Primitive a)
-> Expr ('VectorType (n + 1) a)
forall (n :: Natural) (a1 :: PrimitiveTypes).
(SingI a1, PrintVecElements (n + 1) a1,
 SingI ('VectorType (n + 1) a1)) =>
Proxy ('VectorType (n + 1) a1)
-> Expr ('VectorType n a1)
-> Expr ('Primitive a1)
-> Expr ('VectorType (n + 1) a1)
ExtendVec Proxy ('VectorType (n + 1) a)
forall {k} (t :: k). Proxy t
Proxy

(<>) :: (SingI a, PrintVecElements (n + m) a, KnownNat (n + m)) => Expr (VectorType n a) -> Expr (VectorType m a) -> Expr (VectorType (n + m) a)
<> :: forall (a :: PrimitiveTypes) (n :: Natural) (m :: Natural).
(SingI a, PrintVecElements (n + m) a, KnownNat (n + m)) =>
Expr ('VectorType n a)
-> Expr ('VectorType m a) -> Expr ('VectorType (n + m) a)
(<>) = Proxy ('VectorType (n + m) a)
-> Expr ('VectorType n a)
-> Expr ('VectorType m a)
-> Expr ('VectorType (n + m) a)
forall (n :: Natural) (m :: Natural) (a1 :: PrimitiveTypes).
(SingI a1, PrintVecElements (n + m) a1,
 SingI ('VectorType (n + m) a1)) =>
Proxy ('VectorType (n + m) a1)
-> Expr ('VectorType n a1)
-> Expr ('VectorType m a1)
-> Expr ('VectorType (n + m) a1)
AppendVecs Proxy ('VectorType (n + m) a)
forall {k} (t :: k). Proxy t
Proxy

array :: forall a. (SingI (DesugarExpr a)) => [Expr (DesugarExpr a)] -> Array a
array :: forall a.
SingI (DesugarExpr a) =>
[Expr (DesugarExpr a)] -> Array a
array = Proxy (DesugarExpr a)
-> [Expr (DesugarExpr a)] -> Expr ('ArrayType (DesugarExpr a))
forall (a' :: Types).
SingI a' =>
Proxy a' -> [Expr a'] -> Expr ('ArrayType a')
ConstantArray (forall {k} (t :: k). Proxy t
forall (t :: Types). Proxy t
Proxy @(DesugarExpr a))

(!!) :: (SingI a) => Array (Expr a) -> Uint -> Expr a
!! :: forall (a :: Types). SingI a => Array (Expr a) -> Uint -> Expr a
(!!) = Expr ('ArrayType a) -> Uint -> Expr a
Expr ('ArrayType (DesugarExpr (Expr a))) -> Uint -> Expr a
forall (a :: Types).
SingI ('ArrayType a) =>
Expr ('ArrayType a) -> Uint -> Expr a
Index

type Boolean = Expr TBool

type Int = Expr (Primitive TInt)

type Float = Expr (Primitive TFloat)

type Uint = Expr (Primitive TUInt)

type family DesugarExpr a where
  DesugarExpr (Expr a) = a

type family DesugarPrimitiveExpr a where
  DesugarPrimitiveExpr Int = TInt
  DesugarPrimitiveExpr Float = TFloat
  DesugarPrimitiveExpr Uint = TUInt

type Array b = Expr (ArrayType (DesugarExpr b))

type Vec n a = Expr (VectorType n (DesugarPrimitiveExpr a))

type Vec2 a = Vec 2 a

type Vec3 a = Vec 3 a

type Vec4 a = Vec 4 a

type Mat n m = Expr (MatrixType n m)

mat :: forall n m. (KnownNat n, KnownNat m, PrintMatElements n m) => MatInit n m -> Mat n m
mat :: forall (n :: Natural) (m :: Natural).
(KnownNat n, KnownNat m, PrintMatElements n m) =>
MatInit n m -> Mat n m
mat = Proxy ('MatrixType n m) -> MatInit n m -> Expr ('MatrixType n m)
forall (n :: Natural) (m :: Natural).
(PrintMatElements n m, SingI ('MatrixType n m)) =>
Proxy ('MatrixType n m) -> MatInit n m -> Expr ('MatrixType n m)
ConstantMat (forall {k} (t :: k). Proxy t
forall (t :: Types). Proxy t
Proxy @(MatrixType n m))

class (SingI b, SingI a) => Cast b a where
  cast :: Expr a -> Expr b

instance (SingI (Primitive a), SingI (Primitive b)) => Cast (Primitive a) (Primitive b) where
  cast :: Expr ('Primitive b) -> Expr ('Primitive a)
cast = Proxy ('Primitive a) -> Expr ('Primitive b) -> Expr ('Primitive a)
forall (a1 :: Types) (a :: Types).
(SingI a1, SingI a) =>
Proxy a -> Expr a1 -> Expr a
Cast (forall {k} (t :: k). Proxy t
forall (t :: Types). Proxy t
Proxy @(Primitive a))

instance (SingI (VectorType n a), SingI (VectorType n b)) => Cast (VectorType n a) (VectorType n b) where
  cast :: Expr ('VectorType n b) -> Expr ('VectorType n a)
cast = Proxy ('VectorType n a)
-> Expr ('VectorType n b) -> Expr ('VectorType n a)
forall (n :: Natural) (a1 :: PrimitiveTypes) (b :: PrimitiveTypes).
(SingI ('VectorType n a1), SingI ('VectorType n b)) =>
Proxy ('VectorType n b)
-> Expr ('VectorType n a1) -> Expr ('VectorType n b)
CastVec (forall {k} (t :: k). Proxy t
forall (t :: Types). Proxy t
Proxy @(VectorType n a))

type family Fractional a where
  Fractional Float = True
  Fractional (Expr (VectorType n a)) = Fractional (Expr (Primitive a))
  Fractional a = False

type family Algebric (a :: Types) where
  Algebric (Primitive a) = True
  Algebric (VectorType n a) = True
  Algebric a = False

type family Integral (a :: Types) where
  Integral (Primitive TInt) = True
  Integral (Primitive TUInt) = True
  Integral (VectorType n TInt) = True
  Integral (VectorType n TUInt) = True
  Integral a = False

diag :: forall n a. (n <= 4, KnownNat n, SingI a) => Expr (Primitive a) -> Expr (VectorType n a)
diag :: forall (n :: Natural) (a :: PrimitiveTypes).
(n <= 4, KnownNat n, SingI a) =>
Expr ('Primitive a) -> Expr ('VectorType n a)
diag Expr ('Primitive a)
a = case forall (a :: Natural) (b :: Natural).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT @2 @n of
  Just 2 :~: n
Refl -> Proxy ('VectorType n a) -> VecInit n a -> Expr ('VectorType n a)
forall (n :: Natural) (a' :: PrimitiveTypes).
(PrintVecElements n a', SingI ('VectorType n a')) =>
Proxy ('VectorType n a') -> VecInit n a' -> Expr ('VectorType n a')
ConstantVec (forall {k} (t :: k). Proxy t
forall (t :: Types). Proxy t
Proxy @(VectorType 2 a)) (Expr ('Primitive a)
a, Expr ('Primitive a)
a)
  Maybe (2 :~: n)
Nothing -> case forall (a :: Natural) (b :: Natural).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT @3 @n of
    Just 3 :~: n
Refl -> Proxy ('VectorType n a) -> VecInit n a -> Expr ('VectorType n a)
forall (n :: Natural) (a' :: PrimitiveTypes).
(PrintVecElements n a', SingI ('VectorType n a')) =>
Proxy ('VectorType n a') -> VecInit n a' -> Expr ('VectorType n a')
ConstantVec (forall {k} (t :: k). Proxy t
forall (t :: Types). Proxy t
Proxy @(VectorType 3 a)) (Expr ('Primitive a)
a, Expr ('Primitive a)
a, Expr ('Primitive a)
a)
    Maybe (3 :~: n)
Nothing -> case forall (a :: Natural) (b :: Natural).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT @4 @n of
      Just 4 :~: n
Refl -> Proxy ('VectorType n a) -> VecInit n a -> Expr ('VectorType n a)
forall (n :: Natural) (a' :: PrimitiveTypes).
(PrintVecElements n a', SingI ('VectorType n a')) =>
Proxy ('VectorType n a') -> VecInit n a' -> Expr ('VectorType n a')
ConstantVec (forall {k} (t :: k). Proxy t
forall (t :: Types). Proxy t
Proxy @(VectorType 4 a)) (Expr ('Primitive a)
a, Expr ('Primitive a)
a, Expr ('Primitive a)
a, Expr ('Primitive a)
a)
      Maybe (4 :~: n)
Nothing -> Expr ('VectorType n a)
forall a. HasCallStack => a
undefined

type family SwizzleChar c n where
  SwizzleChar 'x' 1 = True
  SwizzleChar 'r' 1 = True
  SwizzleChar 'y' 2 = True
  SwizzleChar 'g' 2 = True
  SwizzleChar 'z' 3 = True
  SwizzleChar 'b' 3 = True
  SwizzleChar 'w' 4 = True
  SwizzleChar 'a' 4 = True
  SwizzleChar c 1 = False
  SwizzleChar c n = SwizzleChar c (n - 1)

type family Swizzle (s :: Symbol) (n :: Nat) :: Maybe Nat where
  Swizzle "" n = Just 0
  Swizzle s n = SwizzleRec (UnconsSymbol s) n

type family SwizzleRec (s :: Maybe (Char, Symbol)) (n :: Nat) :: Maybe Nat where
  SwizzleRec (Just '(c, s)) n = SwizzleRec' s n (SwizzleChar c n)
  SwizzleRec s n = Nothing

type family SwizzleRec' s n b where
  SwizzleRec' s n True = IncMaybe (Swizzle s n)
  SwizzleRec' s n b = Nothing

type family IncMaybe (n :: Maybe Nat) :: Maybe Nat where
  IncMaybe (Just n) = Just (n + 1)
  IncMaybe Nothing = Nothing

type family SwizzleOutput (a :: PrimitiveTypes) (n :: Nat) where
  SwizzleOutput a 1 = Primitive a
  SwizzleOutput a n = VectorType n a

instance (KnownSymbol s, KnownNat n, SingI a, Swizzle s n ~ Just m, m >= 1, m <= 4, SwizzleOutput a m ~ out) => HasField s (Expr (VectorType n a)) (Expr out) where
  getField :: Expr ('VectorType n a) -> Expr out
getField Expr ('VectorType n a)
v = Expr ('VectorType n a) -> Text -> Expr out
forall (a1 :: Types) (a :: Types).
SingI a1 =>
Expr a1 -> Text -> Expr a
AccessField Expr ('VectorType n a)
v (String -> Text
T.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ (Char -> Char) -> String -> String
forall a b. (a -> b) -> [a] -> [b]
map Char -> Char
m (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ Proxy s -> String
forall (n :: Symbol) (proxy :: Symbol -> *).
KnownSymbol n =>
proxy n -> String
symbolVal (Proxy s -> String) -> Proxy s -> String
forall a b. (a -> b) -> a -> b
$ forall {k} (t :: k). Proxy t
forall (t :: Symbol). Proxy t
Proxy @s)
    where
      m :: Char -> Char
m Char
'r' = Char
'x'
      m Char
'g' = Char
'y'
      m Char
'b' = Char
'z'
      m Char
'a' = Char
'w'
      m Char
a = Char
a