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

import Prelude hiding (replicate, map, length, sequence_)

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

{-# 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 = sequence_ (imapM f str)

{-# 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 sequence_ #-}
sequence_ :: Monad m => Stream m a -> m ()
sequence_ (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 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