Types.hs

{-# 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 = ConstantFloat

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

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

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

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

u32 :: Prelude.Int -> Uint
u32 = 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 a b = ConstantVec Proxy (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 a b c = ConstantVec Proxy (a, b, 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 a b c d = ConstantVec Proxy (a, b, c, d)

(+:) :: (SingI a, PrintVecElements (n + 1) a, KnownNat (n + 1)) => Expr (VectorType n a) -> Expr (Primitive a) -> Expr (VectorType (n + 1) a)
(+:) = ExtendVec Proxy

(<>) :: (SingI a, PrintVecElements (n + m) a, KnownNat (n + m)) => Expr (VectorType n a) -> Expr (VectorType m a) -> Expr (VectorType (n + m) a)
(<>) = AppendVecs Proxy

array :: forall a. (SingI (DesugarExpr a)) => [Expr (DesugarExpr a)] -> Array a
array = ConstantArray (Proxy @(DesugarExpr a))

(!!) :: (SingI a) => Array (Expr 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 = ConstantMat (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 = Cast (Proxy @(Primitive a))

instance (SingI (VectorType n a), SingI (VectorType n b)) => Cast (VectorType n a) (VectorType n b) where
  cast = CastVec (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 a = case eqT @2 @n of
  Just Refl -> ConstantVec (Proxy @(VectorType 2 a)) (a, a)
  Nothing -> case eqT @3 @n of
    Just Refl -> ConstantVec (Proxy @(VectorType 3 a)) (a, a, a)
    Nothing -> case eqT @4 @n of
      Just Refl -> ConstantVec (Proxy @(VectorType 4 a)) (a, a, a, a)
      Nothing -> 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 v = AccessField v (T.pack $ map m $ symbolVal $ Proxy @s)
    where
      m 'r' = 'x'
      m 'g' = 'y'
      m 'b' = 'z'
      m 'a' = 'w'
      m a = a