{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE InstanceSigs               #-}
{-# LANGUAGE PolyKinds                  #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE TypeApplications           #-}

module Control.Monad.Trans.Ix
  ( Ix (..)
  , liftIx
  , unsafeLiftIx
  ) where

import           Control.Monad (MonadPlus (..))
import           Control.Monad.Indexed
import           Data.Coerce (coerce)
import           Data.Kind (Type)
import qualified Prelude as P
import           Prelude hiding (Monad (..), pure)


------------------------------------------------------------------------------
-- | The free indexed monad generated from a monad 'm'. Users are not expected
-- to use 'Ix' directly, but to newtype over it, specializing the kinds of 'i'
-- and 'j' as necessary.
--
-- GeneralizedNewtypeDeriving can be used to get the instances of 'IxFunctor',
-- 'IxPointed', 'IxApplicative', 'IxMonad', 'IxMonadZero' and 'IxMonadPlus' for
-- free.
newtype Ix (m :: Type -> Type) i j a = Ix
  { forall {k} {k} (m :: * -> *) (i :: k) (j :: k) a. Ix m i j a -> m a
runIx :: m a
  }
  deriving (forall a b. a -> Ix m i j b -> Ix m i j a
forall a b. (a -> b) -> Ix m i j a -> Ix m i j b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
forall (m :: * -> *) k (i :: k) k (j :: k) a b.
Functor m =>
a -> Ix m i j b -> Ix m i j a
forall (m :: * -> *) k (i :: k) k (j :: k) a b.
Functor m =>
(a -> b) -> Ix m i j a -> Ix m i j b
<$ :: forall a b. a -> Ix m i j b -> Ix m i j a
$c<$ :: forall (m :: * -> *) k (i :: k) k (j :: k) a b.
Functor m =>
a -> Ix m i j b -> Ix m i j a
fmap :: forall a b. (a -> b) -> Ix m i j a -> Ix m i j b
$cfmap :: forall (m :: * -> *) k (i :: k) k (j :: k) a b.
Functor m =>
(a -> b) -> Ix m i j a -> Ix m i j b
Functor, forall a. a -> Ix m i j a
forall a b. Ix m i j a -> Ix m i j b -> Ix m i j a
forall a b. Ix m i j a -> Ix m i j b -> Ix m i j b
forall a b. Ix m i j (a -> b) -> Ix m i j a -> Ix m i j b
forall a b c.
(a -> b -> c) -> Ix m i j a -> Ix m i j b -> Ix m i j c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
forall {m :: * -> *} {k} {i :: k} {k} {j :: k}.
Applicative m =>
Functor (Ix m i j)
forall (m :: * -> *) k (i :: k) k (j :: k) a.
Applicative m =>
a -> Ix m i j a
forall (m :: * -> *) k (i :: k) k (j :: k) a b.
Applicative m =>
Ix m i j a -> Ix m i j b -> Ix m i j a
forall (m :: * -> *) k (i :: k) k (j :: k) a b.
Applicative m =>
Ix m i j a -> Ix m i j b -> Ix m i j b
forall (m :: * -> *) k (i :: k) k (j :: k) a b.
Applicative m =>
Ix m i j (a -> b) -> Ix m i j a -> Ix m i j b
forall (m :: * -> *) k (i :: k) k (j :: k) a b c.
Applicative m =>
(a -> b -> c) -> Ix m i j a -> Ix m i j b -> Ix m i j c
<* :: forall a b. Ix m i j a -> Ix m i j b -> Ix m i j a
$c<* :: forall (m :: * -> *) k (i :: k) k (j :: k) a b.
Applicative m =>
Ix m i j a -> Ix m i j b -> Ix m i j a
*> :: forall a b. Ix m i j a -> Ix m i j b -> Ix m i j b
$c*> :: forall (m :: * -> *) k (i :: k) k (j :: k) a b.
Applicative m =>
Ix m i j a -> Ix m i j b -> Ix m i j b
liftA2 :: forall a b c.
(a -> b -> c) -> Ix m i j a -> Ix m i j b -> Ix m i j c
$cliftA2 :: forall (m :: * -> *) k (i :: k) k (j :: k) a b c.
Applicative m =>
(a -> b -> c) -> Ix m i j a -> Ix m i j b -> Ix m i j c
<*> :: forall a b. Ix m i j (a -> b) -> Ix m i j a -> Ix m i j b
$c<*> :: forall (m :: * -> *) k (i :: k) k (j :: k) a b.
Applicative m =>
Ix m i j (a -> b) -> Ix m i j a -> Ix m i j b
pure :: forall a. a -> Ix m i j a
$cpure :: forall (m :: * -> *) k (i :: k) k (j :: k) a.
Applicative m =>
a -> Ix m i j a
Applicative, forall a. a -> Ix m i j a
forall a b. Ix m i j a -> Ix m i j b -> Ix m i j b
forall a b. Ix m i j a -> (a -> Ix m i j b) -> Ix m i j b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
forall {m :: * -> *} {k} {i :: k} {k} {j :: k}.
Monad m =>
Applicative (Ix m i j)
forall (m :: * -> *) k (i :: k) k (j :: k) a.
Monad m =>
a -> Ix m i j a
forall (m :: * -> *) k (i :: k) k (j :: k) a b.
Monad m =>
Ix m i j a -> Ix m i j b -> Ix m i j b
forall (m :: * -> *) k (i :: k) k (j :: k) a b.
Monad m =>
Ix m i j a -> (a -> Ix m i j b) -> Ix m i j b
return :: forall a. a -> Ix m i j a
$creturn :: forall (m :: * -> *) k (i :: k) k (j :: k) a.
Monad m =>
a -> Ix m i j a
>> :: forall a b. Ix m i j a -> Ix m i j b -> Ix m i j b
$c>> :: forall (m :: * -> *) k (i :: k) k (j :: k) a b.
Monad m =>
Ix m i j a -> Ix m i j b -> Ix m i j b
>>= :: forall a b. Ix m i j a -> (a -> Ix m i j b) -> Ix m i j b
$c>>= :: forall (m :: * -> *) k (i :: k) k (j :: k) a b.
Monad m =>
Ix m i j a -> (a -> Ix m i j b) -> Ix m i j b
P.Monad)

instance Functor m => IxFunctor (Ix m) where
  imap :: forall a b (j :: k) (k2 :: k1).
(a -> b) -> Ix m j k2 a -> Ix m j k2 b
imap = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap

instance Applicative m => IxPointed (Ix m) where
  ireturn :: forall a (i :: k). a -> Ix m i i a
ireturn = forall (f :: * -> *) a. Applicative f => a -> f a
P.pure

instance Applicative m => IxApplicative (Ix m) where
  iap
      :: forall i j k a b
       . Ix m i j (a -> b)
      -> Ix m j k a
      -> Ix m i k b
  iap :: forall {k} {k} {k} (i :: k) (j :: k) (k :: k) a b.
Ix m i j (a -> b) -> Ix m j k a -> Ix m i k b
iap = coerce :: forall a b. Coercible a b => a -> b
coerce forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
(<*>) @m @a @b

instance P.Monad m => IxMonad (Ix m) where
  ibind
      :: forall i j k a b
       . (a -> Ix m j k b)
      -> Ix m i j a
      -> Ix m i k b
  ibind :: forall {k} {k} {k} (i :: k) (j :: k) (k :: k) a b.
(a -> Ix m j k b) -> Ix m i j a -> Ix m i k b
ibind = coerce :: forall a b. Coercible a b => a -> b
coerce forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
(=<<) @m @a @b

instance MonadPlus m => IxMonadZero (Ix m) where
  imzero
      :: forall i j a
       . Ix m i j a
  imzero :: forall {k} {k} (i :: k) (j :: k) a. Ix m i j a
imzero = coerce :: forall a b. Coercible a b => a -> b
coerce forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadPlus m => m a
mzero @m @a

instance MonadPlus m => IxMonadPlus (Ix m) where
  implus
      :: forall i j a
       . Ix m i j a
      -> Ix m i j a
      -> Ix m i j a
  implus :: forall {k} {k} (i :: k) (j :: k) a.
Ix m i j a -> Ix m i j a -> Ix m i j a
implus = coerce :: forall a b. Coercible a b => a -> b
coerce forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
mplus @m @a


------------------------------------------------------------------------------
-- | Lift an 'm' action into 'Ix m', maintaining the current index.
liftIx :: m a -> Ix m i i a
liftIx :: forall {k} (m :: * -> *) a (i :: k). m a -> Ix m i i a
liftIx = coerce :: forall a b. Coercible a b => a -> b
coerce

------------------------------------------------------------------------------
-- | Lift an 'm' action into 'Ix m', changing the current index. 'unsafeLiftIx'
-- is obviously unsafe due to the fact that it can arbitrarily change the
-- index.
unsafeLiftIx :: m a -> Ix m i j a
unsafeLiftIx :: forall {k} {k} (m :: * -> *) a (i :: k) (j :: k). m a -> Ix m i j a
unsafeLiftIx = coerce :: forall a b. Coercible a b => a -> b
coerce