Functions.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_GHC -Wno-redundant-constraints #-}

module Mischief.Shader.Functions where

import Data.Data
import Data.Singletons (SingI)
import Data.Text (Text)
import GHC.TypeLits hiding (Div)
import Mischief.Render.Core
import Mischief.Render.Texture
import Mischief.Shader.Internal.Bindings
import Mischief.Shader.Internal.Singletons
import Mischief.Shader.Internal.State
import Mischief.Shader.Types
import Prelude (Bool (..), undefined, ($), type (~))

sample :: Expr TTexture -> Expr TSampler -> Vec 2 Float -> Vec 4 Float
sample tex sampler v =
  Function
    "textureSample"
    [ Param tex,
      Param sampler,
      Param v
    ]

unsafeCall :: (SingI a) => Text -> [Param] -> Expr a
unsafeCall = Function

-- Binary
(&) :: (Algebric a ~ True, Integral a ~ True) => Expr a -> Expr a -> Expr a
(&) = CustomOperator "&"

-- Arithmetics

class Num a where
  (+) :: Expr a -> Expr a -> Expr a
  (+) = CustomOperator "+"

  (-) :: Expr a -> Expr a -> Expr a
  (-) = CustomOperator "-"

  (*) :: Expr a -> Expr a -> Expr a
  (*) = CustomOperator "*"

instance Num (Primitive a)

instance Num (VectorType n a)

instance Num (MatrixType n m)

(^+) :: Expr (VectorType n a) -> Expr (Primitive a) -> Expr (VectorType n a)
(^+) = CustomOperator "+"

(+^) :: Expr (Primitive a) -> Expr (VectorType n a) -> Expr (VectorType n a)
(+^) = CustomOperator "+"

(^-) :: Expr (VectorType n a) -> Expr (Primitive a) -> Expr (VectorType n a)
(^-) = CustomOperator "-"

(-^) :: Expr (Primitive a) -> Expr (VectorType n a) -> Expr (VectorType n a)
(-^) = CustomOperator "-"

(^*) :: Expr (VectorType n a) -> Expr (Primitive a) -> Expr (VectorType n a)
(^*) = CustomOperator "*"

(*^) :: Expr (Primitive a) -> Expr (VectorType n a) -> Expr (VectorType n a)
(*^) = CustomOperator "*"

(*!) :: Float -> Mat n m -> Mat n m
(*!) = CustomOperator "*"

(!*) :: Mat n m -> Float -> Mat n m
(!*) = CustomOperator "*"

(^*!) :: Vec m Float -> Mat n m -> Vec n Float
(^*!) = CustomOperator "*"

(!*^) :: Mat n m -> Vec n Float -> Vec m Float
(!*^) = CustomOperator "*"

(!*!) :: Mat k r -> Mat c k -> Mat c r
(!*!) = CustomOperator "*"

class Div a where
  (/) :: Expr a -> Expr a -> Expr a
  (/) = CustomOperator "/"

instance {-# OVERLAPPABLE #-} (Algebric a ~ True) => Div a

(^/) :: Expr (VectorType n a) -> Expr (Primitive a) -> Expr (VectorType n a)
(^/) = CustomOperator "/"

(/^) :: Expr (Primitive a) -> Expr (VectorType n a) -> Expr (VectorType n a)
(/^) = CustomOperator "/"

class Modulo a where
  (%) :: Expr a -> Expr a -> Expr a
  (%) = CustomOperator "%"

instance {-# OVERLAPPABLE #-} (Integral a ~ True) => Modulo a

(^%) :: (Integral (Primitive a) ~ True) => Expr (VectorType n a) -> Expr (Primitive a) -> Expr (VectorType n a)
(^%) = CustomOperator "%"

(%^) :: (Integral (Primitive a) ~ True) => Expr (Primitive a) -> Expr (VectorType n a) -> Expr (VectorType n a)
(%^) = CustomOperator "%"

-- Numeric Functions

abs :: (SingI a, Algebric a ~ True) => Expr a -> Expr a
abs = Abs

acos :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
acos x = Function "acos" [Param x]

acosh :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
acosh x = Function "acosh" [Param x]

asin :: (SingI a, Algebric a ~ True) => Expr a -> Expr a
asin x = Function "asin" [Param x]

asinh :: (SingI a, Algebric a ~ True) => Expr a -> Expr a
asinh x = Function "asinh" [Param x]

atan :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
atan x = Function "atan" [Param x]

atanh :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
atanh x = Function "atanh" [Param x]

atan2 :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a -> Expr a
atan2 x y = Function "atan2" [Param x, Param y]

ceil :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
ceil x = Function "ceil" [Param x]

clamp :: (SingI a, Algebric a ~ True) => Expr a -> Expr a -> Expr a -> Expr a
clamp e low high = Function "clamp" [Param e, Param low, Param high]

cos :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
cos x = Function "cos" [Param x]

cosh :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
cosh x = Function "cosh" [Param x]

countLeadingZeros :: (SingI a, Integral a ~ True) => Expr a -> Expr a
countLeadingZeros x = Function "countLeadingZeros" [Param x]

countOneBits :: (SingI a, Integral a ~ True) => Expr a -> Expr a
countOneBits x = Function "countOneBits" [Param x]

countTrailingZeros :: (SingI a, Integral a ~ True) => Expr a -> Expr a
countTrailingZeros x = Function "countTrailingZeros" [Param x]

cross :: (SingI (VectorType 3 a), Fractional (Expr (Primitive a)) ~ True) => Expr (VectorType 3 a) -> Expr (VectorType 3 a) -> Expr (VectorType 3 a)
cross a b = Function "cross" [Param a, Param b]

degrees :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
degrees x = Function "degrees" [Param x]

determinant :: (SingI n) => Expr (MatrixType n n) -> Float
determinant x = Function "determinant" [Param x]

distance :: (SingI a, Fractional (Expr a) ~ True, VectorOrSingleOf a ~ b) => Expr a -> Expr a -> Expr (Primitive b)
distance a b = Function "distance" [Param a, Param b]

dot :: (SingI (Expr (VectorType n a)), SingI a, SingI n, Algebric (Primitive a) ~ True) => Expr (VectorType n a) -> Expr (VectorType n a) -> Expr (Primitive a)
dot a b = Function "dot" [Param a, Param b]

dot4U8Packed :: Uint -> Uint -> Uint
dot4U8Packed a b = Function "dot4U8Packed" [Param a, Param b]

dot4I8Packed :: Uint -> Uint -> Int
dot4I8Packed a b = Function "dot4I8Packed" [Param a, Param b]

exp :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
exp a = Function "exp" [Param a]

exp2 :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
exp2 a = Function "exp2" [Param a]

extractBits :: (SingI a, Integral a ~ True) => Expr a -> Uint -> Uint -> Expr a
extractBits e offset count = Function "extractBits" [Param e, Param offset, Param count]

faceForward :: (SingI a, SingI n, Fractional (Expr (Primitive a)) ~ True) => Expr (VectorType n a) -> Expr (VectorType n a) -> Expr (VectorType n a) -> Expr (VectorType n a)
faceForward a b c = Function "faceForward" [Param a, Param b, Param c]

firstLeadingBit :: (SingI a, SingI n, Integral a ~ True) => Expr a -> Expr a
firstLeadingBit a = Function "firstLeadingBit" [Param a]

firstTrailingBit :: (SingI a, SingI n, Integral a ~ True) => Expr a -> Expr a
firstTrailingBit a = Function "firstTrailingBit" [Param a]

floor :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
floor x = Function "floor" [Param x]

fma :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a -> Expr a -> Expr a
fma a b c = Function "fma" [Param a, Param b, Param c]

fract :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
fract a = Function "fract" [Param a]

-- TODO: frexp (custom struct)

insertBits :: (SingI a, Integral a ~ True) => Expr a -> Expr a -> Uint -> Uint -> Expr a
insertBits e newbits offset count = Function "insertBits" [Param e, Param newbits, Param offset, Param count]

inverseSqrt :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
inverseSqrt a = Function "inverseSqrt" [Param a]

-- TODO: ldexp

length :: (SingI a, SingI b, SingI n, VectorOrSingleOf b ~ a) => Expr b -> Expr (Primitive a)
length a = Function "length" [Param a]

log :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
log a = Function "log" [Param a]

log2 :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
log2 a = Function "log2" [Param a]

max :: (SingI a, Algebric a ~ True) => Expr a -> Expr a -> Expr a
max a b = Function "max" [Param a, Param b]

min :: (SingI a, Algebric a ~ True) => Expr a -> Expr a -> Expr a
min a b = Function "min" [Param a, Param b]

class Mix a b

instance Mix a a

instance Mix (VectorType n a) (Primitive a)

mix :: (SingI a, SingI b, Fractional (Expr a) ~ True, Fractional (Expr b) ~ True, Mix a b) => Expr a -> Expr a -> Expr b -> Expr a
mix a b c = Function "mix" [Param a, Param b, Param c]

-- TODO: modf (custom struct)

normalize :: (SingI n, SingI a, Fractional (Expr (VectorType n a)) ~ True) => Expr (VectorType n a) -> Expr (VectorType n a)
normalize a = Function "normalize" [Param a]

pow :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a -> Expr a
pow a b = Function "pow" [Param a, Param b]

quantizeToF16 :: (SingI a, VectorOrSingleOf a ~ TFloat) => Expr a -> Expr a
quantizeToF16 a = Function "quantizeToF16" [Param a]

radians :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
radians a = Function "radians" [Param a]

reflect :: (SingI n, SingI a, Fractional (Expr (VectorType n a)) ~ True) => Expr (VectorType n a) -> Expr (VectorType n a) -> Expr (VectorType n a)
reflect a b = Function "reflect" [Param a, Param b]

refract :: (SingI a, SingI b, VectorOrSingleOf a ~ b) => Expr a -> Expr a -> Expr (Primitive b) -> Expr a
refract a b c = Function "refract" [Param a, Param b, Param c]

reverseBits :: (SingI a, Integral a ~ True) => Expr a -> Expr a
reverseBits a = Function "reverseBits" [Param a]

round :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
round a = Function "round" [Param a]

saturate :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
saturate a = Function "saturate" [Param a]

sign :: (SingI a, Algebric a ~ True) => Expr a -> Expr a
sign a = Function "sign" [Param a]

sin :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
sin a = Function "sin" [Param a]

sinh :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
sinh a = Function "sinh" [Param a]

smoothstep :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a -> Expr a -> Expr a
smoothstep a b c = Function "smoothstep" [Param a, Param b, Param c]

sqrt :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
sqrt a = Function "sqrt" [Param a]

step :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a -> Expr a
step a b = Function "step" [Param a, Param b]

tan :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
tan a = Function "tan" [Param a]

tanh :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
tanh a = Function "tanh" [Param a]

transpose :: (SingI n, SingI m) => Expr (MatrixType n m) -> Expr (MatrixType m n)
transpose a = Function "transpose" [Param a]

trunc :: (SingI a, Fractional (Expr a) ~ True) => Expr a -> Expr a
trunc a = Function "trunc" [Param a]

-- TODO: Bit Reinterpretation Functions

-- TODO: Logical Functions

-- TODO: Array Functions

-- TODO: Derivative Functions

-- TODO: Texture Functions

-- TODO: Atomic Functions

-- TODO: Data Packing Functions

-- TODO: Data Unpacking Functions

-- TODO: Synchronization Functions