{-# LANGUAGE GADTs, NoMonomorphismRestriction #-}

module Control.RMonad.AsMonad (AsMonad, embed, unEmbed) where

import Control.Monad
import Control.RMonad hiding (fmap, (>>=), return, fail, mzero, mplus)
import qualified Control.RMonad as RM (fmap, (>>=), return, fail, mzero, mplus)

-- |Turn a restricted monad into a normal monad
embed :: Suitable m a => m a -> AsMonad m a
embed = Embed

-- |A wrapper type that can be used to embed an arbitrary restricted
-- monad as a normal monad.
--
-- Note that in general there can be an efficiency loss in using
-- 'AsMonad' where the underlying restricted monad could have been used
-- instead. For example when using 'Set' with 'AsMonad', intermediate results
-- will not have duplicates removed. On the other hand, it becomes legal
-- to use intermediate results that do not have 'Ord' instances.
data AsMonad m a where
   Embed :: Suitable m a => m a -> AsMonad m a
   FMap :: RFunctor m => (a -> b) -> AsMonad m a -> AsMonad m b
   Return :: RMonad m => a -> AsMonad m a
   Bind :: RMonad m => AsMonad m a -> (a -> AsMonad m b) -> AsMonad m b
   Fail :: RMonad m => String -> AsMonad m a
   MZero :: RMonadPlus m => AsMonad m a
   MPlus :: RMonadPlus m => AsMonad m a -> AsMonad m a -> AsMonad m a

instance RFunctor f => Functor (AsMonad f) where
   fmap = FMap

instance RMonad m => Monad (AsMonad m) where
   return = Return
   (>>=) = Bind
   fail = Fail

instance RMonadPlus m => MonadPlus (AsMonad m) where
   mzero = MZero
   mplus = MPlus

-- |Unwrap an 'AsMonad' value into the enclosed restricted monad
unEmbed :: Suitable m a => AsMonad m a -> m a
unEmbed (Embed m) = m
unEmbed (Return a) = RM.return a
unEmbed MZero = RM.mzero
unEmbed (MPlus m1 m2) = RM.mplus (unEmbed m1) (unEmbed m2)
unEmbed (Bind (Embed m) f) = (RM.>>=) m (\a -> unEmbed (f a))
unEmbed (Bind (Return a) f) = unEmbed (f a)
unEmbed (Bind (Bind m f) g) = unEmbed (Bind m (\x -> Bind (f x) g))
unEmbed (Bind (Fail s) _f) = unEmbed (Fail s)
unEmbed (Bind MZero _f) = unEmbed MZero
unEmbed (Bind (MPlus m1 m2) f) = unEmbed (MPlus (Bind m1 f) (Bind m2 f))
unEmbed (Bind (FMap f m1) g) = unEmbed (Bind m1 (g . f))
unEmbed (Fail s) = RM.fail s
unEmbed (FMap f (Embed m)) = RM.fmap f m
-- for all these cases, we could just translate the FMap into a Bind/Return
-- combination, since the RMonad instance will be available
unEmbed (FMap f (Return a)) = unEmbed (Return (f a))
unEmbed (FMap f (Bind m g)) = unEmbed (Bind m (FMap f . g))
unEmbed (FMap _f (Fail s)) = unEmbed (Fail s)
unEmbed (FMap _f MZero) = unEmbed MZero
unEmbed (FMap f (MPlus m1 m2)) = unEmbed (MPlus (FMap f m1) (FMap f m2))
unEmbed (FMap f (FMap g m)) = unEmbed (FMap (f . g) m)