{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}

module Data.Variant.EGADT where

import Unsafe.Coerce
import GHC.TypeLits
import Data.Kind
import Control.Monad

import Data.Variant
import Data.Variant.VariantF
import Data.Variant.Types
import Data.Variant.Functor

-- $setup
-- >>> :seti -XDataKinds
-- >>> :seti -XTypeApplications
-- >>> :seti -XTypeOperators
-- >>> :seti -XFlexibleContexts
-- >>> :seti -XTypeFamilies
-- >>> :seti -XPatternSynonyms
-- >>> :seti -XDeriveFunctor
-- >>> :seti -XGADTs
-- >>> :seti -XPolyKinds
-- >>> :seti -XPartialTypeSignatures
-- >>>
-- >>> :{
-- >>> data LamF (ast :: Type -> Type) t where
-- >>>   LamF :: ( ast a -> ast b ) -> LamF ast ( a -> b )
-- >>> 
-- >>> data AppF ast t where
-- >>>   AppF :: ast ( a -> b ) -> ast a -> AppF ast b
-- >>>
-- >>> data VarF ast t where
-- >>>   VarF :: String -> VarF ast Int
-- >>>
-- >>> type AST a = EGADT '[LamF,AppF,VarF] a
-- >>>
-- >>> :}
--
-- >>> let y = VF @(AST Int) (VarF "a")
-- >>> :t y
-- y :: EGADT [LamF, AppF, VarF] Int
--
-- >>> :{
-- >>> case y of
-- >>>   VF (VarF x) -> print x
-- >>>   _           -> putStrLn "Not a VarF"
-- >>> :}
-- "a"
--
-- >>> :{
-- >>> f :: AST Int -> AST Int
-- >>> f (VF (VarF x)) = VF (VarF "zz")
-- >>> f _             = error "Unhandled case"
-- >>> :}
--
-- >>> let z = VF (AppF (VF (LamF f)) (VF (VarF "a")))
-- >>> :t z
-- z :: EGADT [LamF, AppF, VarF] Int
--


-- | An EADT with an additional type parameter
newtype EGADT fs t = EGADT (HVariantF fs (EGADT fs) t)

newtype HVariantF (fs :: [ (k -> Type) -> ( k -> Type) ]) (ast :: k -> Type) (t :: k)
  = HVariantF (VariantF (ApplyAll ast fs) t)

toHVariantAt
  :: forall i fs ast a
  .  KnownNat i
  => (Index i fs) ast a -> VariantF (ApplyAll ast fs) a
{-# INLINABLE toHVariantAt #-}
toHVariantAt :: forall {t} {t} (i :: Nat) (fs :: [t -> t -> *]) (ast :: t)
       (a :: t).
KnownNat i =>
Index i fs ast a -> VariantF (ApplyAll ast fs) a
toHVariantAt Index i fs ast a
a = V (ApplyAll a (ApplyAll ast fs)) -> VariantF (ApplyAll ast fs) a
forall t (xs :: [t -> *]) (e :: t).
V (ApplyAll e xs) -> VariantF xs e
VariantF (Word -> Any -> V (ApplyAll a (ApplyAll ast fs))
forall (l :: [*]). Word -> Any -> V l
Variant (forall (n :: Nat). KnownNat n => Word
natValue' @i) (Index i fs ast a -> Any
forall a b. a -> b
unsafeCoerce Index i fs ast a
a))

fromHVariantAt
  :: forall i fs ast a
  .  KnownNat i
  => VariantF (ApplyAll ast fs) a -> Maybe ((Index i fs) ast a)
{-# INLINABLE fromHVariantAt #-}
fromHVariantAt :: forall {t} {t} (i :: Nat) (fs :: [t -> t -> *]) (ast :: t)
       (a :: t).
KnownNat i =>
VariantF (ApplyAll ast fs) a -> Maybe (Index i fs ast a)
fromHVariantAt (VariantF (Variant Word
t Any
a)) = do
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Word
t Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
== forall (n :: Nat). KnownNat n => Word
natValue' @i)
  Index i fs ast a -> Maybe (Index i fs ast a)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (Any -> Index i fs ast a
forall a b. a -> b
unsafeCoerce Any
a)

type instance HBase (EGADT xs) = HVariantF xs

instance HFunctor (HVariantF xs) => HRecursive (EGADT xs) where
  hproject :: HCoalgebra (HBase (EGADT xs)) (EGADT xs)
hproject (EGADT HVariantF xs (EGADT xs) a
a) = HBase (EGADT xs) (EGADT xs) a
HVariantF xs (EGADT xs) a
a

instance HFunctor (HVariantF xs) => HCorecursive (EGADT xs) where
  hembed :: HAlgebra (HBase (EGADT xs)) (EGADT xs)
hembed = HBase (EGADT xs) (EGADT xs) a -> EGADT xs a
HVariantF xs (EGADT xs) a -> EGADT xs a
forall {k} (fs :: [(k -> *) -> k -> *]) (t :: k).
HVariantF fs (EGADT fs) t -> EGADT fs t
EGADT

type family f :<! fs :: Constraint where
  f :<! fs = ( MemberAtIndex (IndexOf f fs) f fs )

type family MemberAtIndex (i :: Nat) f fs :: Constraint where
  MemberAtIndex i f fs = ( KnownNat i, Index i fs ~ f )

type family (:<<!) xs ys :: Constraint where
  '[]       :<<! ys = ()
  (x ': xs) :<<! ys = (x :<! ys, xs :<<! ys)

-- | Pattern-match in an extensible GADT
pattern VF :: forall e a f fs.
  ( e ~ EGADT fs a  -- allow easy use of TypeApplication to set the EGADT type
  , f :<! fs
  ) => f (EGADT fs) a -> EGADT fs a
pattern $mVF :: forall {r} {k} {e} {a :: k} {f :: (k -> *) -> k -> *}
       {fs :: [(k -> *) -> k -> *]}.
(e ~ EGADT fs a, f :<! fs) =>
EGADT fs a -> (f (EGADT fs) a -> r) -> ((# #) -> r) -> r
$bVF :: forall {k} e (a :: k) (f :: (k -> *) -> k -> *)
       (fs :: [(k -> *) -> k -> *]).
(e ~ EGADT fs a, f :<! fs) =>
f (EGADT fs) a -> EGADT fs a
VF x <- ( ( \ ( EGADT (HVariantF VariantF (ApplyAll (EGADT fs) fs) a
v) ) -> forall (i :: Nat) (fs :: [(k -> *) -> k -> *]) (ast :: k -> *)
       (a :: k).
KnownNat i =>
VariantF (ApplyAll ast fs) a -> Maybe (Index i fs ast a)
forall {t} {t} (i :: Nat) (fs :: [t -> t -> *]) (ast :: t)
       (a :: t).
KnownNat i =>
VariantF (ApplyAll ast fs) a -> Maybe (Index i fs ast a)
fromHVariantAt @(IndexOf f fs) @fs VariantF (ApplyAll (EGADT fs) fs) a
v ) -> Just x )
  where
    VF f (EGADT fs) a
x = HVariantF fs (EGADT fs) a -> EGADT fs a
forall {k} (fs :: [(k -> *) -> k -> *]) (t :: k).
HVariantF fs (EGADT fs) t -> EGADT fs t
EGADT (VariantF (ApplyAll (EGADT fs) fs) a -> HVariantF fs (EGADT fs) a
forall k (fs :: [(k -> *) -> k -> *]) (ast :: k -> *) (t :: k).
VariantF (ApplyAll ast fs) t -> HVariantF fs ast t
HVariantF (forall (i :: Nat) (fs :: [(k -> *) -> k -> *]) (ast :: k -> *)
       (a :: k).
KnownNat i =>
Index i fs ast a -> VariantF (ApplyAll ast fs) a
forall {t} {t} (i :: Nat) (fs :: [t -> t -> *]) (ast :: t)
       (a :: t).
KnownNat i =>
Index i fs ast a -> VariantF (ApplyAll ast fs) a
toHVariantAt @(IndexOf f fs) @fs f (EGADT fs) a
Index (IndexOf f fs) fs (EGADT fs) a
x))