{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE LambdaCase #-}

-- | Utils for Monads
module Haskus.Utils.Monad
   ( MonadInIO (..)
   , module Control.Monad
   , module Control.Monad.IO.Class
   , module Control.Monad.Trans.Class
   , whileM
   , loop
   , loopM
   , whenM
   , unlessM
   , ifM
   , notM
   , anyM
   , allM
   , orM
   , andM
   )
where

import Control.Monad.IO.Class
import Control.Monad.Trans.Class
import Control.Monad
import Control.Monad.State

class MonadIO m => MonadInIO m where
   -- | Lift with*-like functions into IO (alloca, etc.)
   liftWith :: (forall c. (a -> IO c) -> IO c) -> (a -> m b) -> m b

   -- | Lift with*-like functions into IO (alloca, etc.)
   liftWith2 :: (forall c. (a -> b -> IO c) -> IO c) -> (a -> b -> m e) -> m e

instance MonadInIO IO where
   {-# INLINABLE liftWith #-}
   liftWith :: (forall c. (a -> IO c) -> IO c) -> (a -> IO b) -> IO b
liftWith forall c. (a -> IO c) -> IO c
wth a -> IO b
f = (a -> IO b) -> IO b
forall c. (a -> IO c) -> IO c
wth a -> IO b
f

   {-# INLINABLE liftWith2 #-}
   liftWith2 :: (forall c. (a -> b -> IO c) -> IO c) -> (a -> b -> IO e) -> IO e
liftWith2 forall c. (a -> b -> IO c) -> IO c
wth a -> b -> IO e
f = (a -> b -> IO e) -> IO e
forall c. (a -> b -> IO c) -> IO c
wth a -> b -> IO e
f

instance MonadInIO m => MonadInIO (StateT s m) where
   {-# INLINABLE liftWith #-}
   liftWith :: (forall c. (a -> IO c) -> IO c)
-> (a -> StateT s m b) -> StateT s m b
liftWith forall c. (a -> IO c) -> IO c
wth a -> StateT s m b
f =
      (s -> m (b, s)) -> StateT s m b
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT ((s -> m (b, s)) -> StateT s m b)
-> (s -> m (b, s)) -> StateT s m b
forall a b. (a -> b) -> a -> b
$ \s
s -> do
         (forall c. (a -> IO c) -> IO c) -> (a -> m (b, s)) -> m (b, s)
forall (m :: * -> *) a b.
MonadInIO m =>
(forall c. (a -> IO c) -> IO c) -> (a -> m b) -> m b
liftWith forall c. (a -> IO c) -> IO c
wth (\a
a -> StateT s m b -> s -> m (b, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (a -> StateT s m b
f a
a) s
s)

   {-# INLINABLE liftWith2 #-}
   liftWith2 :: (forall c. (a -> b -> IO c) -> IO c)
-> (a -> b -> StateT s m e) -> StateT s m e
liftWith2 forall c. (a -> b -> IO c) -> IO c
wth a -> b -> StateT s m e
f =
      (s -> m (e, s)) -> StateT s m e
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT ((s -> m (e, s)) -> StateT s m e)
-> (s -> m (e, s)) -> StateT s m e
forall a b. (a -> b) -> a -> b
$ \s
s ->
         (forall c. (a -> b -> IO c) -> IO c)
-> (a -> b -> m (e, s)) -> m (e, s)
forall (m :: * -> *) a b e.
MonadInIO m =>
(forall c. (a -> b -> IO c) -> IO c) -> (a -> b -> m e) -> m e
liftWith2 forall c. (a -> b -> IO c) -> IO c
wth (\a
a b
b -> StateT s m e -> s -> m (e, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (a -> b -> StateT s m e
f a
a b
b) s
s)

-- | Keep running an operation until it becomes 'False'. As an example:
--
-- @
-- whileM $ do sleep 0.1; notM $ doesFileExist "foo.txt"
-- readFile "foo.txt"
-- @
--
--   If you need some state persisted between each test, use 'loopM'.
whileM :: Monad m => m Bool -> m ()
whileM :: m Bool -> m ()
whileM m Bool
act = do
    Bool
b <- m Bool
act
    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
b (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ m Bool -> m ()
forall (m :: * -> *). Monad m => m Bool -> m ()
whileM m Bool
act

-- Looping

-- | A looping operation, where the predicate returns 'Left' as a seed for the next loop
--   or 'Right' to abort the loop.
--
-- > loop (\x -> if x < 10 then Left $ x * 2 else Right $ show x) 1 == "16"
loop :: (a -> Either a b) -> a -> b
loop :: (a -> Either a b) -> a -> b
loop a -> Either a b
act a
x = case a -> Either a b
act a
x of
    Left a
x' -> (a -> Either a b) -> a -> b
forall a b. (a -> Either a b) -> a -> b
loop a -> Either a b
act a
x'
    Right b
v -> b
v

-- | A monadic version of 'loop', where the predicate returns 'Left' as a seed for the next loop
--   or 'Right' to abort the loop.
loopM :: Monad m => (a -> m (Either a b)) -> a -> m b
loopM :: (a -> m (Either a b)) -> a -> m b
loopM a -> m (Either a b)
act a
x = a -> m (Either a b)
act a
x m (Either a b) -> (Either a b -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
   Left a
x' -> (a -> m (Either a b)) -> a -> m b
forall (m :: * -> *) a b.
Monad m =>
(a -> m (Either a b)) -> a -> m b
loopM a -> m (Either a b)
act a
x'
   Right b
v -> b -> m b
forall (m :: * -> *) a. Monad m => a -> m a
return b
v


-- | Like 'when', but where the test can be monadic.
whenM :: Monad m => m Bool -> m () -> m ()
whenM :: m Bool -> m () -> m ()
whenM m Bool
b m ()
t = m Bool -> m () -> m () -> m ()
forall (m :: * -> *) a. Monad m => m Bool -> m a -> m a -> m a
ifM m Bool
b m ()
t (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())

-- | Like 'unless', but where the test can be monadic.
unlessM :: Monad m => m Bool -> m () -> m ()
unlessM :: m Bool -> m () -> m ()
unlessM m Bool
b m ()
f = m Bool -> m () -> m () -> m ()
forall (m :: * -> *) a. Monad m => m Bool -> m a -> m a -> m a
ifM m Bool
b (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()) m ()
f

-- | Like @if@, but where the test can be monadic.
ifM :: Monad m => m Bool -> m a -> m a -> m a
ifM :: m Bool -> m a -> m a -> m a
ifM m Bool
mb m a
t m a
f = do
   Bool
b <- m Bool
mb
   if Bool
b then m a
t else m a
f

-- | Like 'not', but where the test can be monadic.
notM :: Functor m => m Bool -> m Bool
notM :: m Bool -> m Bool
notM = (Bool -> Bool) -> m Bool -> m Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Bool -> Bool
not

-- | A version of 'any' lifted to a monad. Retains the short-circuiting behaviour.
--
-- > anyM Just [False,True ,undefined] == Just True
-- > anyM Just [False,False,undefined] == undefined
-- > \(f :: Int -> Maybe Bool) xs -> anyM f xs == orM (map f xs)
anyM :: Monad m => (a -> m Bool) -> [a] -> m Bool
anyM :: (a -> m Bool) -> [a] -> m Bool
anyM a -> m Bool
_ []     = Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
anyM a -> m Bool
p (a
x:[a]
xs) = m Bool -> m Bool -> m Bool -> m Bool
forall (m :: * -> *) a. Monad m => m Bool -> m a -> m a -> m a
ifM (a -> m Bool
p a
x) (Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True) ((a -> m Bool) -> [a] -> m Bool
forall (m :: * -> *) a. Monad m => (a -> m Bool) -> [a] -> m Bool
anyM a -> m Bool
p [a]
xs)

-- | A version of 'all' lifted to a monad. Retains the short-circuiting behaviour.
--
-- > allM Just [True,False,undefined] == Just False
-- > allM Just [True,True ,undefined] == undefined
-- > \(f :: Int -> Maybe Bool) xs -> anyM f xs == orM (map f xs)
allM :: Monad m => (a -> m Bool) -> [a] -> m Bool
allM :: (a -> m Bool) -> [a] -> m Bool
allM a -> m Bool
_ []     = Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
allM a -> m Bool
p (a
x:[a]
xs) = m Bool -> m Bool -> m Bool -> m Bool
forall (m :: * -> *) a. Monad m => m Bool -> m a -> m a -> m a
ifM (a -> m Bool
p a
x) ((a -> m Bool) -> [a] -> m Bool
forall (m :: * -> *) a. Monad m => (a -> m Bool) -> [a] -> m Bool
allM a -> m Bool
p [a]
xs) (Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False)

-- | A version of 'or' lifted to a monad. Retains the short-circuiting behaviour.
--
-- > orM [Just False,Just True ,undefined] == Just True
-- > orM [Just False,Just False,undefined] == undefined
-- > \xs -> Just (or xs) == orM (map Just xs)
orM :: Monad m => [m Bool] -> m Bool
orM :: [m Bool] -> m Bool
orM = (m Bool -> m Bool) -> [m Bool] -> m Bool
forall (m :: * -> *) a. Monad m => (a -> m Bool) -> [a] -> m Bool
anyM m Bool -> m Bool
forall a. a -> a
id

-- | A version of 'and' lifted to a monad. Retains the short-circuiting behaviour.
--
-- > andM [Just True,Just False,undefined] == Just False
-- > andM [Just True,Just True ,undefined] == undefined
-- > \xs -> Just (and xs) == andM (map Just xs)
andM :: Monad m => [m Bool] -> m Bool
andM :: [m Bool] -> m Bool
andM = (m Bool -> m Bool) -> [m Bool] -> m Bool
forall (m :: * -> *) a. Monad m => (a -> m Bool) -> [a] -> m Bool
allM m Bool -> m Bool
forall a. a -> a
id