{-# LANGUAGE GADTs #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module AST.Sparse.Types where

import AST.Types

import Data.Kind (Type, Constraint)
import Data.Type.Equality


data Sparse t t' where
  SpSparse :: Sparse t t' -> Sparse t (TMaybe t')
  SpAbsent :: Sparse t TNil

  SpPair :: Sparse a a' -> Sparse b b' -> Sparse (TPair a b) (TPair a' b')
  SpLEither :: Sparse a a' -> Sparse b b' -> Sparse (TLEither a b) (TLEither a' b')
  SpMaybe :: Sparse t t' -> Sparse (TMaybe t) (TMaybe t')
  SpArr :: Sparse t t' -> Sparse (TArr n t) (TArr n t')
  SpScal :: Sparse (TScal t) (TScal t)
deriving instance Show (Sparse t t')

class ApplySparse f where
  applySparse :: Sparse t t' -> f t -> f t'

instance ApplySparse STy where
  applySparse (SpSparse s) t = STMaybe (applySparse s t)
  applySparse SpAbsent _ = STNil
  applySparse (SpPair s1 s2) (STPair t1 t2) = STPair (applySparse s1 t1) (applySparse s2 t2)
  applySparse (SpLEither s1 s2) (STLEither t1 t2) = STLEither (applySparse s1 t1) (applySparse s2 t2)
  applySparse (SpMaybe s) (STMaybe t) = STMaybe (applySparse s t)
  applySparse (SpArr s) (STArr n t) = STArr n (applySparse s t)
  applySparse SpScal t = t

instance ApplySparse SMTy where
  applySparse (SpSparse s) t = SMTMaybe (applySparse s t)
  applySparse SpAbsent _ = SMTNil
  applySparse (SpPair s1 s2) (SMTPair t1 t2) = SMTPair (applySparse s1 t1) (applySparse s2 t2)
  applySparse (SpLEither s1 s2) (SMTLEither t1 t2) = SMTLEither (applySparse s1 t1) (applySparse s2 t2)
  applySparse (SpMaybe s) (SMTMaybe t) = SMTMaybe (applySparse s t)
  applySparse (SpArr s) (SMTArr n t) = SMTArr n (applySparse s t)
  applySparse SpScal t = t


class IsSubType s where
  type IsSubTypeSubject (s :: k -> k -> Type) (f :: k -> Type) :: Constraint
  subtApply :: IsSubTypeSubject s f => s t t' -> f t -> f t'
  subtTrans :: s a b -> s b c -> s a c
  subtFull :: IsSubTypeSubject s f => f t -> s t t

instance IsSubType (:~:) where
  type IsSubTypeSubject (:~:) f = ()
  subtApply = gcastWith
  subtTrans = trans
  subtFull _ = Refl

instance IsSubType Sparse where
  type IsSubTypeSubject Sparse f = f ~ SMTy
  subtApply = applySparse

  subtTrans s1 (SpSparse s2) = SpSparse (subtTrans s1 s2)
  subtTrans _ SpAbsent = SpAbsent
  subtTrans (SpPair s1a s1b) (SpPair s2a s2b) = SpPair (subtTrans s1a s2a) (subtTrans s1b s2b)
  subtTrans (SpLEither s1a s1b) (SpLEither s2a s2b) = SpLEither (subtTrans s1a s2a) (subtTrans s1b s2b)
  subtTrans (SpSparse s1) (SpMaybe s2) = SpSparse (subtTrans s1 s2)
  subtTrans (SpMaybe s1) (SpMaybe s2) = SpMaybe (subtTrans s1 s2)
  subtTrans (SpArr s1) (SpArr s2) = SpArr (subtTrans s1 s2)
  subtTrans SpScal SpScal = SpScal

  subtFull = spDense

spDense :: SMTy t -> Sparse t t
spDense SMTNil = SpAbsent
spDense (SMTPair t1 t2) = SpPair (spDense t1) (spDense t2)
spDense (SMTLEither t1 t2) = SpLEither (spDense t1) (spDense t2)
spDense (SMTMaybe t) = SpMaybe (spDense t)
spDense (SMTArr _ t) = SpArr (spDense t)
spDense (SMTScal _) = SpScal

isDense :: SMTy t -> Sparse t t' -> Maybe (t :~: t')
isDense SMTNil SpAbsent = Just Refl
isDense _ SpSparse{} = Nothing
isDense _ SpAbsent = Nothing
isDense (SMTPair t1 t2) (SpPair s1 s2)
  | Just Refl <- isDense t1 s1, Just Refl <- isDense t2 s2 = Just Refl
  | otherwise = Nothing
isDense (SMTLEither t1 t2) (SpLEither s1 s2)
  | Just Refl <- isDense t1 s1, Just Refl <- isDense t2 s2 = Just Refl
  | otherwise = Nothing
isDense (SMTMaybe t) (SpMaybe s)
  | Just Refl <- isDense t s = Just Refl
  | otherwise = Nothing
isDense (SMTArr _ t) (SpArr s)
  | Just Refl <- isDense t s = Just Refl
  | otherwise = Nothing
isDense (SMTScal _) SpScal = Just Refl

isAbsent :: Sparse t t' -> Bool
isAbsent (SpSparse s) = isAbsent s
isAbsent SpAbsent = True
isAbsent (SpPair s1 s2) = isAbsent s1 && isAbsent s2
isAbsent (SpLEither s1 s2) = isAbsent s1 && isAbsent s2
isAbsent (SpMaybe s) = isAbsent s
isAbsent (SpArr s) = isAbsent s
isAbsent SpScal = False