{-# LANGUAGE BangPatterns, ExistentialQuantification #-}
module Data.RangeMin.Fusion.Stream.Monadic (
	Stream (..),
	length,
	iunfoldMN,
	imapM,
	imapM_,
	cap,
	generateM,
	snocM,
	ifoldlM,
	imapAccumLM,
	unbox,
	fromListN) where

import Prelude hiding (replicate, map, length)

import Data.Vector.Fusion.Util (Box (..))

import Data.RangeMin.Common.Types
import Data.RangeMin.Common.Combinators

data Stream m a = forall s . Stream (Index -> s -> m (Step s a)) s {-# UNPACK #-} !Length
data Step s a = Done | Last a | Yield a s

instance Functor (Step s) where
	{-# INLINE fmap #-}
	fmap _ Done = Done
	fmap f (Last a) = Last (f a)
	fmap f (Yield a s) = Yield (f a) s

instance Monad m => Functor (Stream m) where
	{-# INLINE fmap #-}
	fmap f (Stream suc s0 n) = Stream (liftM (fmap f) .: suc) s0 n

{-# INLINE length #-}
length :: Stream m a -> Length
length (Stream _ _ n) = n

{-# INLINE iunfoldMN #-}
iunfoldMN :: Monad m => Length -> (Index -> b -> m (Maybe (a, b))) -> b -> Stream m a
iunfoldMN !n f s0 = Stream suc s0 n where
	suc i s = do
	  unf <- f i s
	  case unf of
	    Nothing	 -> return Done
	    Just (a, s') -> return (Yield a s')

{-# INLINE imapM #-}
imapM :: Monad m => (Index -> a -> m b) -> Stream m a -> Stream m b
imapM f (Stream suc s0 n) = Stream suc' s0 n where
	suc' i s = do
	  step <- suc i s
	  case step of
	    Done	-> return Done
	    Last a	-> Last `liftM` f i a
	    Yield a s'	-> do
		    b <- f i a
		    return (Yield b s')

{-# INLINE imapM_ #-}
imapM_ :: Monad m => (Index -> a -> m b) -> Stream m a -> m ()
imapM_ f str = consume (imapM f str)

{-# INLINE cap #-}
cap :: Monad m => Stream m a -> Stream m a
cap (Stream suc s0 n) = Stream (\ i s -> if i < n then suc i s else return Done) s0 n

{-# INLINE generateM #-}
generateM :: Monad m => Length -> (Index -> m a) -> Stream m a
generateM n f = Stream suc () n where
	suc i _ = if i < n then do
	    x <- f i
	    return (Yield x ())
	  else return Done

{-# INLINE consume #-}
consume :: Monad m => Stream m a -> m ()
consume (Stream suc s0 _) = go 0 s0 where
	go !i s = do
	  step <- suc i s
	  case step of
	    Yield _ s' -> go (i+1) s'
	    _ -> return ()

{-# INLINE snocM #-}
snocM :: Monad m => Stream m a -> m a -> Stream m a
Stream suc s0 n `snocM` z = Stream suc' s0 (n+1) where
	suc' i s = if i < n then suc i s else Last `liftM` z

{-# INLINE ifoldlM #-}
ifoldlM :: Monad m => (b -> Index -> a -> m b) -> b -> Stream m a -> m b
ifoldlM f z0 (Stream suc s0 _) = go 0 s0 z0 where
	go !i s z = do
	  step <- suc i s
	  case step of
	    Done -> return z
	    Last a -> f z i a
	    Yield a s' ->
	    	go (i+1) s' =<< f z i a

{-# INLINE imapAccumLM #-}
imapAccumLM :: Monad m => (b -> Index -> a -> m (c, b)) -> b -> Stream m a -> Stream m c
imapAccumLM f z0 (Stream suc s0 n) = Stream suc' (s0, z0) n where
	suc' i (s, z) = do
	  step <- suc i s
	  case step of
	    Done -> return Done
	    Last a -> (Last . fst) `liftM` f z i a
	    Yield a s' -> do
	    	(y, z') <- f z i a
	    	return (Yield y (s', z'))

{-# INLINE unbox #-}
unbox :: Monad m => Stream m (Box a) -> Stream m a
unbox (Stream suc s0 n) = Stream suc' s0 n where
	{-# INLINE [0] suc' #-}
	suc' i s = do
	    step <- suc i s
	    case step of
	    	Yield (Box x) s' -> return (Yield x s')
	    	Last (Box x) -> return (Last x)
	    	Done -> return Done

{-# INLINE fromListN #-}
fromListN :: Monad m => Length -> [a] -> Stream m a
fromListN n xs = Stream suc xs n where
	suc _ (x:xs) = return (Yield x xs)
	suc _ [] = return Done