{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE StandaloneKindSignatures #-}

module GHC.Utils.Monad.Codensity
  ( Codensity(..), toCodensity, fromCodensity )
  where

import Data.Kind ( Type )

import GHC.Prelude
import GHC.Exts ( oneShot )

import Control.Monad.Fix
import Control.Monad.IO.Class
import Control.Monad.Trans.Class
import Control.Concurrent.MVar ( newEmptyMVar, readMVar, putMVar )
import Control.Exception
import GHC.IO.Exception
import GHC.IO.Unsafe ( unsafeDupableInterleaveIO )

--------------------------------------------------------------------------------

type Codensity :: (Type -> Type) -> Type -> Type
newtype Codensity m a = Codensity { forall (m :: * -> *) a.
Codensity m a -> forall r. (a -> m r) -> m r
runCodensity :: forall r. (a -> m r) -> m r }
instance Functor (Codensity k) where
  fmap :: forall a b. (a -> b) -> Codensity k a -> Codensity k b
fmap a -> b
f (Codensity forall r. (a -> k r) -> k r
m) = (forall r. (b -> k r) -> k r) -> Codensity k b
forall (m :: * -> *) a.
(forall r. (a -> m r) -> m r) -> Codensity m a
Codensity ((forall r. (b -> k r) -> k r) -> Codensity k b)
-> (forall r. (b -> k r) -> k r) -> Codensity k b
forall a b. (a -> b) -> a -> b
$ ((b -> k r) -> k r) -> (b -> k r) -> k r
forall a b. (a -> b) -> a -> b
oneShot (\b -> k r
k -> (a -> k r) -> k r
forall r. (a -> k r) -> k r
m ((a -> k r) -> k r) -> (a -> k r) -> k r
forall a b. (a -> b) -> a -> b
$ (a -> k r) -> a -> k r
forall a b. (a -> b) -> a -> b
oneShot (\a
x -> b -> k r
k (b -> k r) -> b -> k r
forall a b. (a -> b) -> a -> b
$ a -> b
f a
x))
  {-# INLINE fmap #-}
instance Applicative (Codensity f) where
  pure :: forall a. a -> Codensity f a
pure a
x = (forall r. (a -> f r) -> f r) -> Codensity f a
forall (m :: * -> *) a.
(forall r. (a -> m r) -> m r) -> Codensity m a
Codensity ((forall r. (a -> f r) -> f r) -> Codensity f a)
-> (forall r. (a -> f r) -> f r) -> Codensity f a
forall a b. (a -> b) -> a -> b
$ ((a -> f r) -> f r) -> (a -> f r) -> f r
forall a b. (a -> b) -> a -> b
oneShot (\a -> f r
k -> a -> f r
k a
x)
  {-# INLINE pure #-}
  Codensity forall r. ((a -> b) -> f r) -> f r
f <*> :: forall a b. Codensity f (a -> b) -> Codensity f a -> Codensity f b
<*> Codensity forall r. (a -> f r) -> f r
g =
    (forall r. (b -> f r) -> f r) -> Codensity f b
forall (m :: * -> *) a.
(forall r. (a -> m r) -> m r) -> Codensity m a
Codensity ((forall r. (b -> f r) -> f r) -> Codensity f b)
-> (forall r. (b -> f r) -> f r) -> Codensity f b
forall a b. (a -> b) -> a -> b
$ ((b -> f r) -> f r) -> (b -> f r) -> f r
forall a b. (a -> b) -> a -> b
oneShot (\b -> f r
bfr -> ((a -> b) -> f r) -> f r
forall r. ((a -> b) -> f r) -> f r
f (((a -> b) -> f r) -> f r) -> ((a -> b) -> f r) -> f r
forall a b. (a -> b) -> a -> b
$ ((a -> b) -> f r) -> (a -> b) -> f r
forall a b. (a -> b) -> a -> b
oneShot (\a -> b
ab -> (a -> f r) -> f r
forall r. (a -> f r) -> f r
g ((a -> f r) -> f r) -> (a -> f r) -> f r
forall a b. (a -> b) -> a -> b
$ (a -> f r) -> a -> f r
forall a b. (a -> b) -> a -> b
oneShot (\a
x -> b -> f r
bfr (a -> b
ab a
x))))
  {-# INLINE (<*>) #-}
instance Monad (Codensity f) where
  return :: forall a. a -> Codensity f a
return = a -> Codensity f a
forall a. a -> Codensity f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  {-# INLINE return #-}
  Codensity f a
m >>= :: forall a b. Codensity f a -> (a -> Codensity f b) -> Codensity f b
>>= a -> Codensity f b
k =
    (forall r. (b -> f r) -> f r) -> Codensity f b
forall (m :: * -> *) a.
(forall r. (a -> m r) -> m r) -> Codensity m a
Codensity ((forall r. (b -> f r) -> f r) -> Codensity f b)
-> (forall r. (b -> f r) -> f r) -> Codensity f b
forall a b. (a -> b) -> a -> b
$ ((b -> f r) -> f r) -> (b -> f r) -> f r
forall a b. (a -> b) -> a -> b
oneShot (\b -> f r
c -> Codensity f a -> forall r. (a -> f r) -> f r
forall (m :: * -> *) a.
Codensity m a -> forall r. (a -> m r) -> m r
runCodensity Codensity f a
m ((a -> f r) -> f r) -> (a -> f r) -> f r
forall a b. (a -> b) -> a -> b
$ (a -> f r) -> a -> f r
forall a b. (a -> b) -> a -> b
oneShot (\a
a -> Codensity f b -> forall r. (b -> f r) -> f r
forall (m :: * -> *) a.
Codensity m a -> forall r. (a -> m r) -> m r
runCodensity (a -> Codensity f b
k a
a) b -> f r
c))
  {-# INLINE (>>=) #-}
instance MonadTrans Codensity where
  lift :: forall (m :: * -> *) a. Monad m => m a -> Codensity m a
lift m a
m = (forall r. (a -> m r) -> m r) -> Codensity m a
forall (m :: * -> *) a.
(forall r. (a -> m r) -> m r) -> Codensity m a
Codensity ((forall r. (a -> m r) -> m r) -> Codensity m a)
-> (forall r. (a -> m r) -> m r) -> Codensity m a
forall a b. (a -> b) -> a -> b
$ ((a -> m r) -> m r) -> (a -> m r) -> m r
forall a b. (a -> b) -> a -> b
oneShot (m a
m m a -> (a -> m r) -> m r
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=)
  {-# INLINE lift #-}
instance MonadIO m => MonadIO (Codensity m) where
  liftIO :: forall a. IO a -> Codensity m a
liftIO = m a -> Codensity m a
forall (m :: * -> *) a. Monad m => m a -> Codensity m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m a -> Codensity m a) -> (IO a -> m a) -> IO a -> Codensity m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO a -> m a
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO
  {-# INLINE liftIO #-}
instance MonadIO m => MonadFix (Codensity m) where
  mfix :: forall a. (a -> Codensity m a) -> Codensity m a
mfix a -> Codensity m a
f = (forall r. (a -> m r) -> m r) -> Codensity m a
forall (m :: * -> *) a.
(forall r. (a -> m r) -> m r) -> Codensity m a
Codensity ((forall r. (a -> m r) -> m r) -> Codensity m a)
-> (forall r. (a -> m r) -> m r) -> Codensity m a
forall a b. (a -> b) -> a -> b
$ ((a -> m r) -> m r) -> (a -> m r) -> m r
forall a b. (a -> b) -> a -> b
oneShot (((a -> m r) -> m r) -> (a -> m r) -> m r)
-> ((a -> m r) -> m r) -> (a -> m r) -> m r
forall a b. (a -> b) -> a -> b
$ \ a -> m r
k -> do
    promise <- IO (MVar a) -> m (MVar a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MVar a) -> m (MVar a)) -> IO (MVar a) -> m (MVar a)
forall a b. (a -> b) -> a -> b
$ IO (MVar a)
forall a. IO (MVar a)
newEmptyMVar
    ans     <- liftIO $ unsafeDupableInterleaveIO
                      $ readMVar promise
                          `catch`
                        (\ BlockedIndefinitelyOnMVar
BlockedIndefinitelyOnMVar -> FixIOException -> IO a
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO FixIOException
FixIOException)
    runCodensity (f ans) $ oneShot $ \ a
a -> do
      IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ MVar a -> a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar a
promise a
a
      a -> m r
k a
a
  {-# INLINE mfix #-}

toCodensity :: Monad m => m a -> Codensity m a
toCodensity :: forall (m :: * -> *) a. Monad m => m a -> Codensity m a
toCodensity m a
m = (forall r. (a -> m r) -> m r) -> Codensity m a
forall (m :: * -> *) a.
(forall r. (a -> m r) -> m r) -> Codensity m a
Codensity ((forall r. (a -> m r) -> m r) -> Codensity m a)
-> (forall r. (a -> m r) -> m r) -> Codensity m a
forall a b. (a -> b) -> a -> b
$ ((a -> m r) -> m r) -> (a -> m r) -> m r
forall a b. (a -> b) -> a -> b
oneShot (m a
m m a -> (a -> m r) -> m r
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=)
{-# INLINE toCodensity #-}

fromCodensity :: Monad m => Codensity m a -> m a
fromCodensity :: forall (m :: * -> *) a. Monad m => Codensity m a -> m a
fromCodensity Codensity m a
c = Codensity m a -> forall r. (a -> m r) -> m r
forall (m :: * -> *) a.
Codensity m a -> forall r. (a -> m r) -> m r
runCodensity Codensity m a
c a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return
{-# INLINE fromCodensity #-}