Input.hs

--------------------------------------------------------------------
-- !
-- Module     : System.IO.Lazy.Input
-- Copyright  : (c) Nicolas Pouillard 2009
-- License    : BSD3
--
-- Maintainer : Nicolas Pouillard <nicolas.pouillard@gmail.com>
-- Stability  : provisional
-- Portability:
--
--------------------------------------------------------------------

module System.IO.Lazy.Input
(
  -- * Types
  LI,

  -- * Running
  run, -- :: (NFData sa) => LI sa -> IO sa
  run', -- :: (NFData sa) => LI (SIO a) -> IO sa
  runAsSIO, -- :: (NFData sa) => LI sa -> SIO sa
  runAsSIO', -- :: (NFData sa) => LI sa -> SIO sa

  -- * Basic inputs
  pureLI, -- :: a -> LI a
  hGetContents, -- :: Handle -> LI String
  getContents, -- :: LI String
  readFile, -- :: FilePath -> LI String
  getChanContents, -- :: Chan a -> LI [a]

  -- * Combining multiple inputs lazily

  -- ** Combining them in sequence
  append, -- :: (NFData sa) => LI [sa] -> LI [sa] -> LI [sa]
  concat, -- :: (NFData sa) => [LI [sa]] -> LI [sa]

  -- ** Combining them in parallel
  interleave, -- :: (NFData sa) => LI [sa] -> LI [sa] -> LI [sa]
  interleaveEither, -- :: (NFData sa, NFData sb) => LI [sa] -> LI [sb] -> LI [Either sa sb]
  -- interleaveHandles, -- :: Handle -> Handle -> LI [Either Char Char]
  zip, -- :: (NFData sa, NFData sb) => LI [sa] -> LI [sb] -> LI [(sa, sb)]
  zipWith, -- :: (NFData sa, NFData sb) => (sa -> sb -> c) -> LI [sa] -> LI [sb] -> LI [c]
  zipMaybesWith, -- :: (NFData sa, NFData sb) => (Maybe sa -> Maybe sb -> c) -> LI [sa] -> LI [sb] -> LI [c]
  zipHandles, -- :: Handle -> Handle -> LI [(Char, Char)]

  -- * Debugging
  trace -- :: String -> LI a -> LI a
)
where

import Prelude hiding (zip, zipWith, readFile, concat, sequence, getContents)
import qualified Data.List as L
import System.IO (Handle)
import qualified System.IO as IO
import System.IO.Unsafe (unsafeInterleaveIO)
import qualified System.IO.Strict as SIO
import qualified System.IO.Strict.Internals as SIO.Internals
import System.IO.Lazy.Input.Internals
import System.IO.Unsafe.GetContents (unsafeHGetContents)
import System.IO.Strict (SIO, return')
import Data.Function
import Control.Parallel.Strategies (NFData(..))
import Control.Concurrent.Chan (Chan)
import qualified Control.Concurrent.Chan as CH
import Control.Applicative
import Control.Monad hiding (sequence)

-- | Any pure data can lifted as lazy input.
pureLI :: a -> LI a
pureLI = LI . return . pure

-- | Extract the data from a lazy input, this is commonly
-- used to actually run the given process over lazy inputs.
-- As in all the functions that requires a 'NFData' instance
-- this means the result will forced using 'rnf'.
run :: NFData sa => LI sa -> IO sa
run = run' . fmap return'

-- | Pretty much as 'run' expect that one can use strict
-- /IO/s ("System.IO.Strict") to produce the final result.
run' :: NFData sa => LI (SIO sa) -> IO sa
run' (LI startX) = startX >>= finalize . fmap SIO.run

-- | Pretty much as 'run' but live in the 'SIO' monad instead of 'IO'.
runAsSIO :: NFData sa => LI sa -> SIO sa
runAsSIO = SIO.Internals.wrap1 run

-- | Pretty much as 'run'' but live in the 'SIO' monad instead of 'IO'.
runAsSIO' :: NFData sa => LI (SIO sa) -> SIO sa
runAsSIO' = SIO.Internals.wrap1 run'

-- | Returns the contents of the given handle lazily.
hGetContents :: Handle -> LI String
hGetContents h = unsafeHGetContents h `finallyLI` IO.hClose h

-- | Returns the contents of standard input lazily.
getContents :: LI String
getContents = hGetContents IO.stdin

-- | Add debugging messages using the given string.
-- This will returns the same lazy input but more verbose.
trace :: String -> LI a -> LI a
trace msg (LI startA) = LI $ do putStrLn $ "Starting " ++ msg
                                a `Finally` releaseA <- startA
                                putStrLn $ "Started " ++ msg
                                return $ Finally a $ do
                                  putStrLn $ "Stopping " ++ msg
                                  releaseA
                                  putStrLn $ "Stopped " ++ msg

-- | Sequence two lazy inputs that produces lists as one only
-- list. Note that the resource management is precise. As
-- soon as the beginning of the second input is required,
-- the resource of the first input is released and the
-- the second resource is acquired.
append :: NFData sa => LI [sa] -> LI [sa] -> LI [sa]
append (LI startA) (LI startB) = LI $ do
  xs `Finally` releaseA    <- startA
  ~(ys `Finally` releaseB) <- unsafeInterleaveIO $ releaseA >> startB
  return $ (rnfList xs ++ ys) `Finally` releaseB

-- | Same as 'append' but for a list of inputs.
concat :: (NFData sa) => [LI [sa]] -> LI [sa]
concat = foldr append (pureLI [])

-- | Takes two lazy inputs and returns a single interleaved lazy input.
-- Note that this function is left biased, this is always the left
-- canal that is read first. This function is rarely used directly
-- with file contents since it mixes the two contents, one generally
-- use some tagging to separate them back. Look at 'interleaveEither'
-- for such a function.
interleave :: (NFData sa) => LI [sa] -> LI [sa] -> LI [sa]
interleave (LI startA) (LI startB) = LI $ do
  xs0 `Finally` releaseA <- startA
  ys0 `Finally` releaseB <- startB
  lazyReleaseA <- unsafeInterleaveIO releaseA
  lazyReleaseB <- unsafeInterleaveIO releaseB
  let loopLeft  (x:xs) ys = rnf x `seq` (x : loopRight xs ys)
      loopLeft  []     ys = lazyReleaseA `seq` ys
      loopRight xs (y:ys) = rnf y `seq` (y : loopLeft xs ys)
      loopRight xs []     = lazyReleaseB `seq` xs
  return $ loopLeft xs0 ys0 `Finally` (lazyReleaseA `seq` lazyReleaseB `seq` return ())

-- | Like 'interleave' but it starts by tagging the left input by 'Left' and right
-- input by 'Right' leading to lazy input of 'Either's.
interleaveEither :: (NFData sa, NFData sb) => LI [sa] -> LI [sb] -> LI [Either sa sb]
interleaveEither a b = interleave (map Left <$> a) (map Right <$> b)
 
-- use select
{-
-- | This function is very close to a combination of 'interleaveEither' and two
-- 'hGetContents', however it will wait for the first input that is ready
-- to be read. So this one is not left biased.
interleaveHandles :: Handle -> Handle -> LI [Either Char Char]
interleaveHandles h1 h2 = LI $ loop >>= (`Finally` (IO.hClose h1 >> IO.hClose h2))
  where loop = unsafeInterleaveIO $
                   (do h1ready <- IO.hReady h1
                       if h1ready then readLeft else readRight
                   ) `catchEOF` (map Right <$> unsafeHGetContents h2)
        readRight = (do c    <- IO.hGetChar h2
                        rest <- loop
                        return (Right c : rest)
                    ) `catchEOF` (map Left <$> unsafeHGetContents h1)
        readLeft = do c    <- IO.hGetChar h1
                      rest <- loop
                      return (Left c : rest)
-}

-- | Combine two lazy inputs as a single lazy input of pairs.
--
-- Note that if one input list is short, excess elements of the longer list are discarded.
zip :: (NFData sa, NFData sb) => LI [sa] -> LI [sb] -> LI [(sa, sb)]
zip = zipWith (,)

-- | 'zipWith' generalize 'zip' with any combining function.
zipWith :: (NFData sa, NFData sb) => (sa -> sb -> c) -> LI [sa] -> LI [sb] -> LI [c]
zipWith f (LI startA) (LI startB) = LI $ do
  xs `Finally` releaseA <- startA
  ys `Finally` releaseB <- startB
  let f' x y = rnf x `seq` rnf y `seq` f x y
  return $ L.zipWith f' xs ys `Finally` (releaseA >> releaseB)

zipMaybesWith :: (NFData sa, NFData sb) => (Maybe sa -> Maybe sb -> c) -> LI [sa] -> LI [sb] -> LI [c]
zipMaybesWith f xs ys =
    map (uncurry f) . takeWhile someJust <$> zip (prolongate <$> xs) (prolongate <$> ys)
  where someJust (Nothing, Nothing) = False
        someJust          _         = True
        prolongate :: [a] -> [Maybe a]
        prolongate zs = map Just zs ++ repeat Nothing

-- | A shorthand for @\\h1 h2-> zip (hGetContents h1) (hGetContents h2)@.
zipHandles :: Handle -> Handle -> LI [(Char, Char)]
zipHandles = zip `on` hGetContents

-- | Like 'hGetContents' but it takes a 'FilePath'.
readFile :: FilePath -> LI String
readFile fp = LI $ do
  h <- IO.openFile fp IO.ReadMode
  x <- unsafeHGetContents h
  return $ x `Finally` IO.hClose h

-- |Return a lazy list representing the contents of the supplied
-- 'Chan', much like 'System.IO.hGetContents'.
getChanContents :: Chan a -> LI [a]
getChanContents ch = nonFinalized go
  where go = unsafeInterleaveIO $ do
                x  <- CH.readChan ch
                xs <- go
                return (x:xs)

{-
joinLISIO :: LI (SIO a) -> LI a
joinLISIO x = x !>>= id

runSIOinIO :: NFData a => LI (SIO a) -> IO a
runSIOinIO = runIO . joinLISIO
-}