    {-# LANGUAGE UnboxedTuples #-}

    -- | Extra functions for "Control.Monad".
    --   These functions provide looping, list operations and booleans.
    -- If you need a wider selection of monad loops and list generalisations,
    -- see <http://hackage.haskell.org/package/monad-loops monad-loops>.
    module Control.Monad.Extra(
        module Control.Monad,
        whenJust, whenJustM,
        unit,
        -- * Loops
        loopM, whileM,
        -- * Lists
        partitionM, concatMapM, mapMaybeM, findM, firstJustM,
        -- * Booleans
        whenM, unlessM, ifM, notM, (||^), (&&^), orM, andM, anyM, allM
        ) where

    import Control.Monad
    import Control.Applicative
    import Data.Maybe
    import Prelude

    -- General utilities

    -- | Perform some operation on 'Just', given the field inside the 'Just'.
    --
    -- > whenJust Nothing  print == return ()
    -- > whenJust (Just 1) print == print 1
    whenJust :: Applicative m => Maybe a -> (a -> m ()) -> m ()
    whenJust mg f = maybe (pure ()) f mg

    -- | Like 'whenJust', but where the test can be monadic.
    whenJustM :: Monad m => m (Maybe a) -> (a -> m ()) -> m ()
    whenJustM mg f = maybe (return ()) f =<< mg

    -- | The identity function which requires the inner argument to be @()@. Useful for functions
    --   with overloaded return types.
    --
    -- > \(x :: Maybe ()) -> unit x == x
    unit :: m () -> m ()
    unit = id

    -- Data.List for Monad

    -- | A version of 'partition' that works with a monadic predicate.
    --
    -- > partitionM (Just . even) [1,2,3] == Just ([2], [1,3])
    -- > partitionM (const Nothing) [1,2,3] == Nothing
    partitionM :: Monad m => (a -> m Bool) -> [a] -> m ([a], [a])
    partitionM f [] = return ([], [])
    partitionM f (x:xs) = do
        res <- f x
        (as,bs) <- partitionM f xs
        return ([x | res]++as, [x | not res]++bs)

    -- | A version of 'concatMap' that works with a monadic predicate.
    concatMapM :: Monad m => (a -> m [b]) -> [a] -> m [b]
    {-# INLINE concatMapM #-}
    concatMapM op = foldr f (return [])
        where f x xs = do x <- op x; if null x then xs else do xs <- xs; return $ x++xs

    -- | A version of 'mapMaybe' that works with a monadic predicate.
    mapMaybeM :: Monad m => (a -> m (Maybe b)) -> [a] -> m [b]
    {-# INLINE mapMaybeM #-}
    mapMaybeM op = foldr f (return [])
        where f x xs = do x <- op x; case x of Nothing -> xs; Just x -> do xs <- xs; return $ x:xs

    -- Looping

    -- | A looping operation, 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 act x = do
        res <- act x
        case res of
            Left x -> loopM act x
            Right v -> return v

    -- | 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 act = do
        b <- act
        when b $ whileM act

    -- Booleans

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

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

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

    -- | Like 'not', but where the test can be monadic.
    notM :: Functor m => m Bool -> m Bool
    notM = fmap not

    -- | The lazy '||' operator lifted to a monad. If the first
    --   argument evaluates to 'True' the second argument will not
    --   be evaluated.
    --
    -- > Just True  ||^ undefined  == Just True
    -- > Just False ||^ Just True  == Just True
    -- > Just False ||^ Just False == Just False
    (||^) :: Monad m => m Bool -> m Bool -> m Bool
    (||^) a b = ifM a (return True) b

    -- | The lazy '&&' operator lifted to a monad. If the first
    --   argument evaluates to 'False' the second argument will not
    --   be evaluated.
    --
    -- > Just False &&^ undefined  == Just False
    -- > Just True  &&^ Just True  == Just True
    -- > Just True  &&^ Just False == Just False
    (&&^) :: Monad m => m Bool -> m Bool -> m Bool
    (&&^) a b = ifM a b (return False)

    -- | 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 p [] = return False
    anyM p (x:xs) = ifM (p x) (return True) (anyM p 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 p [] = return True
    allM p (x:xs) = ifM (p x) (allM p xs) (return 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 = anyM 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 = allM id

    -- Searching

    -- | Like 'find', but where the test can be monadic.
    --
    -- > findM (Just . isUpper) "teST"             == Just (Just 'S')
    -- > findM (Just . isUpper) "test"             == Just Nothing
    -- > findM (Just . const True) ["x",undefined] == Just (Just "x")
    findM :: Monad m => (a -> m Bool) -> [a] -> m (Maybe a)
    findM p [] = return Nothing
    findM p (x:xs) = ifM (p x) (return $ Just x) (findM p xs)

    -- | Like 'findM', but also allows you to compute some additional information in the predicate.
    firstJustM :: Monad m => (a -> m (Maybe b)) -> [a] -> m (Maybe b)
    firstJustM p [] = return Nothing
    firstJustM p (x:xs) = maybe (firstJustM p xs) (return . Just) =<< p x