-- Functor/Applicative/Monad typeclasses with an extra type parameter.
-- The extra parameter can be used to enforce ordering constraints on
-- MultiApplicative and MultiMonad actions, or to accumulate information
-- about the actions. When two actions are sequenced, the result type is
-- governed by a type family; when the type family is partial, this also
-- constrains the types of actions which can be sequenced.

{-# LANGUAGE KindSignatures, PolyKinds, TypeFamilies #-}
{-# LANGUAGE FlexibleContexts, FlexibleInstances     #-}

-- For the MIdentity example.
{-# LANGUAGE DataKinds #-}

import Control.Applicative
import Control.Monad
import Data.Void
import Data.Proxy

class MultiFunctor f where
   mfmap :: (a -> b) -> f p a -> f p b

(<!$>) :: MultiFunctor f => (a -> b) -> f p a -> f p b
(<!$>) = mfmap
infixl 4 <!$>

class MultiFunctor f => MultiApplicative f where
   type MAppResult f (p :: k) (q :: k)

   mpure :: a -> f p a

   (<!*>) :: (MAppResult f p q ~ Proxy r) => f p (a -> b) -> f q a -> f r b
   infixl 4 <!*>

   (!*>) :: (MAppResult f p q ~ Proxy r) => f p a -> f q b -> f r b
   a !*> b = flip const <!$> a <!*> b
   infixl 4 !*>

   (<!*) :: (MAppResult f p q ~ Proxy r) => f p a -> f q b -> f r a
   a <!* b = const <!$> a <!*> b
   infixl 4 <!*

class MultiApplicative m => MultiMonad m where
   (!>>=) :: (MAppResult m p q ~ Proxy r) => m p a -> (a -> m q b) -> m r b
   m !>>= f = mjoin (mfmap f m)
   infixl 1 !>>=

   mjoin :: (MAppResult m p q ~ Proxy r) => m p (m q a) -> m r a
   mjoin m = m !>>= id

newtype MultiWrap f (p :: k) a = MW { unMW :: f a }

instance Functor f => MultiFunctor (MultiWrap f) where
   mfmap f (MW x) = MW (fmap f x)

instance Applicative f => MultiApplicative (MultiWrap f) where
   type MAppResult (MultiWrap f) p q = Proxy Void

   mpure          = MW . pure
   MW f <!*> MW x = MW (f <*> x)
   MW f  !*> MW x = MW (f  *> x)
   MW f <!*  MW x = MW (f <*  x)

instance Monad m => MultiMonad (MultiWrap m) where
   MW m !>>= f  = MW $ m >>= unMW . f
   mjoin (MW m) = MW $ join $ fmap unMW m

-- A simple example of how MultiApplicative can be used to enforce ordering.
-- Works like the Identity applicative, but the "phase" type must alternate.
-- P1,P2,P1 and P2,P1,P2 work as usual, but P1,P1 results in a type error as
-- `MIdentityResult P1 P1` does not reduce to a `Proxy r` type.

data Phase = P1 | P2 deriving (Show)
newtype MIdentity (p :: Phase) a = MIdentity { getMIdentity :: a }
   deriving (Show,Eq,Ord)

instance MultiFunctor MIdentity where
   mfmap f (MIdentity x) = MIdentity (f x)

type family MIdentityResult p q :: * where
   MIdentityResult P1 P2 = Proxy P2
   MIdentityResult P2 P1 = Proxy P1

instance MultiApplicative MIdentity where
   type MAppResult MIdentity p q = MIdentityResult p q

   mpure = MIdentity
   MIdentity f <!*> MIdentity x = MIdentity (f x)

phase1 = mpure :: a -> MIdentity 'P1 a
phase2 = mpure :: a -> MIdentity 'P2 a

instance MultiMonad MIdentity where
   MIdentity a !>>= f = MIdentity $ getMIdentity $ f a

-- Examples:
--   (++) <!$> phase1 "hello" !*> phase2 " world" ===> MIdentity "hello world"
--   (++) <!$> phase2 "good"  !*> phase1 " day"   ===> MIdentity "good day"
--   (++) <!$> phase1 "wrong" !*> phase1 " phase" ===> <type error>