content: {-# LANGUAGE LambdaCase, RankNTypes, EmptyCase,
               StandaloneDeriving, FlexibleContexts,
               DeriveDataTypeable, DeriveFoldable,
               DeriveFunctor, DeriveTraversable,
               ScopedTypeVariables, BangPatterns #-}
  {-# LANGUAGE UndecidableInstances #-} -- for Streaming show instance
  module Streaming.Internal where
  
  import Control.Monad
  import Control.Monad.Trans
  import Control.Monad.Trans.Class
  import Control.Applicative
  import Data.Data ( Data, Typeable )
  import Data.Foldable ( Foldable )
  import Data.Traversable
  import Control.Monad.Morph
  import Data.Monoid
  import Data.Functor.Identity
  import GHC.Exts ( build )
  import Prelude hiding (splitAt)
  
  -- | A left-strict pair; the base functor for streams of individual elements.
  data Of a b = !a :> b
      deriving (Data, Eq, Foldable, Functor, Ord,
                Read, Show, Traversable, Typeable)
  infixr 4 :>
  
  -- | Curry a function of left-strict pairs
  kurry :: (Of a b -> c) -> a -> b -> c
  kurry f = \a b -> f (a :> b)
  {-# INLINE kurry #-}
  
  -- | Uncurry a function into a function on left-strict pairs 
  
  unkurry :: (a -> b -> c) -> Of a b -> c
  unkurry f = \(a :> b) -> f a b
  {-# INLINE unkurry #-}
  
  -- | @Stream@ (\'FreeT\') data type. The constructors are exported by @Streaming.Internal@
  data Stream f m r = Step !(f (Stream f m r))
                    | Delay (m (Stream f m r))
                    | Return r
                    deriving (Typeable)
  
  deriving instance (Show r, Show (m (Stream f m r))
                    , Show (f (Stream f m r))) => Show (Stream f m r)
  deriving instance (Eq r, Eq (m (Stream f m r))
                    , Eq (f (Stream f m r))) => Eq (Stream f m r)
  deriving instance (Typeable f, Typeable m, Data r, Data (m (Stream f m r))
                    , Data (f (Stream f m r))) => Data (Stream f m r)
  
  instance (Functor f, Monad m) => Functor (Stream f m) where
    fmap f = buildStream . fmap f . foldStream
    {-# INLINE fmap #-}
      -- loop = \case Step f  -> Step (fmap loop f)
      --              Delay m       -> Delay (liftM loop m)
      --              Return r       -> Return (f r)
  
  instance (Functor f, Monad m) => Monad (Stream f m) where
    return = Return
    {-# INLINE return #-}
    (>>) = \phi psi -> buildStream $ Folding (augmentFolding_ (getFolding (foldStream phi)) 
                                               (getFolding (foldStream psi)))
       where
        augmentFolding_ ::
             (forall r'.  (f r' -> r') -> (m r' -> r') -> (s -> r') -> r')
          -> (forall r'.  (f r' -> r') -> (m r' -> r') -> (r -> r') -> r')
          -> (forall r'.  (f r' -> r') -> (m r' -> r') -> (r -> r') -> r')
        augmentFolding_ = \phi psi construct wrap done -> 
                  phi construct 
                      wrap 
                      (\x  -> psi construct
                                  wrap
                                  done)
        {-# INLINE augmentFolding_ #-}                  
    {-# INLINE (>>) #-}                               
    s >>= f = buildStream (foldBind (foldStream . f) (foldStream s))
    {-# INLINE (>>=) #-}
      -- loop lst where
      -- loop = \case Step f -> Step (fmap loop f)
      --              Delay m      -> Delay (liftM loop m)
      --              Return r      -> f r
  
  instance (Functor f, Monad m) => Applicative (Stream f m) where
    pure = buildStream . return
    {-# INLINE pure #-}
    x <*> y = buildStream $ Folding $ \construct wrap done -> 
        getFolding (foldStream x) 
               construct 
               wrap 
               (\f ->  getFolding (foldStream y) 
                          construct 
                          wrap 
                          (\s -> done (f s)) )
    {-# INLINE (<*>) #-}    
    
  instance Functor f => MonadTrans (Stream f) where
    lift = buildStream . lift 
    {-# INLINE lift #-}
  
  instance Functor f => MFunctor (Stream f) where
    hoist trans = buildStream . hoist trans . foldStream
    {-# INLINE hoist #-}
      -- loop where
      -- loop = \case Step f -> Step (fmap loop f)
      --              Delay m      -> Delay (trans (liftM loop m))
      --              Return r      -> Return r
  
  instance (MonadIO m, Functor f) => MonadIO (Stream f m) where
    liftIO = buildStream . liftIO
    {-# INLINE liftIO #-}
  
  -- | Map streaming layers of one functor to another with a natural transformation
  maps :: (Monad m, Functor f) => (forall x . f x -> g x) -> Stream f m r -> Stream g m r
  maps phi = buildStream . mapsF phi . foldStream
  {-# INLINE maps #-}
  
  
  mapsM :: (Monad m, Functor f) => (forall x . f x -> m (g x)) -> Stream f m r -> Stream g m r
  mapsM phi = buildStream . mapsMF phi . foldStream
  {-# INLINE mapsM #-}
  
  maps' :: (Monad m, Functor f) 
            => (forall x . f x -> m (a, x)) 
            -> Stream f m r 
            -> Stream (Of a) m r
  maps' phi = loop where
    loop stream = case stream of 
      Return r -> Return r
      Delay m -> Delay $ liftM loop m
      Step fs -> Delay $ liftM (Step . uncurry (:>)) (phi (fmap loop fs))
  {-# INLINABLE maps' #-}
  
  -- church encodings:
  -- ----- unwrapped synonym:
  type Folding_ f m r = forall r'
                     .  (f r' -> r')
                     -> (m r' -> r')
                     -> (r -> r')
                     -> r'
  -- ------ wrapped:
  newtype Folding f m r = Folding {getFolding :: Folding_ f m r  }
  
  -- these should perhaps be expressed with
  -- predefined combinators for Folding_
  instance Functor (Folding f m) where
    fmap f phi = Folding (\construct wrap done ->
      getFolding phi construct
                  wrap
                  (done . f))
  
  instance Monad (Folding f m) where
    return r = Folding (\construct wrap done -> done r)
    (>>=) = flip foldBind
    {-# INLINE (>>=) #-}
  
  foldBind f phi = Folding (\construct wrap done ->
    getFolding phi construct
                wrap
                (\a -> getFolding (f a) construct
                                     wrap
                                     done))
  {-# INLINE foldBind #-}
  
  instance Applicative (Folding f m) where
    pure r = Folding (\construct wrap done -> done r)
    phi <*> psi = Folding (\construct wrap done ->
      getFolding phi construct
                  wrap
                  (\f -> getFolding psi construct
                                     wrap
                                     (\a -> done (f a))))
  
  instance MonadTrans (Folding f) where
    lift ma = Folding (\constr wrap done -> wrap (liftM done ma))
    {-# INLINE lift #-}
    
  instance Functor f => MFunctor (Folding f) where
    hoist trans phi = Folding (\construct wrap done ->
      getFolding phi construct (wrap . trans) done)
    {-# INLINE hoist #-}
  instance (MonadIO m, Functor f) => MonadIO (Folding f m) where
    liftIO io = Folding (\construct wrap done ->
               wrap (liftM done (liftIO io))
                  )
    {-# INLINE liftIO #-}
  
  
  mapsF :: (forall x . f x -> g x) -> Folding f m r -> Folding g m r
  mapsF morph (Folding phi) = Folding $ \construct wrap done -> 
      phi (construct . morph)
          wrap
          done
  {-# INLINE mapsF #-}
  
  mapsMF :: (Monad m) => (forall x . f x -> m (g x)) -> Folding f m r -> Folding g m r
  mapsMF morph (Folding phi) = Folding $ \construct wrap done -> 
      phi (wrap . liftM construct . morph)
          wrap
          done
  {-# INLINE mapsMF #-}
  
  
  mapsFoldF :: (Monad m) 
            => (forall x . f x -> m (a, x)) 
            -> Folding f m r 
            -> Folding (Of a) m r
  mapsFoldF crush = mapsMF (liftM (\(a,b) -> a :> b) . crush) 
  {-# INLINE mapsFoldF #-}
  -- -------------------------------------
  -- optimization operations: wrapped case
  -- -------------------------------------
  
  --
  
  -- `foldStream` is a flipped and wrapped variant of Atkey's
  -- effectfulFolding :: (Functor f, Monad m) =>
  --    (m x -> x) -> (r -> x) -> (f x -> x) -> Stream f m r -> x
  -- modulo the 'Return' constructor, which implicitly restricts the
  -- available class of Functors.
  -- See http://bentnib.org/posts/2012-01-06-streams.html and
  -- the (nightmarish) associated paper.
  
  -- Our plan is thus where possible to replace the datatype Stream with
  -- the associated effectfulFolding itself, wrapped as Folding
  
  foldStream  :: (Functor f, Monad m) => Stream f m t -> Folding f m t
  foldStream lst = Folding (destroy lst)
  {-# INLINE[0] foldStream  #-}
  
  buildStream :: Folding f m r -> Stream f m r
  buildStream (Folding phi) = phi Step Delay Return
  {-# INLINE[0] buildStream #-}
  
  
  -- The compiler has no difficulty with the rule for the wrapped case.
  -- I have not investigated whether the remaining newtype
  -- constructor is acting as an impediment. The stage [0] or [1]
  -- seems irrelevant in either case.
  
  {-# RULES
    "foldStream/buildStream" forall phi.
      foldStream (buildStream phi) = phi
      #-}
  
  -- -------------------------------------
  -- optimization operations: unwrapped case
  -- -------------------------------------
  destroy 
    :: (Functor f, Monad m) =>
       Stream f m r -> (f b -> b) -> (m b -> b) -> (r -> b) -> b
  destroy = \lst construct wrap done ->
     let loop = \case Delay mlst -> wrap (liftM loop mlst)
                      Step flst  -> construct (fmap loop flst)
                      Return r   -> done r 
     in loop lst
  {-# INLINABLE destroy #-}
  
  construct
    :: (forall b . (f b -> b) -> (m b -> b) -> (r -> b) -> b) ->  Stream f m r
  construct = \phi -> phi Step Delay Return
  {-# INLINE construct #-}
  
  
  foldStreamx
    :: (Functor f, Monad m) =>
       Stream f m t -> (f b -> b) -> (m b -> b) -> (t -> b) -> b
  foldStreamx = \lst construct wrap done ->
     let loop = \case Delay mlst -> wrap (liftM loop mlst)
                      Step flst  -> construct (fmap loop flst)
                      Return r   -> done r
     in  loop lst
  {-# INLINE[1] foldStreamx #-}
  
  
  buildStreamx = \phi -> phi Step Delay Return
  {-# INLINE[1] buildStreamx #-}
  
  -- The compiler seems to have trouble seeing these rules as applicable,
  -- unlike those for foldStream & buildStream. Opaque arity is
  -- a plausible hypothesis when you know nothing yet.
  -- When additional arguments are given to a rule,
  -- the most saturated is the one that fires,
  -- but it only fires where this one would.
  
  {-# RULES
  
    "foldStreamx/buildStreamx" forall phi.
      foldStreamx (buildStreamx phi) = phi
  
      #-}
  
  
  buildList_ :: Folding_ (Of a) Identity () -> [a]
  buildList_ phi = phi (\(a :> as) -> a : as)
                       (\(Identity xs) -> xs)
                       (\() -> [])
  {-# INLINE buildList_ #-}
  
  buildListM_ :: Monad m => Folding_ (Of a) m () -> m [a]
  buildListM_ phi = phi (\(a :> mas) -> liftM (a :) mas)
                        (>>= id)
                        (\() -> return [])
  {-# INLINE buildListM_ #-}
  
  foldList_ :: Monad m => [a] -> Folding_ (Of a) m ()
  foldList_ xs = \construct wrap done ->
             foldr (\x r -> construct (x:>r)) (done ()) xs
  {-# INLINE foldList_ #-}
  
  
  buildList :: Folding (Of a) Identity () -> [a]
  buildList = \(Folding phi) -> buildList_ phi
  {-# INLINE[0] buildList #-}
  
  foldList :: Monad m => [a] -> Folding (Of a) m ()
  foldList = \xs -> Folding (foldList_ xs)
  {-# INLINE[0] foldList #-}
  
  {-# RULES
    "foldr/buildList" forall phi op seed .
      foldr op seed (buildList phi) =
         getFolding phi (unkurry op) runIdentity (\() -> seed)
      #-}
      
  {-# RULES
    "foldr/buildList_" forall (phi :: Folding_ (Of a) Identity ()) 
                              (op :: a -> b -> b) 
                              (seed :: b) .
      foldr op seed (buildList_ phi) =
         phi (unkurry op) runIdentity (\() -> seed)
      #-}
  
  {-# RULES
    "foldList/buildList" forall phi.
      foldList(buildList phi) = phi
      #-}               
      
      
  --
  intercalates :: (Monad m, Monad (t m), MonadTrans t) =>
       t m a -> Stream (t m) m b -> t m b
  intercalates sep = go0
    where
      go0 f = case f of 
        Return r -> return r 
        Delay m -> lift m >>= go0 
        Step fstr -> do
                  f' <- fstr
                  go1 f'
      go1 f = case f of 
        Return r -> return r 
        Delay m     -> lift m >>= go1
        Step fstr ->  do
                  sep
                  f' <- fstr
                  go1 f'
  {-# INLINABLE intercalates #-}
  
  intercalates' :: (Monad m, Monad (t m), MonadTrans t) =>
       t m a -> Stream (t m) m b -> t m b
  intercalates' sep s = getFolding (foldStream s)  
     (\tmstr -> do 
       str <- tmstr
       sep
       str
       )
     (join . lift)
     return
  {-# INLINE intercalates' #-}
  
  iterTM ::
    (Functor f, Monad m, MonadTrans t,
     Monad (t m)) =>
    (f (t m a) -> t m a) -> Stream f m a -> t m a
  iterTM out str = getFolding (foldStream str) out (join . lift) return
  {-# INLINE iterTM #-}
  
  iterT ::
    (Functor f, Monad m) => (f (m a) -> m a) -> Stream f m a -> m a
  iterT out str = getFolding (foldStream str) out join return
  {-# INLINE iterT #-}
  
  concats ::
      (MonadTrans t, Monad (t m), Monad m) =>
      Stream (t m) m a -> t m a
  concats str = getFolding (foldStream str) join (join . lift) return
  {-# INLINE concats #-}
  
  
  splitAt :: (Monad m, Functor f) => Int -> Stream f m r -> Stream f m (Stream f m r)
  splitAt = loop where
    loop !n stream 
      | n <= 1 = Return stream
      | otherwise = case stream of
          Return r       -> Return (Return r)
          Delay m        -> Delay (liftM (loop n) m)
          Step fs        -> case n of 
            0 -> Return (Step fs)
            _ -> Step (fmap (loop (n-1)) fs)
  {-# INLINABLE splitAt #-}                        
  
  chunksOf :: (Monad m, Functor f) => Int -> Stream f m r -> Stream (Stream f m) m r
  chunksOf n0 = loop where
    loop stream = case stream of
      Return r       -> Return r
      Delay m        -> Delay (liftM loop m)
      Step fs        -> Step $ Step $ fmap (fmap loop . splitAt n0) fs
  {-# INLINABLE chunksOf #-}                        
  