content: {-# LANGUAGE GADTs #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE LinearTypes #-}
  {-# LANGUAGE QualifiedDo #-}
  {-# LANGUAGE RebindableSyntax #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE NoImplicitPrelude #-}
  {-# OPTIONS_GHC -Wno-name-shadowing #-}
  {-# OPTIONS_HADDOCK hide #-}
  
  -- | This module provides all functions which produce a
  -- 'Stream (Of a) m r' from some given non-stream inputs.
  module Streaming.Linear.Internal.Produce
    ( -- * Constructing Finite 'Stream's
      yield,
      each',
      unfoldr,
      fromHandle,
      readFile,
      replicate,
      replicateM,
      replicateZip,
      untilRight,
  
      -- * Working with infinite 'Stream's
      stdinLnN,
      stdinLnUntil,
      stdinLnUntilM,
      stdinLnZip,
      readLnN,
      readLnUntil,
      readLnUntilM,
      readLnZip,
      iterateN,
      iterateZip,
      iterateMN,
      iterateMZip,
      cycleN,
      cycleZip,
      enumFromN,
      enumFromZip,
      enumFromThenN,
      enumFromThenZip,
    )
  where
  
  import qualified Control.Functor.Linear as Control
  import Data.Text (Text)
  import qualified Data.Text as Text
  import Data.Unrestricted.Linear
  import GHC.Stack
  import Prelude.Linear (($))
  import Streaming.Linear.Internal.Consume (effects)
  import Streaming.Linear.Internal.Process
  import Streaming.Linear.Internal.Type
  import qualified System.IO as System
  import System.IO.Linear
  import System.IO.Resource.Linear
  import Prelude
    ( Bool (..),
      Either (..),
      Enum,
      Eq (..),
      FilePath,
      Int,
      Num (..),
      Ord (..),
      Read,
      fromEnum,
      otherwise,
      toEnum,
    )
  import qualified Prelude
  
  -- # The Finite Stream Constructors
  -------------------------------------------------------------------------------
  
  -- | A singleton stream
  --
  -- @
  -- \>\>\> stdoutLn $ yield "hello"
  -- hello
  -- @
  --
  -- @
  -- \>\>\> S.sum $ do {yield 1; yield 2; yield 3}
  -- 6 :> ()
  -- @
  yield :: (Control.Monad m) => a -> Stream (Of a) m ()
  yield x = Step $ x :> Return ()
  {-# INLINE yield #-}
  
  -- | Stream the elements of a pure, foldable container.
  --
  -- @
  -- \>\>\> S.print $ each' [1..3]
  -- 1
  -- 2
  -- 3
  -- @
  each' :: (Control.Monad m) => [a] -> Stream (Of a) m ()
  each' xs = Prelude.foldr (\a stream -> Step $ a :> stream) (Return ()) xs
  {-# INLINEABLE each' #-}
  
  -- | Build a @Stream@ by unfolding steps starting from a seed. In particular note
  --    that @S.unfoldr S.next = id@.
  unfoldr ::
    (Control.Monad m) =>
    (s %1 -> m (Either r (Ur a, s))) ->
    s %1 ->
    Stream (Of a) m r
  unfoldr step s = unfoldr' step s
    where
      unfoldr' ::
        (Control.Monad m) =>
        (s %1 -> m (Either r (Ur a, s))) ->
        s %1 ->
        Stream (Of a) m r
      unfoldr' step s =
        Effect $
          step s Control.>>= \case
            Left r -> Control.return $ Return r
            Right (Ur a, s') ->
              Control.return $ Step $ a :> unfoldr step s'
  {-# INLINEABLE unfoldr #-}
  
  -- Note: we use the RIO monad from linear base to enforce
  -- the protocol of file handles and file I/O
  fromHandle :: Handle %1 -> Stream (Of Text) RIO ()
  fromHandle h = loop h
    where
      loop :: Handle %1 -> Stream (Of Text) RIO ()
      loop h = Control.do
        (Ur isEOF, h') <- Control.lift $ hIsEOF h
        case isEOF of
          True -> Control.do
            Control.lift $ hClose h'
            Control.return ()
          False -> Control.do
            (Ur text, h'') <- Control.lift $ hGetLine h'
            yield text
            fromHandle h''
  {-# INLINEABLE fromHandle #-}
  
  -- | Read the lines of a file given the filename.
  readFile :: FilePath -> Stream (Of Text) RIO ()
  readFile path = Control.do
    handle <- Control.lift $ openFile path System.ReadMode
    fromHandle handle
  
  -- | Repeat an element several times.
  replicate :: (HasCallStack, Control.Monad m) => Int -> a -> Stream (Of a) m ()
  replicate n a
    | n < 0 = Prelude.error "Cannot replicate a stream of negative length"
    | otherwise = loop n a
    where
      loop :: (Control.Monad m) => Int -> a -> Stream (Of a) m ()
      loop n a
        | n == 0 = Return ()
        | otherwise = Effect $ Control.return $ Step $ a :> loop (n - 1) a
  {-# INLINEABLE replicate #-}
  
  -- | Repeat an action several times, streaming its results.
  --
  -- @
  -- \>\>\> import qualified Unsafe.Linear as Unsafe
  -- \>\>\> import qualified Data.Time as Time
  -- \>\>\> let getCurrentTime = fromSystemIO (Unsafe.coerce Time.getCurrentTime)
  -- \>\>\> S.print $ S.replicateM 2 getCurrentTime
  -- 2015-08-18 00:57:36.124508 UTC
  -- 2015-08-18 00:57:36.124785 UTC
  -- @
  replicateM ::
    (Control.Monad m) =>
    Int ->
    m (Ur a) ->
    Stream (Of a) m ()
  replicateM n ma
    | n < 0 = Prelude.error "Cannot replicate a stream of negative length"
    | otherwise = loop n ma
    where
      loop :: (Control.Monad m) => Int -> m (Ur a) -> Stream (Of a) m ()
      loop n ma
        | n == 0 = Return ()
        | otherwise = Effect $ Control.do
            Ur a <- ma
            Control.return $ Step $ a :> (replicateM (n - 1) ma)
  
  -- | Replicate a constant element and zip it with the finite stream which
  -- is the first argument.
  replicateZip ::
    (Control.Monad m) =>
    Stream (Of x) m r ->
    a ->
    Stream (Of (a, x)) m r
  replicateZip stream a = map ((,) a) stream
  {-# INLINEABLE replicateZip #-}
  
  untilRight ::
    forall m a r.
    (Control.Monad m) =>
    m (Either (Ur a) r) ->
    Stream (Of a) m r
  untilRight mEither = Effect loop
    where
      loop :: m (Stream (Of a) m r)
      loop = Control.do
        either <- mEither
        case either of
          Left (Ur a) ->
            Control.return $ Step $ a :> (untilRight mEither)
          Right r -> Control.return $ Return r
  {-# INLINEABLE untilRight #-}
  
  -- # The \"Affine\" 'Stream'
  -------------------------------------------------------------------------------
  
  -- | An *affine stream is represented with a state of type @x@,
  -- a possibly terminating step function of type @(x %1-> m (Either (f x) r))@,
  -- and a stop-short function @(x %1-> m r)@.
  --
  -- This mirrors the unfold of a normal stream:
  --
  -- > data Stream f m r where
  -- >   Stream :: x %1-> (x %1-> m (Either (f x) r)) -> Stream f m r
  --
  -- *Though referred to as an \"affine stream\" this might not be the correct
  -- definition for affine streams. Sorting this out requires a bit more
  -- careful thought.
  data AffineStream f m r where
    AffineStream ::
      x %1 ->
      (x %1 -> m (Either (f x) r)) ->
      (x %1 -> m r) ->
      AffineStream f m r
  
  -- | Take @n@ number of elements from the affine stream, for non-negative
  -- @n@. (Negative @n@ is treated as 0.)
  take ::
    forall f m r.
    (Control.Monad m, Control.Functor f) =>
    Int ->
    AffineStream f m r %1 ->
    Stream f m r
  take = loop
    where
      loop :: Int -> AffineStream f m r %1 -> Stream f m r
      loop n (AffineStream s step end)
        | n <= 0 = Effect $ Control.fmap Control.return $ end s
        | otherwise = Effect $ Control.do
            next <- step s
            case next of
              Right r -> Control.return (Return r)
              Left fx ->
                Control.return $
                  Step $
                    Control.fmap (\x -> loop (n - 1) (AffineStream x step end)) fx
  {-# INLINEABLE take #-}
  
  -- | Run an affine stream until it ends or a monadic test succeeds.
  -- Drop the element it succeeds on.
  untilM ::
    forall a m r.
    (Control.Monad m) =>
    (a -> m Bool) ->
    AffineStream (Of a) m r %1 ->
    Stream (Of a) m r
  untilM = loop
    where
      loop :: (a -> m Bool) -> AffineStream (Of a) m r %1 -> Stream (Of a) m r
      loop test (AffineStream s step end) = Effect $ Control.do
        next <- step s
        case next of
          Right r -> Control.return (Return r)
          Left (a :> next) -> Control.do
            testResult <- test a
            case testResult of
              False ->
                Control.return $
                  Step $
                    a :> loop test (AffineStream next step end)
              True -> Control.fmap Control.return $ end next
  {-# INLINEABLE untilM #-}
  
  -- | Like 'untilM' but without the monadic test.
  until ::
    forall a m r.
    (Control.Monad m) =>
    (a -> Bool) ->
    AffineStream (Of a) m r %1 ->
    Stream (Of a) m r
  until = loop
    where
      loop :: (a -> Bool) -> AffineStream (Of a) m r %1 -> Stream (Of a) m r
      loop test (AffineStream s step end) = Effect $ Control.do
        next <- step s
        case next of
          Right r -> Control.return (Return r)
          Left (a :> next) -> case test a of
            True -> Control.fmap Control.return $ end next
            False ->
              Control.return $
                Step $
                  a :> loop test (AffineStream next step end)
  {-# INLINEABLE until #-}
  
  -- | Zip a finite stream with an affine stream.
  zip ::
    forall a x m r1 r2.
    (Control.Monad m) =>
    Stream (Of x) m r1 %1 ->
    AffineStream (Of a) m r2 %1 ->
    Stream (Of (x, a)) m (r1, r2)
  zip = loop
    where
      loop ::
        Stream (Of x) m r1 %1 ->
        AffineStream (Of a) m r2 %1 ->
        Stream (Of (x, a)) m (r1, r2)
      loop stream (AffineStream s step end) =
        case stream of
          Return r1 ->
            Effect $
              Control.fmap (\r2 -> Control.return $ (r1, r2)) $
                end s
          Effect m ->
            Effect $
              Control.fmap (\str -> loop str (AffineStream s step end)) m
          Step (x :> rest) -> Effect $ Control.do
            next <- step s
            case next of
              Right r2 -> Control.do
                r1 <- effects rest
                Control.return (Return (r1, r2))
              Left (a :> rest') ->
                Control.return $
                  Step $
                    (x, a) :> loop rest (AffineStream rest' step end)
  {-# INLINEABLE zip #-}
  
  -- | An affine stream of standard input lines.
  stdinLn :: AffineStream (Of Text) IO ()
  stdinLn = AffineStream () getALine Control.pure
    where
      getALine :: () %1 -> IO (Either (Of Text ()) ())
      getALine () = Control.do
        Ur line <- fromSystemIOU System.getLine
        Control.return $ Left (Text.pack line :> ())
  
  -- | An affine stream of reading lines, crashing on failed parse.
  readLn :: (Read a) => AffineStream (Of a) IO ()
  readLn = AffineStream () readALine Control.pure
    where
      readALine :: (Read a) => () %1 -> IO (Either (Of a ()) ())
      readALine () = Control.do
        Ur line <- fromSystemIOU System.getLine
        Control.return $ Left (Prelude.read line :> ())
  
  -- | An affine stream iterating an initial state forever.
  iterate ::
    forall a m.
    (Control.Monad m) =>
    a ->
    (a -> a) ->
    AffineStream (Of a) m ()
  iterate a step =
    AffineStream (Ur a) stepper (\x -> Control.return $ consume x)
    where
      stepper :: Ur a %1 -> m (Either (Of a (Ur a)) ())
      stepper (Ur a) =
        Control.return $
          Left $
            a :> Ur (step a)
  
  -- | An affine stream monadically iterating an initial state forever.
  iterateM ::
    forall a m.
    (Control.Monad m) =>
    m (Ur a) ->
    (a -> m (Ur a)) ->
    AffineStream (Of a) m ()
  iterateM ma step =
    AffineStream ma stepper (Control.fmap consume)
    where
      stepper :: m (Ur a) %1 -> m (Either (Of a (m (Ur a))) ())
      stepper ma = Control.do
        Ur a <- ma
        Control.return $ Left $ a :> (step a)
  
  -- Remark. In order to implement the affine break function, which is the third
  -- argument of the constructor, we need to specify the functor as @Of@.
  -- Approaches to keeping it functor general seem messy.
  
  -- | An affine stream cycling through a given finite stream forever.
  cycle ::
    forall a m r.
    (Control.Monad m, Consumable r) =>
    Stream (Of a) m r ->
    AffineStream (Of a) m r
  cycle stream =
    -- Note. The state is (original stream, stream_in_current_cycle)
    AffineStream (Ur stream, stream) stepStream leftoverEffects
    where
      leftoverEffects ::
        (Ur (Stream (Of a) m r), Stream (Of a) m r) %1 -> m r
      leftoverEffects (Ur _, str) = effects str
  
      stepStream ::
        (Control.Functor f) =>
        (Ur (Stream f m r), Stream f m r) %1 ->
        m (Either (f (Ur (Stream f m r), Stream f m r)) r)
      stepStream (Ur s, str) =
        case str of
          Return r -> lseq r $ stepStream (Ur s, s)
          Effect m ->
            m Control.>>= (\stream -> stepStream (Ur s, stream))
          Step f ->
            Control.return $
              Left $
                Control.fmap ((,) (Ur s)) f
  
  -- | An affine stream iterating an enumerated stream forever.
  enumFrom :: (Control.Monad m, Enum e) => e -> AffineStream (Of e) m ()
  enumFrom e = iterate e Prelude.succ
  
  -- | An affine stream iterating an enumerated stream forever, using the
  -- first two elements to determine the gap to skip by.
  -- E.g., @enumFromThen  3 5@ is like @[3,5..]@.
  enumFromThen ::
    forall e m.
    (Control.Monad m, Enum e) =>
    e ->
    e ->
    AffineStream (Of e) m ()
  enumFromThen e e' = iterate e enumStep
    where
      enumStep :: e -> e
      enumStep enum =
        toEnum $
          (fromEnum enum) + ((fromEnum e') - (fromEnum e))
  
  -- Think:  \enum -> enum + stepSize where stepSize = (e1 - e0)
  
  -- # Working with infinite 'Stream's
  -------------------------------------------------------------------------------
  
  -- | @stdinLnN n@ is a stream of @n@ lines from standard input
  stdinLnN :: Int -> Stream (Of Text) IO ()
  stdinLnN n = take n stdinLn
  {-# INLINE stdinLnN #-}
  
  -- | Provides a stream of standard input and omits the first line
  -- that satisfies the predicate, possibly requiring IO
  stdinLnUntilM :: (Text -> IO Bool) -> Stream (Of Text) IO ()
  stdinLnUntilM test = untilM test stdinLn
  {-# INLINE stdinLnUntilM #-}
  
  -- | Provides a stream of standard input and omits the first line
  -- that satisfies the predicate
  stdinLnUntil :: (Text -> Bool) -> Stream (Of Text) IO ()
  stdinLnUntil test = until test stdinLn
  {-# INLINE stdinLnUntil #-}
  
  -- | Given a finite stream, provide a stream of lines of standard input
  -- zipped with that finite stream
  stdinLnZip :: Stream (Of x) IO r %1 -> Stream (Of (x, Text)) IO r
  stdinLnZip stream = Control.fmap (\(r, ()) -> r) $ zip stream stdinLn
  {-# INLINE stdinLnZip #-}
  
  readLnN :: (Read a) => Int -> Stream (Of a) IO ()
  readLnN n = take n readLn
  {-# INLINE readLnN #-}
  
  readLnUntilM :: (Read a) => (a -> IO Bool) -> Stream (Of a) IO ()
  readLnUntilM test = untilM test readLn
  {-# INLINE readLnUntilM #-}
  
  readLnUntil :: (Read a) => (a -> Bool) -> Stream (Of a) IO ()
  readLnUntil test = until test readLn
  {-# INLINE readLnUntil #-}
  
  readLnZip :: (Read a) => Stream (Of x) IO r %1 -> Stream (Of (x, a)) IO r
  readLnZip stream = Control.fmap (\(r, ()) -> r) $ zip stream readLn
  {-# INLINE readLnZip #-}
  
  -- | Iterate a pure function from a seed value,
  -- streaming the results forever.
  iterateN :: (Control.Monad m) => Int -> (a -> a) -> a -> Stream (Of a) m ()
  iterateN n step a = take n $ iterate a step
  {-# INLINE iterateN #-}
  
  iterateZip ::
    (Control.Monad m) =>
    Stream (Of x) m r ->
    (a -> a) ->
    a ->
    Stream (Of (x, a)) m r
  iterateZip stream step a =
    Control.fmap (\(r, ()) -> r) $ zip stream $ iterate a step
  {-# INLINE iterateZip #-}
  
  -- | Iterate a monadic function from a seed value,
  -- streaming the results forever.
  iterateMN ::
    (Control.Monad m) =>
    Int ->
    (a -> m (Ur a)) ->
    m (Ur a) ->
    Stream (Of a) m ()
  iterateMN n step ma = take n $ iterateM ma step
  {-# INLINE iterateMN #-}
  
  iterateMZip ::
    (Control.Monad m) =>
    Stream (Of x) m r %1 ->
    (a -> m (Ur a)) ->
    m (Ur a) ->
    Stream (Of (x, a)) m r
  iterateMZip stream step ma =
    Control.fmap (\(r, ()) -> r) $ zip stream $ iterateM ma step
  {-# INLINE iterateMZip #-}
  
  -- | Cycle a stream a finite number of times
  cycleN ::
    (Control.Monad m, Consumable r) =>
    Int ->
    Stream (Of a) m r ->
    Stream (Of a) m r
  cycleN n stream = take n $ cycle stream
  {-# INLINE cycleN #-}
  
  -- | @cycleZip s1 s2@ will cycle @s2@ just enough to zip with the given finite
  -- stream @s1@. Note that we consume all the effects of the remainder of the
  -- cycled stream @s2@. That is, we consume @s2@ the smallest natural number of
  -- times we need to zip.
  cycleZip ::
    (Control.Monad m, Consumable s) =>
    Stream (Of a) m r %1 ->
    Stream (Of b) m s ->
    Stream (Of (a, b)) m (r, s)
  cycleZip str stream = zip str $ cycle stream
  {-# INLINE cycleZip #-}
  
  -- | An finite sequence of enumerable values at a fixed distance, determined
  --   by the first and second values.
  --
  -- @
  -- \>\>\> S.print $ S.enumFromThenN 3 100 200
  -- 100
  -- 200
  -- 300
  -- @
  enumFromThenN :: (Control.Monad m, Enum e) => Int -> e -> e -> Stream (Of e) m ()
  enumFromThenN n e e' = take n $ enumFromThen e e'
  {-# INLINE enumFromThenN #-}
  
  -- | A finite sequence of enumerable values at a fixed distance determined
  -- by the first and second values. The length is limited by zipping
  -- with a given finite stream, i.e., the first argument.
  enumFromThenZip ::
    (Control.Monad m, Enum e) =>
    Stream (Of a) m r %1 ->
    e ->
    e ->
    Stream (Of (a, e)) m r
  enumFromThenZip stream e e' =
    Control.fmap (\(r, ()) -> r) $ zip stream $ enumFromThen e e'
  {-# INLINE enumFromThenZip #-}
  
  -- | Like 'enumFromThenN' but where the next element in the enumeration is just
  -- the successor @succ n@ for a given enum @n@.
  enumFromN :: (Control.Monad m, Enum e) => Int -> e -> Stream (Of e) m ()
  enumFromN n e = take n $ enumFrom e
  {-# INLINE enumFromN #-}
  
  -- | Like 'enumFromThenZip' but where the next element in the enumeration is just
  -- the successor @succ n@ for a given enum @n@.
  enumFromZip ::
    (Control.Monad m, Enum e) =>
    Stream (Of a) m r %1 ->
    e ->
    Stream (Of (a, e)) m r
  enumFromZip str e =
    Control.fmap (\(r, ()) -> r) $ zip str $ enumFrom e
  {-# INLINE enumFromZip #-}
  