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

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

import Control.Monad.IO.Unlift
  ( MonadIO (liftIO)
  , MonadUnliftIO
  , UnliftIO (unliftIO)
  , askUnliftIO
  )
import Data.List qualified as L
import System.IO.Unsafe qualified 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 = [m a] -> m [a]
forall (m :: * -> *) a. MonadUnliftIO m => [m a] -> m [a]
sequenceM (Int -> m a -> [m a]
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
  f <- m (UnliftIO m)
forall (m :: * -> *). MonadUnliftIO m => m (UnliftIO m)
askUnliftIO
  liftIO $ sequenceIO (L.map (unliftIO f) as)

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

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 = s -> (s -> a -> m (s, b)) -> [a] -> m [b]
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 = 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 :: forall a. [IO a] -> IO [a]
go [] = [a] -> IO [a]
forall a. 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 a b. IO (a -> b) -> IO a -> IO b
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)