Types.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

module DataFrame.Internal.Types where

import Control.DeepSeq (NFData)
import Data.Int (Int16, Int32, Int64, Int8)
import Data.Kind (Constraint, Type)
import Data.Typeable (Typeable)
import qualified Data.Vector.Unboxed as VU
import Data.Word (Word16, Word32, Word64, Word8)

type Columnable' a = (Typeable a, Show a, Ord a, Eq a, Read a, NFData a)

{- | A type with column representations used to select the
"right" representation when specializing the `toColumn` function.
-}
data Rep
    = RBoxed
    | RUnboxed
    | ROptional

-- | Type-level if statement.
type family If (cond :: Bool) (yes :: k) (no :: k) :: k where
    If 'True yes _ = yes
    If 'False _ no = no

-- | All unboxable types (according to the `vector` package).
type family Unboxable (a :: Type) :: Bool where
    Unboxable Int = 'True
    Unboxable Int8 = 'True
    Unboxable Int16 = 'True
    Unboxable Int32 = 'True
    Unboxable Int64 = 'True
    Unboxable Word = 'True
    Unboxable Word8 = 'True
    Unboxable Word16 = 'True
    Unboxable Word32 = 'True
    Unboxable Word64 = 'True
    Unboxable Char = 'True
    Unboxable Bool = 'True
    Unboxable Double = 'True
    Unboxable Float = 'True
    Unboxable _ = 'False

type family Numeric (a :: Type) :: Bool where
    Numeric Integer = 'True
    Numeric Int = 'True
    Numeric Int8 = 'True
    Numeric Int16 = 'True
    Numeric Int32 = 'True
    Numeric Int64 = 'True
    Numeric Word = 'True
    Numeric Word8 = 'True
    Numeric Word16 = 'True
    Numeric Word32 = 'True
    Numeric Word64 = 'True
    Numeric Double = 'True
    Numeric Float = 'True
    Numeric _ = 'False

-- | Compute the column representation tag for any ‘a’.
type family KindOf a :: Rep where
    KindOf (Maybe a) = 'ROptional
    KindOf a = If (Unboxable a) 'RUnboxed 'RBoxed

-- | Type-level boolean for constraint/type comparison.
data SBool (b :: Bool) where
    STrue :: SBool 'True
    SFalse :: SBool 'False

-- | The runtime witness for our type-level branching.
class SBoolI (b :: Bool) where
    sbool :: SBool b

instance SBoolI 'True where sbool = STrue
instance SBoolI 'False where sbool = SFalse

-- | Type-level function to determine whether or not a type is unboxa
sUnbox :: forall a. (SBoolI (Unboxable a)) => SBool (Unboxable a)
sUnbox = sbool @(Unboxable a)

sNumeric :: forall a. (SBoolI (Numeric a)) => SBool (Numeric a)
sNumeric = sbool @(Numeric a)

type family When (flag :: Bool) (c :: Constraint) :: Constraint where
    When 'True c = c
    When 'False c = () -- empty constraint

type UnboxIf a = When (Unboxable a) (VU.Unbox a)

type family IntegralTypes (a :: Type) :: Bool where
    IntegralTypes Integer = 'True
    IntegralTypes Int = 'True
    IntegralTypes Int8 = 'True
    IntegralTypes Int16 = 'True
    IntegralTypes Int32 = 'True
    IntegralTypes Int64 = 'True
    IntegralTypes Word = 'True
    IntegralTypes Word8 = 'True
    IntegralTypes Word16 = 'True
    IntegralTypes Word32 = 'True
    IntegralTypes Word64 = 'True
    IntegralTypes _ = 'False

sIntegral :: forall a. (SBoolI (IntegralTypes a)) => SBool (IntegralTypes a)
sIntegral = sbool @(IntegralTypes a)

type IntegralIf a = When (IntegralTypes a) (Integral a)

type family FloatingTypes (a :: Type) :: Bool where
    FloatingTypes Float = 'True
    FloatingTypes Double = 'True
    FloatingTypes _ = 'False

sFloating :: forall a. (SBoolI (FloatingTypes a)) => SBool (FloatingTypes a)
sFloating = sbool @(FloatingTypes a)

type FloatingIf a = When (FloatingTypes a) (Real a, Fractional a)