--------------------------------------------------------------------
-- !
-- 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.Extra
(
  -- ** Various strictier input sequencing
  lift2MayForceFirst, -- :: (NFData sa) => (sa -> b -> c) -> LI sa -> LI b -> LI c
  lift2ForceFirst, -- :: (NFData sa) => (sa -> b -> c) -> LI sa -> LI b -> LI c
  lift2ForceSecond, -- :: (NFData sb) => (a -> sb -> c) -> LI a -> LI sb -> LI c
  lift2ForceBoth, -- :: (NFData sa, NFData sb) => (sa -> sb -> c) -> LI sa -> LI sb -> LI c
  (!>>=), -- :: (NFData sa) => LI sa -> (sa -> LI b) -> LI b
  (=<<!), -- :: (NFData sa) => (sa -> LI b) -> LI sa -> LI b
  ap', -- :: (NFData sa) => LI (sa -> b) -> LI sa -> LI b
  sequence, -- :: (NFData sa) => [LI sa] -> LI [sa]
)
where

import Prelude hiding (zip, zipWith, readFile, concat, sequence, getContents)
import System.IO.Unsafe (unsafeInterleaveIO)
import System.IO.Lazy.Input.Internals
import System.IO.Lazy.Input
import Control.Parallel.Strategies (NFData(..))
import Control.Applicative
import Control.Monad hiding (sequence)

{- | Lift a pure two arguments function, to a function over lazy inputs.

 Note that the only the first argument /may/ be deeply forced.
 In particular it is deeply forced if the function use its second argument.

 The strictness is here to enforce the evaluation order of reading inputs.

 Since too much strictness breaks the interest of lazy inputs, one provides
 more specific but lazier combinators like 'append', 'interleave', and 'zip'.
-}
lift2MayForceFirst :: (NFData sa) => (sa -> b -> c) -> LI sa -> LI b -> LI c
lift2MayForceFirst f (LI startA) (LI startB) = LI $ do
  x `Finally` releaseA <- startA
  y `Finally` releaseB <- startB
  lazyReleaseA <- unsafeInterleaveIO releaseA
  let r = f x (rnf x `seq` lazyReleaseA `seq` y)
  return $ r `Finally` (lazyReleaseA `seq` releaseB)

{- | Lift a pure two arguments function, to a function over lazy inputs.

 Note that the only the first argument is deeply forced before calling the function.

 The strictness is here to enforce the evaluation order of reading inputs.

 This lifting function can be generalized to n-ary functions, all arguments
 but the last one will be deeply forced.

 @
 liftN f mx1 mx2 ... mxN = mx1 !>>= \x1 -> mx2 !>>= \x2 -> ... f x1 x2 <$> mxN
 @
-}
lift2ForceFirst :: (NFData sa) => (sa -> b -> c) -> LI sa -> LI b -> LI c
lift2ForceFirst f mx my = mx !>>= \x -> f x <$> my

{- Like 'lift2ForceFirst' but only force the second argument.

 This lifting function can be generalized to n-ary functions, all arguments
 but the first one will be deeply forced.

 @
 liftN f mx1 mx2 ... mxN = f <$> mx1 `ap'` mx2 `ap'` ... `ap'` mxN
 @
-}
lift2ForceSecond :: (NFData sb) => (a -> sb -> c) -> LI a -> LI sb -> LI c
lift2ForceSecond f mx my = f <$> mx `ap'` my

{- | Lift a pure two arguments function, to a function over lazy inputs.

 Note that both arguments are deeply forced before calling the function.
 See 'lift2ForceFirst' and 'lift2ForceSecond' for lazier versions.

 This one can also be generalized to n-ary functions:

 @
 liftN f mx1 mx2 ... mxN = pureLI f `ap'` mx1 `ap'` mx2 `ap'` ... `ap'` mxN
 @
-}
lift2ForceBoth :: (NFData sa, NFData sb) => (sa -> sb -> c) -> LI sa -> LI sb -> LI c
lift2ForceBoth f mx my = pureLI f `ap'` mx `ap'` my

-- | Combines a function wrapped as a lazy input and an argument.
-- This is like 'ap' or '<*>' but stricter.
--
-- Note that since functions types are not member of 'NFData', this function
-- is the only one dealing with functions wrapped as lazy inputs.
--
-- However as with 'ap' or '<*>', this function generalize 'lift2ForceSecond', 'lift3Fst'...
--
-- Example:
-- @
-- lift3Fst f x y z = f <$> x `ap'` y `ap'` z
--
-- lift3strict f x y z = pureLI f `ap'` x `ap'` y `ap'` z
-- @
--
-- The 'ap'' function only deeply force the second argument, so in the case
-- of chaining, the arguments will be forced from left to right. Note that
-- if one starts the chain by lifting the function using 'pureLI', then all
-- the arguments will be forced. One can let one of the arguments lazy
-- by using note however that if one start the chain with '<$>' (same as
-- 'fmap' or 'liftM') then the first argument would not be forced, but one
-- can start with 'pureLI'
ap' :: (NFData sa) => LI (sa -> b) -> LI sa -> LI b
ap' (LI startF) marg = LI $ do
  f `Finally` releaseF <- startF
  arg <- run marg
  return $ f arg `Finally` releaseF
infixl 4 `ap'`

-- | Turns a list of lazy inputs as an input of list.
sequence :: (NFData sa) => [LI sa] -> LI [sa]
sequence = foldr (lift2ForceFirst (:)) (pureLI [])

-- | A kind of strict /bind/ over lazy inputs.
infixl 1 !>>=
(!>>=) :: (NFData sa) => LI sa -> (sa -> LI b) -> LI b
ma !>>= f = LI $ run ma >>= startLI . f

-- | Same as '!>>=' but with arguments flipped.
infixr 1 =<<!
(=<<!) :: (NFData sa) => (sa -> LI b) -> LI sa -> LI b
(=<<!) = flip (!>>=)