{-# LANGUAGE BangPatterns #-}

module HaskellWorks.Control.Monad.Lazy
  ( replicateM
  , sequenceM
  , unfoldrM
  , traverseM
  , forM
  , forceM
  ) where

import Control.DeepSeq
import Control.Monad ((<$!>))
import Control.Monad.IO.Unlift

import qualified System.IO.Unsafe as IO

replicateM :: MonadUnliftIO m => Int -> m a -> m [a]
replicateM :: Int -> m a -> m [a]
replicateM Int
n m a
f = [m a] -> m [a]
forall (m :: * -> *) a. MonadUnliftIO m => [m a] -> m [a]
sequenceM (Int -> m a -> [m a]
forall a. Int -> a -> [a]
replicate Int
n m a
f)

sequenceM :: MonadUnliftIO m => [m a] -> m [a]
sequenceM :: [m a] -> m [a]
sequenceM [m a]
as = do
  UnliftIO m
f <- m (UnliftIO m)
forall (m :: * -> *). MonadUnliftIO m => m (UnliftIO m)
askUnliftIO
  IO [a] -> m [a]
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO [a] -> m [a]) -> IO [a] -> m [a]
forall a b. (a -> b) -> a -> b
$ [IO a] -> IO [a]
forall a. [IO a] -> IO [a]
sequenceIO ((m a -> IO a) -> [m a] -> [IO a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (UnliftIO m -> forall a. m a -> IO a
forall (m :: * -> *). UnliftIO m -> forall a. m a -> IO a
unliftIO UnliftIO m
f) [m a]
as)

-- | Generates a lazy list of values that are produced by a given monadic function.
--
-- This function is intended to be like the "standard" 'unfoldrM' except
-- that the list is generated lazily.
unfoldrM :: MonadUnliftIO m => (b -> m (Maybe (a, b))) -> b -> m [a]
unfoldrM :: (b -> m (Maybe (a, b))) -> b -> m [a]
unfoldrM b -> m (Maybe (a, b))
f b
z = do
  UnliftIO m
u <- m (UnliftIO m)
forall (m :: * -> *). MonadUnliftIO m => m (UnliftIO m)
askUnliftIO
  IO [a] -> m [a]
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO [a] -> m [a]) -> IO [a] -> m [a]
forall a b. (a -> b) -> a -> b
$ IO [a] -> IO [a]
forall a. IO a -> IO a
IO.unsafeInterleaveIO (UnliftIO m -> b -> IO [a]
go UnliftIO m
u b
z)
  where
    go :: UnliftIO m -> b -> IO [a]
go !UnliftIO m
u !b
b = do
      Maybe (a, b)
m <- UnliftIO m -> m (Maybe (a, b)) -> IO (Maybe (a, b))
forall (m :: * -> *). UnliftIO m -> forall a. m a -> IO a
unliftIO UnliftIO m
u (b -> m (Maybe (a, b))
f b
b)
      case Maybe (a, b)
m of
        Maybe (a, b)
Nothing      -> [a] -> IO [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
        Just (!a
a, b
b') -> do
          [a]
rest <- IO [a] -> IO [a]
forall a. IO a -> IO a
IO.unsafeInterleaveIO (UnliftIO m -> b -> IO [a]
go UnliftIO m
u b
b')
          [a] -> IO [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
aa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
rest)

-- | Traverses the function over the list and produces a lazy list in a
-- monadic context.
--
-- It is intended to be like the "standard" 'traverse' except
-- that the list is generated lazily.
traverseM :: MonadUnliftIO m => (a -> m b) -> [a] -> m [b]
traverseM :: (a -> m b) -> [a] -> m [b]
traverseM a -> m b
f [a]
as = do
  UnliftIO m
u <- m (UnliftIO m)
forall (m :: * -> *). MonadUnliftIO m => m (UnliftIO m)
askUnliftIO
  IO [b] -> m [b]
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO [b] -> m [b]) -> IO [b] -> m [b]
forall a b. (a -> b) -> a -> b
$ IO [b] -> IO [b]
forall a. IO a -> IO a
IO.unsafeInterleaveIO (UnliftIO m -> [a] -> IO [b]
go UnliftIO m
u [a]
as)
  where
    go :: UnliftIO m -> [a] -> IO [b]
go UnliftIO m
_ [] = [b] -> IO [b]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    go !UnliftIO m
u (a
v:[a]
vs) = do
      !b
res <- UnliftIO m -> m b -> IO b
forall (m :: * -> *). UnliftIO m -> forall a. m a -> IO a
unliftIO UnliftIO m
u (a -> m b
f a
v)
      [b]
rest <- IO [b] -> IO [b]
forall a. IO a -> IO a
IO.unsafeInterleaveIO (UnliftIO m -> [a] -> IO [b]
go UnliftIO m
u [a]
vs)
      [b] -> IO [b]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b
resb -> [b] -> [b]
forall a. a -> [a] -> [a]
:[b]
rest)

forM :: MonadUnliftIO m => [a] -> (a -> m b) -> m [b]
forM :: [a] -> (a -> m b) -> m [b]
forM = ((a -> m b) -> [a] -> m [b]) -> [a] -> (a -> m b) -> m [b]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (a -> m b) -> [a] -> m [b]
forall (m :: * -> *) a b.
MonadUnliftIO m =>
(a -> m b) -> [a] -> m [b]
traverseM

forceM :: (Monad m, NFData a) => m a -> m a
forceM :: m a -> m a
forceM = (a -> a
forall a. NFData a => a -> a
force (a -> a) -> m a -> m a
forall (m :: * -> *) a b. Monad m => (a -> b) -> m a -> m b
<$!>)

-- Internal
sequenceIO :: [IO a] -> IO [a]
sequenceIO :: [IO a] -> IO [a]
sequenceIO = IO [a] -> IO [a]
forall a. IO a -> IO a
IO.unsafeInterleaveIO (IO [a] -> IO [a]) -> ([IO a] -> IO [a]) -> [IO a] -> IO [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [IO a] -> IO [a]
forall a. [IO a] -> IO [a]
go
  where go :: [IO a] -> IO [a]
        go :: [IO a] -> IO [a]
go []       = [a] -> IO [a]
forall (m :: * -> *) a. Monad m => a -> m a
return []
        go (IO a
fa:[IO a]
fas) = (:) (a -> [a] -> [a]) -> IO a -> IO ([a] -> [a])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO a
fa IO ([a] -> [a]) -> IO [a] -> IO [a]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> IO [a] -> IO [a]
forall a. IO a -> IO a
IO.unsafeInterleaveIO ([IO a] -> IO [a]
forall a. [IO a] -> IO [a]
go [IO a]
fas)