{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeData #-}
module AST.Types where

import Data.Int (Int32, Int64)
import Data.GADT.Compare
import Data.GADT.Show
import Data.Kind (Type)
import Data.Type.Equality

import Data


type data Ty
  = TNil
  | TPair Ty Ty
  | TEither Ty Ty
  | TLEither Ty Ty
  | TMaybe Ty
  | TArr Nat Ty  -- ^ rank, element type
  | TScal ScalTy
  | TAccum Ty  -- ^ contained type must be a monoid type

type data ScalTy = TI32 | TI64 | TF32 | TF64 | TBool

type STy :: Ty -> Type
data STy t where
  STNil :: STy TNil
  STPair :: STy a -> STy b -> STy (TPair a b)
  STEither :: STy a -> STy b -> STy (TEither a b)
  STLEither :: STy a -> STy b -> STy (TLEither a b)
  STMaybe :: STy a -> STy (TMaybe a)
  STArr :: SNat n -> STy t -> STy (TArr n t)
  STScal :: SScalTy t -> STy (TScal t)
  STAccum :: SMTy t -> STy (TAccum t)
deriving instance Show (STy t)

instance GCompare STy where
  gcompare = \cases
    STNil STNil -> GEQ
    STNil _ -> GLT ; _ STNil -> GGT
    (STPair a b) (STPair a' b') -> gorderingLift2 (gcompare a a') (gcompare b b')
    STPair{} _ -> GLT ; _ STPair{} -> GGT
    (STEither a b) (STEither a' b') -> gorderingLift2 (gcompare a a') (gcompare b b')
    STEither{} _ -> GLT ; _ STEither{} -> GGT
    (STLEither a b) (STLEither a' b') -> gorderingLift2 (gcompare a a') (gcompare b b')
    STLEither{} _ -> GLT ; _ STLEither{} -> GGT
    (STMaybe a) (STMaybe a') -> gorderingLift1 (gcompare a a')
    STMaybe{} _ -> GLT ; _ STMaybe{} -> GGT
    (STArr n t) (STArr n' t') -> gorderingLift2 (gcompare n n') (gcompare t t')
    STArr{} _ -> GLT ; _ STArr{} -> GGT
    (STScal t) (STScal t') -> gorderingLift1 (gcompare t t')
    STScal{} _ -> GLT ; _ STScal{} -> GGT
    (STAccum t) (STAccum t') -> gorderingLift1 (gcompare t t')
    -- STAccum{} _ -> GLT ; _ STAccum{} -> GGT

instance TestEquality STy where testEquality = geq
instance GEq STy where geq = defaultGeq
instance GShow STy where gshowsPrec = defaultGshowsPrec

-- | Monoid types
type SMTy :: Ty -> Type
data SMTy t where
  SMTNil :: SMTy TNil
  SMTPair :: SMTy a -> SMTy b -> SMTy (TPair a b)
  SMTLEither :: SMTy a -> SMTy b -> SMTy (TLEither a b)
  SMTMaybe :: SMTy a -> SMTy (TMaybe a)
  SMTArr :: SNat n -> SMTy t -> SMTy (TArr n t)
  SMTScal :: ScalIsNumeric t ~ True => SScalTy t -> SMTy (TScal t)
deriving instance Show (SMTy t)

instance GCompare SMTy where
  gcompare = \cases
    SMTNil SMTNil -> GEQ
    SMTNil _ -> GLT ; _ SMTNil -> GGT
    (SMTPair a b) (SMTPair a' b') -> gorderingLift2 (gcompare a a') (gcompare b b')
    SMTPair{} _ -> GLT ; _ SMTPair{} -> GGT
    (SMTLEither a b) (SMTLEither a' b') -> gorderingLift2 (gcompare a a') (gcompare b b')
    SMTLEither{} _ -> GLT ; _ SMTLEither{} -> GGT
    (SMTMaybe a) (SMTMaybe a') -> gorderingLift1 (gcompare a a')
    SMTMaybe{} _ -> GLT ; _ SMTMaybe{} -> GGT
    (SMTArr n t) (SMTArr n' t') -> gorderingLift2 (gcompare n n') (gcompare t t')
    SMTArr{} _ -> GLT ; _ SMTArr{} -> GGT
    (SMTScal t) (SMTScal t') -> gorderingLift1 (gcompare t t')
    -- SMTScal{} _ -> GLT ; _ SMTScal{} -> GGT

instance TestEquality SMTy where testEquality = geq
instance GEq SMTy where geq = defaultGeq
instance GShow SMTy where gshowsPrec = defaultGshowsPrec

fromSMTy :: SMTy t -> STy t
fromSMTy = \case
  SMTNil -> STNil
  SMTPair t1 t2 -> STPair (fromSMTy t1) (fromSMTy t2)
  SMTLEither t1 t2 -> STLEither (fromSMTy t1) (fromSMTy t2)
  SMTMaybe t -> STMaybe (fromSMTy t)
  SMTArr n t -> STArr n (fromSMTy t)
  SMTScal sty -> STScal sty

data SScalTy t where
  STI32 :: SScalTy TI32
  STI64 :: SScalTy TI64
  STF32 :: SScalTy TF32
  STF64 :: SScalTy TF64
  STBool :: SScalTy TBool
deriving instance Show (SScalTy t)

instance GCompare SScalTy where
  gcompare = \cases
    STI32 STI32 -> GEQ
    STI32 _ -> GLT ; _ STI32 -> GGT
    STI64 STI64 -> GEQ
    STI64 _ -> GLT ; _ STI64 -> GGT
    STF32 STF32 -> GEQ
    STF32 _ -> GLT ; _ STF32 -> GGT
    STF64 STF64 -> GEQ
    STF64 _ -> GLT ; _ STF64 -> GGT
    STBool STBool -> GEQ
    -- STBool _ -> GLT ; _ STBool -> GGT

instance TestEquality SScalTy where testEquality = geq
instance GEq SScalTy where geq = defaultGeq
instance GShow SScalTy where gshowsPrec = defaultGshowsPrec

scalRepIsShow :: SScalTy t -> Dict (Show (ScalRep t))
scalRepIsShow STI32 = Dict
scalRepIsShow STI64 = Dict
scalRepIsShow STF32 = Dict
scalRepIsShow STF64 = Dict
scalRepIsShow STBool = Dict

type TIx = TScal TI64

tIx :: STy TIx
tIx = STScal STI64

type family ScalRep t where
  ScalRep TI32 = Int32
  ScalRep TI64 = Int64
  ScalRep TF32 = Float
  ScalRep TF64 = Double
  ScalRep TBool = Bool

type family ScalIsNumeric t where
  ScalIsNumeric TI32 = True
  ScalIsNumeric TI64 = True
  ScalIsNumeric TF32 = True
  ScalIsNumeric TF64 = True
  ScalIsNumeric TBool = False

type family ScalIsFloating t where
  ScalIsFloating TI32 = False
  ScalIsFloating TI64 = False
  ScalIsFloating TF32 = True
  ScalIsFloating TF64 = True
  ScalIsFloating TBool = False

type family ScalIsIntegral t where
  ScalIsIntegral TI32 = True
  ScalIsIntegral TI64 = True
  ScalIsIntegral TF32 = False
  ScalIsIntegral TF64 = False
  ScalIsIntegral TBool = False

-- | Returns true for arrays /and/ accumulators.
hasArrays :: STy t' -> Bool
hasArrays STNil = False
hasArrays (STPair a b) = hasArrays a || hasArrays b
hasArrays (STEither a b) = hasArrays a || hasArrays b
hasArrays (STLEither a b) = hasArrays a || hasArrays b
hasArrays (STMaybe t) = hasArrays t
hasArrays STArr{} = True
hasArrays STScal{} = False
hasArrays STAccum{} = True

type family Tup env where
  Tup '[] = TNil
  Tup (t : ts) = TPair (Tup ts) t

mkTup :: f TNil -> (forall a b. f a -> f b -> f (TPair a b))
      -> SList f list -> f (Tup list)
mkTup nil _    SNil = nil
mkTup nil pair (e `SCons` es) = pair (mkTup nil pair es) e

tTup :: SList STy env -> STy (Tup env)
tTup = mkTup STNil STPair

unTup :: (forall a b. c (TPair a b) -> (c a, c b))
      -> SList f list -> c (Tup list) -> SList c list
unTup _ SNil _ = SNil
unTup unpack (_ `SCons` list) tup =
  let (xs, x) = unpack tup
  in x `SCons` unTup unpack list xs

type family InvTup core env where
  InvTup core '[] = core
  InvTup core (t : ts) = InvTup (TPair core t) ts