Monad.hs

{-# LANGUAGE RankNTypes, ExistentialQuantification, FlexibleInstances, UndecidableInstances, TypeFamilyDependencies, MultiParamTypeClasses, ScopedTypeVariables, FlexibleContexts #-}

module Futhask.Monad where

import System.IO.Unsafe
import Control.Monad.Base
import Control.Monad.Trans
import Control.Monad.Trans.Control
import Control.Monad.Identity

import Futhask.Context

-- *The Monad Transformer

-- | The transformer that passes around the Futhark context.
-- The @context@ parameter is library specific and prevents mixing of different generated libraries.
-- The @c@ parameter is skolem and prevents Futhark values from escaping the monad.
newtype FutT context c m a = FutT (context -> m a)

instance MonadTrans (FutT context c) where
    lift a = FutT (\_ -> a)
    {-# INLINEABLE lift #-}

instance Functor m => Functor (FutT context c m) where
    fmap f (FutT a) = FutT (fmap f.a)
    {-# INLINEABLE fmap #-}

instance Applicative m => Applicative (FutT context c m) where
    pure a = FutT (\_ -> pure a)
    (<*>) (FutT a) (FutT b) = FutT (\c -> a c <*> b c)
    {-# INLINEABLE pure #-}
    {-# INLINEABLE (<*>) #-}

instance Monad m => Monad (FutT context c m) where
    (>>=) (FutT a) f = FutT (\c -> a c >>= (\(FutT b) -> b c) . f)
    {-# INLINEABLE (>>=) #-}

instance (MonadBase b m) => MonadBase b (FutT context c m) where
    liftBase = liftBaseDefault
    {-# INLINEABLE liftBase #-}

instance MonadTransControl (FutT context c) where
    type StT (FutT context c) a = a
    liftWith a = FutT (\c -> a (\(FutT a') -> a' c))
    restoreT = lift
    {-# INLINEABLE liftWith #-}
    {-# INLINEABLE restoreT #-}

instance MonadBaseControl b m => MonadBaseControl b (FutT context c m) where
    type StM (FutT context c m) a = ComposeSt (FutT context c) m a
    liftBaseWith = defaultLiftBaseWith
    restoreM = defaultRestoreM
    {-# INLINEABLE liftBaseWith #-}
    {-# INLINEABLE restoreM #-}

-- **Predefined monads

-- | lazy without side effects, no manual memory management possible
type Fut context c = FutT context c Identity

-- | strict with side effects, manual memory management possible
type FutIO context c = FutT context c IO


-- *Running the Transformer

-- | Runs the transformer in a given context
runFutTIn :: (Context context) => context -> (forall c. FutT context c m a) -> m a
runFutTIn context (FutT a) = a context

-- | Runs the transformer in a context with a list of options
{-# NOINLINE runFutTWith #-}
runFutTWith :: (Context context) => [ContextOption] -> (forall c. FutT context c m a) -> m a
runFutTWith options a
    = unsafePerformIO
    $ mkContext options >>= \(c, _) -> return $ runFutTIn c a

-- | Runs the transformer with no context options
runFutT :: (Context context) => (forall c. FutT context c m a) -> m a
runFutT = runFutTWith []

-- | Runs the @Fut@ monad in a given context
runFutIn :: (Context context) => context -> (forall c. Fut context c a) -> a
runFutIn context a = runIdentity $ runFutTIn context $ a


-- | Runs the @Fut@ monad with in a context with a list of options
runFutWith :: (Context context) => [ContextOption] -> (forall c. Fut context c a) -> a
runFutWith options a = runIdentity $ runFutTWith options a

-- | Runs the @Fut@ monad with no context options
runFut :: (Context context) => (forall c. Fut context c a) -> a
runFut = runFutWith []

-- *Misc

-- | Lifts a futhark operation from the identity monad to any other monad
pureFut :: (Context context, Monad m) => Fut context c a -> FutT context c m a
pureFut (FutT a) = FutT (pure . runIdentity . a)

-- | Maps a function to the monad transformer
mapFutT :: (m a -> n b) -> FutT context c m a -> FutT context c n b
mapFutT f (FutT a) = FutT (f.a)

-- | Same as @mapFutT@ but for two arguments
map2FutT :: (m a -> n b -> k c) -> FutT context c' m a -> FutT context c' n b -> FutT context c' k c
map2FutT f (FutT a) (FutT b) = FutT (\c -> f (a c) (b c))

-- *Profiling

-- | Turns on profiling for the argument.
profile :: (Context context, MonadIO m) => FutT context c m a -> FutT context c m a
profile a = do
    FutT (\c -> liftIO $ inContext c unpauseProfiling)
    res <- a
    FutT (\c -> liftIO $ inContext c pauseProfiling)
    return res

-- *Unsafe Lifting

-- | Declares a @FutIO@ value as pure. The value must not consume any arguments or have side effects.
{-# INLINE unsafeFromFutIO #-}
unsafeFromFutIO :: (Context context, Monad m) => FutIO context c a -> FutT context c m a
unsafeFromFutIO (FutT a) = FutT (\ctx -> pure $! unsafePerformIO (a ctx))

-- | Wraps a @(context -> IO a)@ into the generic @FutT c m a@ . Useful for creating new abstractions from the @Raw@ interface functions. Operations should have no external side effects, or the function should be restricted to @MonadIO@.
{-# INLINE unsafeLiftFromIO #-}
unsafeLiftFromIO :: (Context context, Monad m) => (context -> IO a) -> FutT context c m a
unsafeLiftFromIO a = FutT (\ctx -> pure $! unsafePerformIO (a ctx))