{-# language ConstraintKinds       #-}
{-# language FlexibleContexts      #-}
{-# language FlexibleInstances     #-}
{-# language GADTs                 #-}
{-# language MultiParamTypeClasses #-}
{-# language ScopedTypeVariables   #-}
{-# language TypeFamilies          #-}
{-# language TypeInType            #-}
{-# language TypeOperators         #-}
{-# language UndecidableInstances  #-}
-- | Representation of types as constructor + list of types.
module Data.PolyKinded (
  -- * Lists of types and application
  LoT(..), (:@@:), LoT1, LoT2
, HeadLoT, TailLoT, SpineLoT
  -- ** Singleton for list of types
, SLoT(..), SForLoT(..), Proxy(..)
  -- * Splitting types
, SplitF, Nat(..), TyEnv(..), SplitN
) where

import           Data.Kind
import           Data.Proxy

infixr 5 :&&:
-- | @LoT k@ represents a list of types which can be applied
-- to a data type of kind @k@.
data LoT k where
  -- | Empty list of types.
  LoT0    ::                LoT Type
  -- | Cons a type with a list of types.
  (:&&:)  :: k -> LoT ks -> LoT (k -> ks)

-- | List of types with a single element.
type LoT1 a = a ':&&: 'LoT0
-- | List of types with two elements.
type LoT2 a b = a ':&&: b ':&&: 'LoT0

-- | Apply a list of types to a type constructor.
--
-- >>> :kind! Either :@@: (Int :&&: Bool :&&: LoT0)
-- Either Int Bool :: Type
type family (f :: k) :@@: (tys :: LoT k) :: Type where
  f :@@: _  = f
  f :@@: as = f (HeadLoT as) :@@: TailLoT as

-- | Head of a non-empty list of types.
--
-- >>> :kind! HeadLoT (Int :&&: LoT0)
-- Int :: Type
type family HeadLoT (tys :: LoT (k -> k')) :: k where
  HeadLoT (a ':&&: _) = a

-- | Tail of a non-empty list of types.
--
-- >>> :kind! TailLoT (Int :&&: Bool :&&: LoT0)
-- Bool :&&: LoT0 :: LoT (Type -> Type)
type family TailLoT (tys :: LoT (k -> k')) :: LoT k' where
  TailLoT (_ ':&&: as) = as

-- | Construct the spine of a list of types whose length is known.
--
-- It can be useful to introduce unification variables for lists of types which
-- will be fully instantiated during constraint resolution.
-- A constraint @p ~ SpineLoT p@ will thus instantiate the spine of @p@.
--
-- On concrete lists, this is the identity function.
type family SpineLoT (ts :: LoT k) :: LoT k where
  SpineLoT (ts :: LoT (k -> k')) = HeadLoT ts ':&&: SpineLoT (TailLoT ts)
  SpineLoT (ts :: LoT Type)      = 'LoT0

data SLoT (l :: LoT k) where
  SLoT0 :: SLoT 'LoT0
  SLoTA :: Proxy t -> SLoT ts -> SLoT (t ':&&: ts)

class SForLoT (l :: LoT k) where
  slot :: SLoT l
instance (l ~ 'LoT0) => SForLoT (l :: LoT Type) where
  slot :: SLoT l
slot = SLoT l
SLoT 'LoT0
SLoT0
instance (l ~ (t ':&&: ts), SForLoT ts) => SForLoT (l :: LoT (k -> k')) where
  slot :: SLoT l
slot = Proxy t -> SLoT ts -> SLoT (t ':&&: ts)
forall k ks (t :: k) (ts :: LoT ks).
Proxy t -> SLoT ts -> SLoT (t ':&&: ts)
SLoTA Proxy t
forall k (t :: k). Proxy t
Proxy SLoT ts
forall k (l :: LoT k). SForLoT l => SLoT l
slot

-- | Split a type @t@ until the constructor @f@ is found.
--
-- >>> :kind! SplitF (Either Int Bool) Either
-- Int :&&: Bool :&&: LoT0 :: LoT (Type -> Type -> Type)
-- >>> :kind! SplitF (Either Int Bool) (Either Int)
-- Bool :&&: LoT0 :: LoT (Type -> Type)
type SplitF (t :: d) (f :: k) = SplitF' t f 'LoT0
type family SplitF' (t :: d) (f :: k) (p :: LoT l) :: LoT k where
  SplitF' f     f acc = acc
  SplitF' (t a) f acc = SplitF' t f (a ':&&: acc)

-- | Simple natural numbers.
data Nat = Z | S Nat

-- | A type constructor and a list of types that can be applied to it.
data TyEnv where
  TyEnv :: forall k. k -> LoT k -> TyEnv

-- | Split a type @t@ until its list of types has length @n@.
--
-- >>> :kind! SplitN (Either Int Bool) (S (S Z))
-- TyEnv Either (Int :&&: Bool :&&: LoT0) :: TyEnv
-- >>> :kind! SplitF (Either Int Bool) (S Z)
-- TyEnv (Either Int) (Bool :&&: LoT0) :: TyEnv
type family SplitN (n :: Nat) t :: TyEnv where
  SplitN n t = SplitN' n t 'LoT0
type family SplitN' (n :: Nat) (t :: d) (p :: LoT d) :: TyEnv where
  SplitN' 'Z     t            acc = 'TyEnv t acc
  SplitN' ('S n) (t (a :: l)) acc = SplitN' n t (a ':&&: acc)