{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Cardano.CLI.IO.Lazy
  ( replicateM
  , sequenceM
  , traverseM
  , traverseStateM
  , forM
  , forStateM
  ) where

import           Control.Monad.IO.Unlift (MonadIO (liftIO), MonadUnliftIO, UnliftIO (unliftIO),
                   askUnliftIO)

import qualified Data.List as L
import qualified System.IO.Unsafe as IO

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

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

-- | 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 :: forall (m :: * -> *) a b.
MonadUnliftIO m =>
(a -> m b) -> [a] -> m [b]
traverseM a -> m b
f [a]
as = do
  UnliftIO m
u <- forall (m :: * -> *). MonadUnliftIO m => m (UnliftIO m)
askUnliftIO
  forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> 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
_ [] = forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    go !UnliftIO m
u (a
v:[a]
vs) = do
      !b
res <- forall (m :: * -> *). UnliftIO m -> forall a. m a -> IO a
unliftIO UnliftIO m
u (a -> m b
f a
v)
      [b]
rest <- forall a. IO a -> IO a
IO.unsafeInterleaveIO (UnliftIO m -> [a] -> IO [b]
go UnliftIO m
u [a]
vs)
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (b
resforall a. a -> [a] -> [a]
:[b]
rest)

traverseStateM :: forall m s a b. MonadUnliftIO m => s -> (s -> a -> m (s, b)) -> [a] -> m [b]
traverseStateM :: forall (m :: * -> *) s a b.
MonadUnliftIO m =>
s -> (s -> a -> m (s, b)) -> [a] -> m [b]
traverseStateM s
s s -> a -> m (s, b)
f [a]
as = do
  UnliftIO m
u <- forall (m :: * -> *). MonadUnliftIO m => m (UnliftIO m)
askUnliftIO
  forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IO a -> IO a
IO.unsafeInterleaveIO (s -> UnliftIO m -> [a] -> IO [b]
go s
s UnliftIO m
u [a]
as)
  where
    go :: s -> UnliftIO m -> [a] -> IO [b]
    go :: s -> UnliftIO m -> [a] -> IO [b]
go s
_ UnliftIO m
_ [] = forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    go s
t !UnliftIO m
u (a
v:[a]
vs) = do
      (s
t', !b
res) <- forall (m :: * -> *). UnliftIO m -> forall a. m a -> IO a
unliftIO UnliftIO m
u (s -> a -> m (s, b)
f s
t a
v)
      [b]
rest <- forall a. IO a -> IO a
IO.unsafeInterleaveIO (s -> UnliftIO m -> [a] -> IO [b]
go s
t' UnliftIO m
u [a]
vs)
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (b
resforall a. a -> [a] -> [a]
:[b]
rest)

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

forStateM :: MonadUnliftIO m => s -> [a] -> (s -> a -> m (s, b)) -> m [b]
forStateM :: forall (m :: * -> *) s a b.
MonadUnliftIO m =>
s -> [a] -> (s -> a -> m (s, b)) -> m [b]
forStateM s
s [a]
as s -> a -> m (s, b)
f = forall (m :: * -> *) s a b.
MonadUnliftIO m =>
s -> (s -> a -> m (s, b)) -> [a] -> m [b]
traverseStateM s
s s -> a -> m (s, b)
f [a]
as

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