|
| 1 | +{-# LANGUAGE BangPatterns #-} |
| 2 | +{-# LANGUAGE ScopedTypeVariables #-} |
| 3 | + |
| 4 | +module Cardano.CLI.IO.Lazy |
| 5 | + ( replicateM |
| 6 | + , sequenceM |
| 7 | + , traverseM |
| 8 | + , traverseStateM |
| 9 | + , forM |
| 10 | + , forStateM |
| 11 | + ) where |
| 12 | + |
| 13 | +import Control.Applicative (Applicative((<*>), pure), (<$>)) |
| 14 | +import Control.Monad (Monad(..)) |
| 15 | +import Control.Monad.IO.Unlift (MonadIO(liftIO), MonadUnliftIO, askUnliftIO, UnliftIO(unliftIO)) |
| 16 | +import Data.Function (($), (.), flip) |
| 17 | +import Data.Int (Int) |
| 18 | +import System.IO (IO) |
| 19 | + |
| 20 | +import qualified Data.List as L |
| 21 | +import qualified System.IO.Unsafe as IO |
| 22 | + |
| 23 | +replicateM :: MonadUnliftIO m => Int -> m a -> m [a] |
| 24 | +replicateM n f = sequenceM (L.replicate n f) |
| 25 | + |
| 26 | +sequenceM :: MonadUnliftIO m => [m a] -> m [a] |
| 27 | +sequenceM as = do |
| 28 | + f <- askUnliftIO |
| 29 | + liftIO $ sequenceIO (L.map (unliftIO f) as) |
| 30 | + |
| 31 | +-- | Traverses the function over the list and produces a lazy list in a |
| 32 | +-- monadic context. |
| 33 | +-- |
| 34 | +-- It is intended to be like the "standard" 'traverse' except |
| 35 | +-- that the list is generated lazily. |
| 36 | +traverseM :: MonadUnliftIO m => (a -> m b) -> [a] -> m [b] |
| 37 | +traverseM f as = do |
| 38 | + u <- askUnliftIO |
| 39 | + liftIO $ IO.unsafeInterleaveIO (go u as) |
| 40 | + where |
| 41 | + go _ [] = pure [] |
| 42 | + go !u (v:vs) = do |
| 43 | + !res <- unliftIO u (f v) |
| 44 | + rest <- IO.unsafeInterleaveIO (go u vs) |
| 45 | + pure (res:rest) |
| 46 | + |
| 47 | +traverseStateM :: forall m s a b. MonadUnliftIO m => s -> (s -> a -> m (s, b)) -> [a] -> m [b] |
| 48 | +traverseStateM s f as = do |
| 49 | + u <- askUnliftIO |
| 50 | + liftIO $ IO.unsafeInterleaveIO (go s u as) |
| 51 | + where |
| 52 | + go :: s -> UnliftIO m -> [a] -> IO [b] |
| 53 | + go _ _ [] = pure [] |
| 54 | + go t !u (v:vs) = do |
| 55 | + (t', !res) <- unliftIO u (f t v) |
| 56 | + rest <- IO.unsafeInterleaveIO (go t' u vs) |
| 57 | + pure (res:rest) |
| 58 | + |
| 59 | +forM :: MonadUnliftIO m => [a] -> (a -> m b) -> m [b] |
| 60 | +forM = flip traverseM |
| 61 | + |
| 62 | +forStateM :: MonadUnliftIO m => s -> [a] -> (s -> a -> m (s, b)) -> m [b] |
| 63 | +forStateM s as f = traverseStateM s f as |
| 64 | + |
| 65 | +-- Internal |
| 66 | +sequenceIO :: [IO a] -> IO [a] |
| 67 | +sequenceIO = IO.unsafeInterleaveIO . go |
| 68 | + where go :: [IO a] -> IO [a] |
| 69 | + go [] = return [] |
| 70 | + go (fa:fas) = (:) <$> fa <*> IO.unsafeInterleaveIO (go fas) |
0 commit comments