{-# LANGUAGE BangPatterns #-}

-- | List conversions and utilities.

module System.IO.Streams.List
 ( -- * List conversions
   fromList
 , toList
 , outputToList
 , writeList

   -- * Utility
 , chunkList
 , chunkListWith
 , concatLists
 , listOutputStream
 ) where

------------------------------------------------------------------------------
import           Control.Concurrent.MVar    (modifyMVar, modifyMVar_, newMVar)
import           Control.Monad.IO.Class     (MonadIO (..))
import           Data.IORef                 (newIORef, readIORef, writeIORef)
import           Prelude                    hiding (read)
------------------------------------------------------------------------------
import           System.IO.Streams.Internal (InputStream, OutputStream, await, connect, fromConsumer, fromGenerator, makeInputStream, read, write, yield)


------------------------------------------------------------------------------
-- | Transforms a list into an 'InputStream' that produces no side effects.
--
-- @
-- ghci> is <- Streams.'fromList' [1, 2]
-- ghci> 'replicateM' 3 (Streams.'read' is)
-- [Just 1, Just 2, Nothing]
-- @
fromList :: [c] -> IO (InputStream c)
fromList :: [c] -> IO (InputStream c)
fromList [c]
inp = [c] -> IO (IORef [c])
forall a. a -> IO (IORef a)
newIORef [c]
inp IO (IORef [c])
-> (IORef [c] -> IO (InputStream c)) -> IO (InputStream c)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IO (Maybe c) -> IO (InputStream c)
forall a. IO (Maybe a) -> IO (InputStream a)
makeInputStream (IO (Maybe c) -> IO (InputStream c))
-> (IORef [c] -> IO (Maybe c)) -> IORef [c] -> IO (InputStream c)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IORef [c] -> IO (Maybe c)
forall a. IORef [a] -> IO (Maybe a)
f
  where
    f :: IORef [a] -> IO (Maybe a)
f IORef [a]
ref = IORef [a] -> IO [a]
forall a. IORef a -> IO a
readIORef IORef [a]
ref IO [a] -> ([a] -> IO (Maybe a)) -> IO (Maybe a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \[a]
l ->
            case [a]
l of
              []     -> Maybe a -> IO (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
              (a
x:[a]
xs) -> IORef [a] -> [a] -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef [a]
ref [a]
xs IO () -> IO (Maybe a) -> IO (Maybe a)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Maybe a -> IO (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Maybe a
forall a. a -> Maybe a
Just a
x)
{-# INLINE fromList #-}


------------------------------------------------------------------------------
-- | 'listOutputStream' returns an 'OutputStream' which stores values fed into
-- it and an action which flushes all stored values to a list.
--
-- The flush action resets the store.
--
-- Note that this function /will/ buffer any input sent to it on the heap.
-- Please don't use this unless you're sure that the amount of input provided
-- is bounded and will fit in memory without issues.
--
-- @
-- ghci> (os, flush) <- Streams.'listOutputStream' :: IO ('OutputStream' Int, IO [Int])
-- ghci> Streams.'writeList' [1, 2] os
-- ghci> flush
-- [1, 2]
-- ghci> Streams.'writeList' [3, 4] os
-- ghci> flush
-- [3, 4]
-- @
listOutputStream :: IO (OutputStream c, IO [c])
listOutputStream :: IO (OutputStream c, IO [c])
listOutputStream = do
    MVar ([c] -> [c])
r <- ([c] -> [c]) -> IO (MVar ([c] -> [c]))
forall a. a -> IO (MVar a)
newMVar [c] -> [c]
forall a. a -> a
id
    OutputStream c
c <- Consumer c () -> IO (OutputStream c)
forall r a. Consumer r a -> IO (OutputStream r)
fromConsumer (Consumer c () -> IO (OutputStream c))
-> Consumer c () -> IO (OutputStream c)
forall a b. (a -> b) -> a -> b
$ MVar ([c] -> [c]) -> Consumer c ()
forall a c. MVar ([a] -> c) -> Consumer a ()
consumer MVar ([c] -> [c])
r
    (OutputStream c, IO [c]) -> IO (OutputStream c, IO [c])
forall (m :: * -> *) a. Monad m => a -> m a
return (OutputStream c
c, MVar ([c] -> [c]) -> IO [c]
forall a. MVar ([a] -> [a]) -> IO [a]
flush MVar ([c] -> [c])
r)

  where
    consumer :: MVar ([a] -> c) -> Consumer a ()
consumer MVar ([a] -> c)
r = Consumer a ()
go
      where
        go :: Consumer a ()
go = Consumer a (Maybe a)
forall r. Consumer r (Maybe r)
await Consumer a (Maybe a) -> (Maybe a -> Consumer a ()) -> Consumer a ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Consumer a () -> (a -> Consumer a ()) -> Maybe a -> Consumer a ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (() -> Consumer a ()
forall (m :: * -> *) a. Monad m => a -> m a
return (() -> Consumer a ()) -> () -> Consumer a ()
forall a b. (a -> b) -> a -> b
$! ()) ((a -> Consumer a ()) -> Maybe a -> Consumer a ())
-> (a -> Consumer a ()) -> Maybe a -> Consumer a ()
forall a b. (a -> b) -> a -> b
$ \a
c -> do
                            IO () -> Consumer a ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> Consumer a ()) -> IO () -> Consumer a ()
forall a b. (a -> b) -> a -> b
$ MVar ([a] -> c) -> (([a] -> c) -> IO ([a] -> c)) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ MVar ([a] -> c)
r ((([a] -> c) -> IO ([a] -> c)) -> IO ())
-> (([a] -> c) -> IO ([a] -> c)) -> IO ()
forall a b. (a -> b) -> a -> b
$ \[a] -> c
dl -> ([a] -> c) -> IO ([a] -> c)
forall (m :: * -> *) a. Monad m => a -> m a
return ([a] -> c
dl ([a] -> c) -> ([a] -> [a]) -> [a] -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a
ca -> [a] -> [a]
forall a. a -> [a] -> [a]
:))
                            Consumer a ()
go)

    flush :: MVar ([a] -> [a]) -> IO [a]
flush MVar ([a] -> [a])
r = MVar ([a] -> [a])
-> (([a] -> [a]) -> IO ([a] -> [a], [a])) -> IO [a]
forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar MVar ([a] -> [a])
r ((([a] -> [a]) -> IO ([a] -> [a], [a])) -> IO [a])
-> (([a] -> [a]) -> IO ([a] -> [a], [a])) -> IO [a]
forall a b. (a -> b) -> a -> b
$ \[a] -> [a]
dl -> ([a] -> [a], [a]) -> IO ([a] -> [a], [a])
forall (m :: * -> *) a. Monad m => a -> m a
return ([a] -> [a]
forall a. a -> a
id, [a] -> [a]
dl [])
{-# INLINE listOutputStream #-}


------------------------------------------------------------------------------
-- | Drains an 'InputStream', converting it to a list. N.B. that this function
-- reads the entire 'InputStream' strictly into memory and as such is not
-- recommended for streaming applications or where the size of the input is not
-- bounded or known.
--
-- @
-- ghci> is <- Streams.'fromList' [1, 2]
-- ghci> Streams.'toList' is
-- [1, 2]
-- @
toList :: InputStream a -> IO [a]
toList :: InputStream a -> IO [a]
toList InputStream a
is = (OutputStream a -> IO ()) -> IO [a]
forall a b. (OutputStream a -> IO b) -> IO [a]
outputToList (InputStream a -> OutputStream a -> IO ()
forall a. InputStream a -> OutputStream a -> IO ()
connect InputStream a
is)
{-# INLINE toList #-}


------------------------------------------------------------------------------
-- | Given an IO action that requires an 'OutputStream', creates one and
-- captures all the output the action sends to it as a list.
--
-- Example:
--
-- @
-- ghci> import "Control.Applicative"
-- ghci> ('connect' <$> 'fromList' [\"a\", \"b\", \"c\"]) >>= 'outputToList'
-- [\"a\",\"b\",\"c\"]
-- @
outputToList :: (OutputStream a -> IO b) -> IO [a]
outputToList :: (OutputStream a -> IO b) -> IO [a]
outputToList OutputStream a -> IO b
f = do
    (OutputStream a
os, IO [a]
getList) <- IO (OutputStream a, IO [a])
forall c. IO (OutputStream c, IO [c])
listOutputStream
    b
_ <- OutputStream a -> IO b
f OutputStream a
os
    IO [a]
getList
{-# INLINE outputToList #-}


------------------------------------------------------------------------------
-- | Feeds a list to an 'OutputStream'. Does /not/ write an end-of-stream to
-- the stream.
--
-- @
-- ghci> os \<- Streams.'unlines' Streams.'System.IO.Streams.stdout' >>= Streams.'System.IO.Streams.contramap' (S.pack . show) :: IO ('OutputStream' Int)
-- ghci> Streams.'writeList' [1, 2] os
-- 1
-- 2
-- ghci> Streams.'writeList' [3, 4] os
-- 3
-- 4
-- @
writeList :: [a] -> OutputStream a -> IO ()
writeList :: [a] -> OutputStream a -> IO ()
writeList [a]
xs OutputStream a
os = (a -> IO ()) -> [a] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((Maybe a -> OutputStream a -> IO ())
-> OutputStream a -> Maybe a -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip Maybe a -> OutputStream a -> IO ()
forall a. Maybe a -> OutputStream a -> IO ()
write OutputStream a
os (Maybe a -> IO ()) -> (a -> Maybe a) -> a -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Maybe a
forall a. a -> Maybe a
Just) [a]
xs
{-# INLINE writeList #-}


------------------------------------------------------------------------------
-- | Splits an input stream into chunks of at most size @n@.
--
-- Example:
--
-- @
-- ghci> 'fromList' [1..14::Int] >>= 'chunkList' 4 >>= 'toList'
-- [[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14]]
-- @
chunkList :: Int                   -- ^ chunk size
          -> InputStream a         -- ^ stream to process
          -> IO (InputStream [a])
chunkList :: Int -> InputStream a -> IO (InputStream [a])
chunkList Int
n InputStream a
input = if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
                      then [Char] -> IO (InputStream [a])
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO (InputStream [a])) -> [Char] -> IO (InputStream [a])
forall a b. (a -> b) -> a -> b
$ [Char]
"chunkList: bad size: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
n
                      else Generator [a] () -> IO (InputStream [a])
forall r a. Generator r a -> IO (InputStream r)
fromGenerator (Generator [a] () -> IO (InputStream [a]))
-> Generator [a] () -> IO (InputStream [a])
forall a b. (a -> b) -> a -> b
$ Int -> ([a] -> [a]) -> Generator [a] ()
go Int
n [a] -> [a]
forall a. a -> a
id
  where
    go :: Int -> ([a] -> [a]) -> Generator [a] ()
go !Int
k [a] -> [a]
dl | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0    = [a] -> Generator [a] ()
forall r. r -> Generator r ()
yield ([a] -> [a]
dl []) Generator [a] () -> Generator [a] () -> Generator [a] ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> ([a] -> [a]) -> Generator [a] ()
go Int
n [a] -> [a]
forall a. a -> a
id
             | Bool
otherwise = do
                   IO (Maybe a) -> Generator [a] (Maybe a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (InputStream a -> IO (Maybe a)
forall a. InputStream a -> IO (Maybe a)
read InputStream a
input) Generator [a] (Maybe a)
-> (Maybe a -> Generator [a] ()) -> Generator [a] ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Generator [a] ()
-> (a -> Generator [a] ()) -> Maybe a -> Generator [a] ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Generator [a] ()
finish a -> Generator [a] ()
chunk
      where
        finish :: Generator [a] ()
finish  = let l :: [a]
l = [a] -> [a]
dl []
                  in if [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
l then () -> Generator [a] ()
forall (m :: * -> *) a. Monad m => a -> m a
return (() -> Generator [a] ()) -> () -> Generator [a] ()
forall a b. (a -> b) -> a -> b
$! () else [a] -> Generator [a] ()
forall r. r -> Generator r ()
yield [a]
l
        chunk :: a -> Generator [a] ()
chunk a
x = Int -> ([a] -> [a]) -> Generator [a] ()
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ([a] -> [a]
dl ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:))


------------------------------------------------------------------------------
-- | Splits an input stream into chunks whenever @p elt count@ returns true.
--
-- Example:
--
-- @
-- ghci> 'fromList' [1..14::Int] >>= 'chunkListWith' (\x n -> n>=4) >>= 'toList'
-- [[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14]]
-- ghci> 'fromList' ['a'..'z'] >>= 'chunkListWith' (\x n -> n>=4 && x `elem` "aeiouy") >>= 'toList'
-- ["abcde","fghi","jklmno","pqrstu","vwxy","z"]
-- @
--
-- /Since: 1.3.3.0./
chunkListWith :: (a -> Int -> Bool)    -- ^ break predicate
              -> InputStream a         -- ^ stream to process
              -> IO (InputStream [a])
chunkListWith :: (a -> Int -> Bool) -> InputStream a -> IO (InputStream [a])
chunkListWith a -> Int -> Bool
p InputStream a
input =
  Generator [a] () -> IO (InputStream [a])
forall r a. Generator r a -> IO (InputStream r)
fromGenerator (Generator [a] () -> IO (InputStream [a]))
-> Generator [a] () -> IO (InputStream [a])
forall a b. (a -> b) -> a -> b
$ Maybe a -> Int -> ([a] -> [a]) -> Generator [a] ()
go Maybe a
forall a. Maybe a
Nothing Int
0 [a] -> [a]
forall a. a -> a
id
  where
    go :: Maybe a -> Int -> ([a] -> [a]) -> Generator [a] ()
go Maybe a
v !Int
k [a] -> [a]
dl
      | Just a
x <- Maybe a
v, a -> Int -> Bool
p a
x Int
k = [a] -> Generator [a] ()
forall r. r -> Generator r ()
yield ([a] -> [a]
dl []) Generator [a] () -> Generator [a] () -> Generator [a] ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Maybe a -> Int -> ([a] -> [a]) -> Generator [a] ()
go Maybe a
forall a. Maybe a
Nothing Int
0 [a] -> [a]
forall a. a -> a
id
      | Bool
otherwise = do
          IO (Maybe a) -> Generator [a] (Maybe a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (InputStream a -> IO (Maybe a)
forall a. InputStream a -> IO (Maybe a)
read InputStream a
input) Generator [a] (Maybe a)
-> (Maybe a -> Generator [a] ()) -> Generator [a] ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Generator [a] ()
-> (a -> Generator [a] ()) -> Maybe a -> Generator [a] ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Generator [a] ()
finish a -> Generator [a] ()
chunk
      where
        finish :: Generator [a] ()
finish =
          let l :: [a]
l = [a] -> [a]
dl []
          in if [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
l
               then () -> Generator [a] ()
forall (m :: * -> *) a. Monad m => a -> m a
return (() -> Generator [a] ()) -> () -> Generator [a] ()
forall a b. (a -> b) -> a -> b
$! ()
               else [a] -> Generator [a] ()
forall r. r -> Generator r ()
yield [a]
l
        chunk :: a -> Generator [a] ()
chunk a
x = Maybe a -> Int -> ([a] -> [a]) -> Generator [a] ()
go (a -> Maybe a
forall a. a -> Maybe a
Just a
x) (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ([a] -> [a]
dl ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
:))


------------------------------------------------------------------------------
-- | Given an input stream containing lists, produces a new input stream that
-- will yield the concatenation of these lists. See 'Prelude.concat'.
--
-- Example:
--
-- @
-- ghci> Streams.'fromList' [[1,2,3::Int], [4,5,6]] >>=
--       Streams.'concatLists' >>=
--       Streams.'toList'
-- [1,2,3,4,5,6]
-- @
concatLists :: InputStream [a] -> IO (InputStream a)
concatLists :: InputStream [a] -> IO (InputStream a)
concatLists InputStream [a]
input = Generator a () -> IO (InputStream a)
forall r a. Generator r a -> IO (InputStream r)
fromGenerator Generator a ()
go
  where
    go :: Generator a ()
go      = IO (Maybe [a]) -> Generator a (Maybe [a])
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (InputStream [a] -> IO (Maybe [a])
forall a. InputStream a -> IO (Maybe a)
read InputStream [a]
input) Generator a (Maybe [a])
-> (Maybe [a] -> Generator a ()) -> Generator a ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Generator a ()
-> ([a] -> Generator a ()) -> Maybe [a] -> Generator a ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (() -> Generator a ()
forall (m :: * -> *) a. Monad m => a -> m a
return (() -> Generator a ()) -> () -> Generator a ()
forall a b. (a -> b) -> a -> b
$! ()) [a] -> Generator a ()
chunk
    chunk :: [a] -> Generator a ()
chunk [a]
l = [Generator a ()] -> Generator a ()
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, Monad m) =>
t (m a) -> m ()
sequence_ ((a -> Generator a ()) -> [a] -> [Generator a ()]
forall a b. (a -> b) -> [a] -> [b]
map a -> Generator a ()
forall r. r -> Generator r ()
yield [a]
l) Generator a () -> Generator a () -> Generator a ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Generator a ()
go