Variadic.hs

{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE PolyKinds              #-}
{-# LANGUAGE TypeFamilyDependencies #-}
{-# LANGUAGE TypeOperators          #-}
{-# LANGUAGE UndecidableInstances   #-}

{-|
Module: Data.Function.Variadic

Type families for variadic functions.
-}

module Data.Function.Variadic where

-- base
import Data.Int
  ( Int8,Int16,Int32,Int64 )
import Data.Kind
  ( Type )
import Data.Word
  ( Word8,Word16,Word32,Word64 )
import GHC.TypeNats
  ( Nat, type (-), type (<=?) )

-- half
import Numeric.Half
  ( Half )

-- fir
import {-# SOURCE #-} Data.Product
  ( HList )
import {-# SOURCE #-} FIR.AST
  ( AST )
import FIR.AST.Type
  ( AugType(Val, (:-->)) )
import FIR.Prim.Array
  ( Array,RuntimeArray )
import FIR.Prim.RayTracing
  ( AccelerationStructure, RayQuery )
import {-# SOURCE #-} FIR.Prim.Struct
  ( Struct )
import Math.Linear
  ( V,M )

------------------------------------------------------------

type NatVariadic
      ( n :: Nat  )
      ( a :: Type )
      ( b :: Type )
  = ( NatVariadic' n a b (1 <=? n) :: AugType )

type family NatVariadic'
              ( n    :: Nat  )
              ( a    :: Type )
              ( b    :: Type )
              ( geq1 :: Bool )
            = ( r    :: AugType )
            | r -> b
            where
  NatVariadic' n _ b 'False = Val b
  NatVariadic' n a b 'True  = Val a :--> NatVariadic (n-1) a b

------------------------------------------------------------

type family AugListVariadic (as :: [AugType]) (b :: Type) = ( r :: AugType ) | r -> as b where
  AugListVariadic '[]       b = Val b
  AugListVariadic (a ': as) b = a :--> AugListVariadic as b

-- explicitly list all primitive types, to allow injectivity annotation
type family ListVariadic (as :: [Type]) (b :: Type) = (r :: Type) | r -> as b where  
  ListVariadic (a ': as) b = a -> ListVariadic as b
  ListVariadic '[] ()     = ()
  ListVariadic '[] Bool   = Bool
  ListVariadic '[] Word   = Word 
  ListVariadic '[] Word8  = Word8
  ListVariadic '[] Word16 = Word16
  ListVariadic '[] Word32 = Word32
  ListVariadic '[] Word64 = Word64
  ListVariadic '[] Int    = Int 
  ListVariadic '[] Int8   = Int8
  ListVariadic '[] Int16  = Int16
  ListVariadic '[] Int32  = Int32
  ListVariadic '[] Int64  = Int64
  ListVariadic '[] Half   = Half
  ListVariadic '[] Float  = Float
  ListVariadic '[] Double = Double
  ListVariadic '[] (V n a) = V n a
  ListVariadic '[] (M m n a) = M m n a
  ListVariadic '[] (Struct as) = Struct as
  ListVariadic '[] (Array n a) = Array n a
  ListVariadic '[] (RuntimeArray a) = RuntimeArray a
  ListVariadic '[] (AST a) = AST a
  ListVariadic '[] (a1,a2) = (a1,a2)
  ListVariadic '[] (a1,a2,a3) = (a1,a2,a3)
  ListVariadic '[] (a1,a2,a3,a4) = (a1,a2,a3,a4)
  ListVariadic '[] (a1,a2,a3,a4,a5) = (a1,a2,a3,a4,a5)
  ListVariadic '[] (a1,a2,a3,a4,a5,a6) = (a1,a2,a3,a4,a5,a6)
  ListVariadic '[] (a1,a2,a3,a4,a5,a6,a7) = (a1,a2,a3,a4,a5,a6,a7)
  ListVariadic '[] (a1,a2,a3,a4,a5,a6,a7,a8) = (a1,a2,a3,a4,a5,a6,a7,a8)
  ListVariadic '[] (HList as) = HList as
  ListVariadic '[] AccelerationStructure = AccelerationStructure
  ListVariadic '[] RayQuery = RayQuery