{-# 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 Haskus.Utils.EGADT where

import Unsafe.Coerce

import Haskus.Utils.Monad
import Haskus.Utils.Variant
import Haskus.Utils.VariantF
import Haskus.Utils.Types

-- $setup
-- >>> :set -XDataKinds
-- >>> :set -XTypeApplications
-- >>> :set -XTypeOperators
-- >>> :set -XFlexibleContexts
-- >>> :set -XTypeFamilies
-- >>> :set -XPatternSynonyms
-- >>> :set -XDeriveFunctor
-- >>> :set -XGADTs
-- >>> :set -XPolyKinds
-- >>> :set -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 :: 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 (KnownNat i => Word
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 :: 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
== KnownNat i => Word
forall (n :: Nat). KnownNat n => Word
natValue' @i)
  Index i fs ast a -> Maybe (Index i fs ast 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 :: EGADT xs a -> HBase (EGADT xs) (EGADT xs) a
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 :: HBase (EGADT xs) (EGADT xs) a -> EGADT xs a
hembed = HBase (EGADT 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 $bVF :: f (EGADT fs) a -> EGADT fs a
$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) -> (Void# -> r) -> r
VF x <- ( ( \ ( EGADT (HVariantF v) ) -> fromHVariantAt @(IndexOf f fs) @fs 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 (Index (IndexOf f fs) fs (EGADT fs) a
-> VariantF (ApplyAll (EGADT fs) 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))