{-# language AllowAmbiguousTypes  #-}
{-# language ConstraintKinds      #-}
{-# language DataKinds            #-}
{-# language GADTs                #-}
{-# language ImpredicativeTypes   #-}
{-# language PolyKinds            #-}
{-# language ScopedTypeVariables  #-}
{-# language TypeApplications     #-}
{-# language TypeFamilies         #-}
{-# language TypeOperators        #-}
{-# language UndecidableInstances #-}
-- | 'Atom's represent the shape of a type, possibly with variables,
-- which can be 'Interpret'ed given a list of types.
module Data.PolyKinded.Atom (
  -- * Atoms
  Atom(..), (:$:), (:~:), (:~~:)
  -- * Type variables
, TyVar(..)
  -- ** Aliases for specific variables
, Var0, Var1, Var2, Var3, Var4, Var5, Var6, Var7, Var8, Var9
  -- * Intepretation
, Interpret, InterpretVar, Satisfies, ContainsTyVar
  -- ** Auxiliary types for interpretation
, ForAllI(..), SuchThatI(..)
  -- ** Records a type as constructor + list of variables
, WrappedI(..), toWrappedI, fromWrappedI
) where

import           Data.Kind
import           Data.PolyKinded
import           GHC.Exts

-- | Well-scoped de Bruijn representation of type variables.
-- @TyVar d@ represents all the possible type variables which
-- can refer to the holes in kind @d@.
--
-- We recommend using the aliases 'Var0', 'Var1', ...
-- instead of the constructors, for further clarity.
data TyVar (d :: Type) (k :: TYPE r) where
  -- | First hole in @d@.
  VZ :: TyVar (x -> xs) x
  -- | Successor hole, increases the hole reference by 1.
  VS :: TyVar xs k -> TyVar (x -> xs) k

type Var0 = 'Var 'VZ
type Var1 = 'Var ('VS 'VZ)
type Var2 = 'Var ('VS ('VS 'VZ))
type Var3 = 'Var ('VS ('VS ('VS 'VZ)))
type Var4 = 'Var ('VS ('VS ('VS ('VS 'VZ))))
type Var5 = 'Var ('VS ('VS ('VS ('VS ('VS 'VZ)))))
type Var6 = 'Var ('VS ('VS ('VS ('VS ('VS ('VS 'VZ))))))
type Var7 = 'Var ('VS ('VS ('VS ('VS ('VS ('VS ('VS 'VZ)))))))
type Var8 = 'Var ('VS ('VS ('VS ('VS ('VS ('VS ('VS ('VS 'VZ))))))))
type Var9 = 'Var ('VS ('VS ('VS ('VS ('VS ('VS ('VS ('VS ('VS 'VZ)))))))))

infixr 5 :&:
infixr 5 :=>>:
-- | Shape of a type, possibly with type variables.
--
-- >>> :kind Kon [] :@: Var0 -- the type [a] for unknown a
-- Kon [] :@: Var0 :: Atom (* -> xs) *
data Atom (d :: Type) (k :: TYPE r) where
  -- | Represents a type variable.
  Var :: TyVar d k -> Atom d k
  -- | Represents a constant type, like 'Int'.
  Kon :: k         -> Atom d k
  -- | Represents type application.
  (:@:) :: Atom d (k1 -> k2) -> Atom d k1 -> Atom d k2
  -- | Represents the conjunction of two constraints.
  (:&:) :: Atom d Constraint -> Atom d Constraint -> Atom d Constraint
  -- | Represents universal quantification.
  ForAll  :: Atom (d1 -> d) Type -> Atom d Type
  -- | Represents constraint requirement, the "thick arrow" @=>@.
  (:=>>:) :: Atom d Constraint -> Atom d Type -> Atom d Type

-- | Represents an applied constructor.
-- Instead of @Kon [] :@: Var0$ you can write @[] :$: Var0$.
type f :$:  x = 'Kon f ':@: x
-- | Represents (homogeneous) type equality.
type a :~:  b = 'Kon (~) ':@: a ':@: b
-- | Represents heterogeneous type equality.
type a :~~: b = 'Kon (~~) ':@: a ':@: b

-- | Obtains the type in the list @tys@ referenced
-- by the type variable @t@.
--
-- >>> :kind! Interpret Var0 (LoT2 Int Bool)
-- Interpret Var0 (LoT2 Int Bool) :: *
-- = Int
-- >>> :kind! Interpret Var1 (LoT2 Int Bool)
-- Interpret Var1 (LoT2 Int Bool) :: *
-- = Bool
type family InterpretVar (t :: TyVar d k) (tys :: LoT d) :: k where
  InterpretVar 'VZ     tys = HeadLoT tys
  InterpretVar ('VS v) tys = InterpretVar v (TailLoT tys)

-- | Replaces the holes in the 'Atom' @t@ by the elements of
-- the list of types @tys@. The amount and kind of types in @tys@
-- must match statically those required by the 'Atom'.
--
-- >>> :kind! Interpret ([] :$: Var0) (LoT1 Int)
-- Interpret ([] :$: Var0) (LoT1 Int) :: *
-- = [Int]
type family Interpret (t :: Atom d k) (tys :: LoT d) :: k where
  Interpret ('Var v)     tys = InterpretVar v tys
  Interpret ('Kon t)     tys = t
  Interpret (f ':@: x)   tys = (Interpret f tys) (Interpret x tys)
  Interpret (c ':&: d)   tys = (Interpret c tys, Interpret d tys)
  Interpret ('ForAll f)  tys = ForAllI f tys
  Interpret (c ':=>>: f) tys = SuchThatI c f tys

-- | Auxiliary type for interpretation of the 'ForAll' atom.
-- Required because a type family like 'Interpret' cannot return
-- a polymorphic type.
newtype ForAllI (f :: Atom (d1 -> d) Type) (tys :: LoT d) where
  ForAllI :: (forall t. Interpret f (t ':&&: tys)) -> ForAllI f tys

-- | Records a value of type @f@ applied to the list @tys@.
--
-- >>> :t WrapI [1] :: WrappedI ([] :$: Var0) (LoT1 Int)
-- WrapI [1] :: WrappedI ([] :$: Var0) (LoT1 Int)
newtype WrappedI (f :: Atom d Type) (tys :: LoT d) =
  WrapI { WrappedI f tys -> Interpret f tys
unwrapI :: Interpret f tys }

toWrappedI :: forall f tys t. ForAllI f tys -> WrappedI f (t ':&&: tys)
toWrappedI :: ForAllI f tys -> WrappedI f (t ':&&: tys)
toWrappedI (ForAllI forall (t :: d1). Interpret f (t ':&&: tys)
x) = Interpret f (t ':&&: tys) -> WrappedI f (t ':&&: tys)
forall d (f :: Atom d *) (tys :: LoT d).
Interpret f tys -> WrappedI f tys
WrapI (Interpret f (t ':&&: tys)
forall (t :: d1). Interpret f (t ':&&: tys)
x @t)

fromWrappedI :: forall f tys. (forall t. WrappedI f (t ':&&: tys)) -> ForAllI f tys
fromWrappedI :: (forall (t :: d1). WrappedI f (t ':&&: tys)) -> ForAllI f tys
fromWrappedI = forall b.
Coercible (forall (t :: d1). WrappedI f (t ':&&: tys)) b =>
(forall (t :: d1). WrappedI f (t ':&&: tys)) -> b
coerce @(forall t. WrappedI f (t ':&&: tys))

-- | Auxiliary type for interpretation of the '(:=>>:)' atom.
-- Required because a type family like 'Interpret' cannot return
-- a type with constraints.
newtype SuchThatI (c :: Atom d Constraint) (f :: Atom d Type) (tys :: LoT d) where
  SuchThatI :: (Interpret c tys => Interpret f tys) -> SuchThatI c f tys

-- | Interprets a list of 'Atoms' representing constraints
-- into the actual constraints. This is a specialization of
-- 'Interpret' for the case of constraints.
--
-- >>> :kind! Satisfies '[Eq :$: Var0, Show :$: Var0] (LoT1 Int)
-- Satisfies '[Eq :$: Var0, Show :$: Var0] (LoT1 Int) :: Constraint
-- = (Eq Int, (Show Int, () :: Constraint))
type family Satisfies (cs :: [Atom d Constraint]) (tys :: LoT d) :: Constraint where
  Satisfies '[]       tys = ()
  Satisfies (c ': cs) tys = (Interpret c tys, Satisfies cs tys)

-- | Determines whether a given type variable @v@ is used within an 'Atom' @t@.
-- If not, we know that the atom is constant with respect to that variable.
type family ContainsTyVar (v :: TyVar d k) (t :: Atom d p) :: Bool where
  ContainsTyVar v ('Var v)     = 'True
  ContainsTyVar v ('Var w)     = 'False
  ContainsTyVar v ('Kon t)     = 'False
  ContainsTyVar v (f ':@: x)   = Or (ContainsTyVar v f) (ContainsTyVar v x)
  ContainsTyVar v (x ':&: y)   = Or (ContainsTyVar v x) (ContainsTyVar v y)
  ContainsTyVar v (c ':=>>: f) = Or (ContainsTyVar v c) (ContainsTyVar v f)
  ContainsTyVar v ('ForAll f)  = ContainsTyVar ('VS v) f

type family Or (x :: Bool) (y :: Bool) :: Bool where
  Or 'True  thing  = 'True
  Or thing 'True   = 'True
  Or 'False 'False = 'False