{-# LANGUAGE CPP, ExistentialQuantification, Rank2Types #-}

#define USE_PRAGMAS

#ifdef USE_PRAGMAS
#define INPRAG(f)  {-# INLINE f #-}
#define INPRAG0(f) {-# INLINE [0] f #-}
#define INPRAG1(f) {-# INLINE [1] f #-}
#else
#define INPRAG(f)
#define INPRAG0(f)
#define INPRAG1(f)
#endif

module Data.Vector.Fusion.Stream.Monadic (
  Stream(..), Step(..),
  size, sized,
  length, null,
  empty, singleton, cons, snoc, replicate, replicateM, generate, generateM, (++),
  head, last, (!!),
  slice, init, tail, take, drop,
  map, mapM, mapM_, trans, unbox, concatMap,
  indexed, indexedR, zipWithM_,
  zipWithM, zipWith3M, zipWith4M, zipWith5M, zipWith6M,
  zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
  zip, zip3, zip4, zip5, zip6,
  filter, filterM, takeWhile, takeWhileM, dropWhile, dropWhileM,
  elem, notElem, find, findM, findIndex, findIndexM,
  foldl, foldlM, foldl1, foldl1M, foldM, fold1M,
  foldlx, foldlMx, foldl1x, foldl1Mx, foldMx, fold1Mx,
  foldr, foldrM, foldr1, foldr1M,
  and, or, concatMapM,
  unfoldr, unfoldrM,
  unfoldrN, unfoldrNM,
  prescanl, prescanlM, prescanlx, prescanlMx,
  postscanl, postscanlM, postscanlx, postscanlMx,
  scanl, scanlM, scanlx, scanlMx,
  scanl1, scanl1M, scanl1x, scanl1Mx,
  enumFromStepN, enumFromTo, enumFromThenTo,
  toList, fromList, fromListN, unsafeFromList
) where

import qualified Prelude as P

import Data.Char      ( ord )
import GHC.Base       ( unsafeChr )
import Control.Monad  ( liftM )
import Prelude hiding ( length, null,
                        replicate, (++),
                        head, last, (!!),
                        init, tail, take, drop,
                        map, mapM, mapM_, concatMap,
                        zipWith, zipWith3, zip, zip3,
                        filter, takeWhile, dropWhile,
                        elem, notElem,
                        foldl, foldl1, foldr, foldr1,
                        and, or,
                        scanl, scanl1,
                        enumFromTo, enumFromThenTo )

data SPEC = SPEC

data Step s a = Yield a s  -- ^ a new element and a new seed
              | Skip    s  -- ^ just a new seed
              | Done       -- ^ end of stream

data Stream m a = forall s. Stream (s -> m (Step s a)) s Size

size :: Stream m a -> Size
INPRAG(size)
size (Stream _ _ sz) = sz

sized :: Stream m a -> Size -> Stream m a
INPRAG1(sized)
sized (Stream step s _) sz = Stream step s sz

length :: Monad m => Stream m a -> m Int
INPRAG1(length)
length s = foldlx (\n _ -> n+1) 0 s

null :: Monad m => Stream m a -> m Bool
INPRAG1(null)
null s = foldr (\_ _ -> False) True s

empty :: Monad m => Stream m a
INPRAG1(empty)
empty = Stream (const (return Done)) () (Exact 0)

singleton :: Monad m => a -> Stream m a
INPRAG1(singleton)
singleton x = Stream (return . step) True (Exact 1)
  where
    INPRAG0(step)
    step True  = Yield x False
    step False = Done

replicate :: Monad m => Int -> a -> Stream m a
INPRAG(replicate)
replicate n x = replicateM n (return x)

replicateM :: Monad m => Int -> m a -> Stream m a
INPRAG1(replicateM)
replicateM n p = Stream step n (Exact (delay_inline max n 0))
  where
    INPRAG0(step)
    step i | i <= 0    = return Done
           | otherwise = do { x <- p; return $ Yield x (i-1) }

generate :: Monad m => Int -> (Int -> a) -> Stream m a
INPRAG(generate)
generate n f = generateM n (return . f)

generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
INPRAG1(generateM)
generateM n f = n `seq` Stream step 0 (Exact (delay_inline max n 0))
  where
    INPRAG0(step)
    step i | i < n     = do
                           x <- f i
                           return $ Yield x (i+1)
           | otherwise = return Done

cons :: Monad m => a -> Stream m a -> Stream m a
INPRAG(cons)
cons x s = singleton x ++ s

snoc :: Monad m => Stream m a -> a -> Stream m a
INPRAG(snoc)
snoc s x = s ++ singleton x

infixr 5 ++
(++) :: Monad m => Stream m a -> Stream m a -> Stream m a
INPRAG1((++))
Stream stepa sa na ++ Stream stepb sb nb = Stream step (Left sa) (na + nb)
  where
    INPRAG0(step)
    step (Left  sa) = do
                        r <- stepa sa
                        case r of
                          Yield x sa' -> return $ Yield x (Left  sa')
                          Skip    sa' -> return $ Skip    (Left  sa')
                          Done        -> return $ Skip    (Right sb)
    step (Right sb) = do
                        r <- stepb sb
                        case r of
                          Yield x sb' -> return $ Yield x (Right sb')
                          Skip    sb' -> return $ Skip    (Right sb')
                          Done        -> return $ Done

head :: Monad m => Stream m a -> m a
INPRAG1(head)
head (Stream step s _) = head_loop SPEC s
  where
    head_loop SPEC s
      = do
          r <- step s
          case r of
            Yield x _  -> return x
            Skip    s' -> head_loop SPEC s'
            Done       -> error "head"

last :: Monad m => Stream m a -> m a
INPRAG1(last)
last (Stream step s _) = last_loop0 SPEC s
  where
    last_loop0 SPEC s
      = do
          r <- step s
          case r of
            Yield x s' -> last_loop1 SPEC x s'
            Skip    s' -> last_loop0 SPEC   s'
            Done       -> error "last"

    last_loop1 SPEC x s
      = do
          r <- step s
          case r of
            Yield y s' -> last_loop1 SPEC y s'
            Skip    s' -> last_loop1 SPEC x s'
            Done       -> return x

(!!) :: Monad m => Stream m a -> Int -> m a
INPRAG((!!))
Stream step s _ !! i | i < 0     = error "!! negative index"
                     | otherwise = index_loop SPEC s i
  where
    index_loop SPEC s i
      = i `seq`
        do
          r <- step s
          case r of
            Yield x s' | i == 0    -> return x
                       | otherwise -> index_loop SPEC s' (i-1)
            Skip    s'             -> index_loop SPEC s' i
            Done                   -> error "!!"

slice :: Monad m => Int
                 -> Int
                 -> Stream m a
                 -> Stream m a
INPRAG(slice)
slice i n s = take n (drop i s)

init :: Monad m => Stream m a -> Stream m a
INPRAG1(init)
init (Stream step s sz) = Stream stepx (Nothing, s) (sz - 1)
  where
    INPRAG0(stepx)
    stepx (Nothing, s) = liftM (\r ->
                           case r of
                             Yield x s' -> Skip (Just x,  s')
                             Skip    s' -> Skip (Nothing, s')
                             Done       -> error "init"
                         ) (step s)

    stepx (Just x,  s) = liftM (\r ->
                           case r of
                             Yield y s' -> Yield x (Just y, s')
                             Skip    s' -> Skip    (Just x, s')
                             Done       -> Done
                         ) (step s)

tail :: Monad m => Stream m a -> Stream m a
INPRAG1(tail)
tail (Stream step s sz) = Stream stepx (Left s) (sz - 1)
  where
    INPRAG0(stepx)
    stepx (Left  s) = liftM (\r ->
                        case r of
                          Yield x s' -> Skip (Right s')
                          Skip    s' -> Skip (Left  s')
                          Done       -> error "tail"
                      ) (step s)

    stepx (Right s) = liftM (\r ->
                        case r of
                          Yield x s' -> Yield x (Right s')
                          Skip    s' -> Skip    (Right s')
                          Done       -> Done
                      ) (step s)

take :: Monad m => Int -> Stream m a -> Stream m a
INPRAG1(take)
take n (Stream step s sz) = Stream stepx (s, 0) (smaller (Exact n) sz)
  where
    INPRAG0(stepx)
    stepx (s, i) | i < n = liftM (\r ->
                             case r of
                               Yield x s' -> Yield x (s', i+1)
                               Skip    s' -> Skip    (s', i)
                               Done       -> Done
                           ) (step s)
    stepx (s, i) = return Done

drop :: Monad m => Int -> Stream m a -> Stream m a
INPRAG1(drop)
drop n (Stream step s sz) = Stream stepx (s, Just n) (sz - Exact n)
  where
    INPRAG0(stepx)
    stepx (s, Just i) | i > 0 = liftM (\r ->
                                case r of
                                   Yield x s' -> Skip (s', Just (i-1))
                                   Skip    s' -> Skip (s', Just i)
                                   Done       -> Done
                                ) (step s)
                      | otherwise = return $ Skip (s, Nothing)

    stepx (s, Nothing) = liftM (\r ->
                           case r of
                             Yield x s' -> Yield x (s', Nothing)
                             Skip    s' -> Skip    (s', Nothing)
                             Done       -> Done
                           ) (step s)

instance Monad m => Functor (Stream m) where
  INPRAG(fmap)
  fmap = map

map :: Monad m => (a -> b) -> Stream m a -> Stream m b
INPRAG(map)
map f = mapM (return . f)

mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
INPRAG1(mapM)
mapM f (Stream step s n) = Stream stepx s n
  where
    INPRAG0(stepx)
    stepx s = do
                r <- step s
                case r of
                  Yield x s' -> liftM  (`Yield` s') (f x)
                  Skip    s' -> return (Skip    s')
                  Done       -> return Done

consume :: Monad m => Stream m a -> m ()
INPRAG1(consume)
consume (Stream step s _) = consume_loop SPEC s
  where
    consume_loop SPEC s
      = do
          r <- step s
          case r of
            Yield _ s' -> consume_loop SPEC s'
            Skip    s' -> consume_loop SPEC s'
            Done       -> return ()

mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
INPRAG1(mapM_)
mapM_ m = consume . mapM m

trans :: (Monad m, Monad m') => (forall a. m a -> m' a)
                             -> Stream m a -> Stream m' a
INPRAG1(trans)
trans f (Stream step s n) = Stream (f . step) s n

unbox :: Monad m => Stream m (Box a) -> Stream m a
INPRAG1(unbox)
unbox (Stream step s n) = Stream stepx s n
  where
    INPRAG0(stepx)
    stepx s = do
                r <- step s
                case r of
                  Yield (Box x) s' -> return $ Yield x s'
                  Skip          s' -> return $ Skip    s'
                  Done             -> return $ Done

indexed :: Monad m => Stream m a -> Stream m (Int,a)
INPRAG1(indexed)
indexed (Stream step s n) = Stream stepx (s,0) n
  where
    INPRAG0(stepx)
    stepx (s,i) = i `seq`
                  do
                    r <- step s
                    case r of
                      Yield x s' -> return $ Yield (i,x) (s', i+1)
                      Skip    s' -> return $ Skip        (s', i)
                      Done       -> return Done

indexedR :: Monad m => Int -> Stream m a -> Stream m (Int,a)
INPRAG1(indexedR)
indexedR m (Stream step s n) = Stream stepx (s,m) n
  where
    INPRAG0(stepx)
    stepx (s,i) = i `seq`
                  do
                    r <- step s
                    case r of
                      Yield x s' -> let i' = i-1
                                    in
                                    return $ Yield (i',x) (s', i')
                      Skip    s' -> return $ Skip         (s', i)
                      Done       -> return Done

zipWithM :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
INPRAG1(zipWithM)
zipWithM f (Stream stepa sa na) (Stream stepb sb nb)
  = Stream step (sa, sb, Nothing) (smaller na nb)
  where
    INPRAG0(step)
    step (sa, sb, Nothing) = liftM (\r ->
                               case r of
                                 Yield x sa' -> Skip (sa', sb, Just x)
                                 Skip    sa' -> Skip (sa', sb, Nothing)
                                 Done        -> Done
                             ) (stepa sa)

    step (sa, sb, Just x)  = do
                               r <- stepb sb
                               case r of
                                 Yield y sb' ->
                                   do
                                     z <- f x y
                                     return $ Yield z (sa, sb', Nothing)
                                 Skip    sb' -> return $ Skip (sa, sb', Just x)
                                 Done        -> return $ Done

zipWithM_ :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> m ()
INPRAG(zipWithM_)
zipWithM_ f sa sb = consume (zipWithM f sa sb)

zipWith3M :: Monad m => (a -> b -> c -> m d) -> Stream m a -> Stream m b -> Stream m c -> Stream m d
INPRAG1(zipWith3M)
zipWith3M f (Stream stepa sa na) (Stream stepb sb nb) (Stream stepc sc nc)
  = Stream step (sa, sb, sc, Nothing) (smaller na (smaller nb nc))
  where
    INPRAG0(step)
    step (sa, sb, sc, Nothing) = do
        r <- stepa sa
        return $ case r of
            Yield x sa' -> Skip (sa', sb, sc, Just (x, Nothing))
            Skip    sa' -> Skip (sa', sb, sc, Nothing)
            Done        -> Done

    step (sa, sb, sc, Just (x, Nothing)) = do
        r <- stepb sb
        return $ case r of
            Yield y sb' -> Skip (sa, sb', sc, Just (x, Just y))
            Skip    sb' -> Skip (sa, sb', sc, Just (x, Nothing))
            Done        -> Done

    step (sa, sb, sc, Just (x, Just y)) = do
        r <- stepc sc
        case r of
            Yield z sc' -> f x y z >>= (\res -> return $ Yield res (sa, sb, sc', Nothing))
            Skip    sc' -> return $ Skip (sa, sb, sc', Just (x, Just y))
            Done        -> return $ Done

zipWith4M :: Monad m => (a -> b -> c -> d -> m e)
                     -> Stream m a -> Stream m b -> Stream m c -> Stream m d
                     -> Stream m e
INPRAG(zipWith4M)
zipWith4M f sa sb sc sd
  = zipWithM (\(a,b) (c,d) -> f a b c d) (zip sa sb) (zip sc sd)

zipWith5M :: Monad m => (a -> b -> c -> d -> e -> m f)
                     -> Stream m a -> Stream m b -> Stream m c -> Stream m d
                     -> Stream m e -> Stream m f
INPRAG(zipWith5M)
zipWith5M f sa sb sc sd se
  = zipWithM (\(a,b,c) (d,e) -> f a b c d e) (zip3 sa sb sc) (zip sd se)

zipWith6M :: Monad m => (a -> b -> c -> d -> e -> f -> m g)
                     -> Stream m a -> Stream m b -> Stream m c -> Stream m d
                     -> Stream m e -> Stream m f -> Stream m g
INPRAG(zipWith6M)
zipWith6M fn sa sb sc sd se sf
  = zipWithM (\(a,b,c) (d,e,f) -> fn a b c d e f) (zip3 sa sb sc)
                                                  (zip3 sd se sf)

zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
INPRAG(zipWith)
zipWith f = zipWithM (\a b -> return (f a b))

zipWith3 :: Monad m => (a -> b -> c -> d)
                    -> Stream m a -> Stream m b -> Stream m c -> Stream m d
INPRAG(zipWith3)
zipWith3 f = zipWith3M (\a b c -> return (f a b c))

zipWith4 :: Monad m => (a -> b -> c -> d -> e)
                    -> Stream m a -> Stream m b -> Stream m c -> Stream m d
                    -> Stream m e
INPRAG(zipWith4)
zipWith4 f = zipWith4M (\a b c d -> return (f a b c d))

zipWith5 :: Monad m => (a -> b -> c -> d -> e -> f)
                    -> Stream m a -> Stream m b -> Stream m c -> Stream m d
                    -> Stream m e -> Stream m f
INPRAG(zipWith5)
zipWith5 f = zipWith5M (\a b c d e -> return (f a b c d e))

zipWith6 :: Monad m => (a -> b -> c -> d -> e -> f -> g)
                    -> Stream m a -> Stream m b -> Stream m c -> Stream m d
                    -> Stream m e -> Stream m f -> Stream m g
INPRAG(zipWith6)
zipWith6 fn = zipWith6M (\a b c d e f -> return (fn a b c d e f))

zip :: Monad m => Stream m a -> Stream m b -> Stream m (a,b)
INPRAG(zip)
zip = zipWith (,)

zip3 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m (a,b,c)
INPRAG(zip3)
zip3 = zipWith3 (,,)

zip4 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m d
                -> Stream m (a,b,c,d)
INPRAG(zip4)
zip4 = zipWith4 (,,,)

zip5 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m d
                -> Stream m e -> Stream m (a,b,c,d,e)
INPRAG(zip5)
zip5 = zipWith5 (,,,,)

zip6 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m d
                -> Stream m e -> Stream m f -> Stream m (a,b,c,d,e,f)
INPRAG(zip6)
zip6 = zipWith6 (,,,,,)

filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
INPRAG(filter)
filter f = filterM (return . f)

filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
INPRAG1(filterM)
filterM f (Stream step s n) = Stream stepx s (toMax n)
  where
    INPRAG0(stepx)
    stepx s = do
                r <- step s
                case r of
                  Yield x s' -> do
                                  b <- f x
                                  return $ if b then Yield x s'
                                                else Skip    s'
                  Skip    s' -> return $ Skip s'
                  Done       -> return $ Done

takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
INPRAG(takeWhile)
takeWhile f = takeWhileM (return . f)

takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
INPRAG1(takeWhileM)
takeWhileM f (Stream step s n) = Stream stepx s (toMax n)
  where
    INPRAG0(stepx)
    stepx s = do
                r <- step s
                case r of
                  Yield x s' -> do
                                  b <- f x
                                  return $ if b then Yield x s' else Done
                  Skip    s' -> return $ Skip s'
                  Done       -> return $ Done

dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
INPRAG(dropWhile)
dropWhile f = dropWhileM (return . f)

data DropWhile s a = DropWhile_Drop s | DropWhile_Yield a s | DropWhile_Next s

dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
INPRAG1(dropWhileM)
dropWhileM f (Stream step s n) = Stream stepx (DropWhile_Drop s) (toMax n)
  where
    INPRAG0(stepx)
    stepx (DropWhile_Drop s)
      = do
          r <- step s
          case r of
            Yield x s' -> do
                            b <- f x
                            return $ if b then Skip (DropWhile_Drop    s')
                                          else Skip (DropWhile_Yield x s')
            Skip    s' -> return $ Skip (DropWhile_Drop    s')
            Done       -> return $ Done

    stepx (DropWhile_Yield x s) = return $ Yield x (DropWhile_Next s)

    stepx (DropWhile_Next s)
      = liftM (\r ->
          case r of
            Yield x s' -> Skip    (DropWhile_Yield x s')
            Skip    s' -> Skip    (DropWhile_Next    s')
            Done       -> Done
        ) (step s)

infix 4 `elem`
elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
INPRAG1(elem)
elem x (Stream step s _) = elem_loop SPEC s
  where
    elem_loop SPEC s
      = do
          r <- step s
          case r of
            Yield y s' | x == y    -> return True
                       | otherwise -> elem_loop SPEC s'
            Skip    s'             -> elem_loop SPEC s'
            Done                   -> return False

infix 4 `notElem`
notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
INPRAG(notElem)
notElem x s = liftM not (elem x s)

find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
INPRAG(find)
find f = findM (return . f)

findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
INPRAG1(findM)
findM f (Stream step s _) = find_loop SPEC s
  where
    find_loop SPEC s
      = do
          r <- step s
          case r of
            Yield x s' -> do
                            b <- f x
                            if b then return $ Just x
                                 else find_loop SPEC s'
            Skip    s' -> find_loop SPEC s'
            Done       -> return Nothing

findIndex :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe Int)
INPRAG1(findIndex)
findIndex f = findIndexM (return . f)

findIndexM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe Int)
INPRAG1(findIndexM)
findIndexM f (Stream step s _) = findIndex_loop SPEC s 0
  where
    findIndex_loop SPEC s i
      = do
          r <- step s
          case r of
            Yield x s' -> do
                            b <- f x
                            if b then return $ Just i
                                 else findIndex_loop SPEC s' (i+1)
            Skip    s' -> findIndex_loop SPEC s' i
            Done       -> return Nothing

foldl :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
INPRAG(foldl)
foldl f = foldlM (\a b -> return (f a b))

foldlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
INPRAG1(foldlM)
foldlM m z (Stream step s _) = foldlM_loop SPEC z s
  where
    foldlM_loop SPEC z s
      = do
          r <- step s
          case r of
            Yield x s' -> do { z' <- m z x; foldlM_loop SPEC z' s' }
            Skip    s' -> foldlM_loop SPEC z s'
            Done       -> return z

foldM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
INPRAG(foldM)
foldM = foldlM

foldl1 :: Monad m => (a -> a -> a) -> Stream m a -> m a
INPRAG(foldl1)
foldl1 f = foldl1M (\a b -> return (f a b))

foldl1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
INPRAG1(foldl1M)
foldl1M f (Stream step s sz) = foldl1M_loop SPEC s
  where
    foldl1M_loop SPEC s
      = do
          r <- step s
          case r of
            Yield x s' -> foldlM f x (Stream step s' (sz - 1))
            Skip    s' -> foldl1M_loop SPEC s'
            Done       -> error "foldl1M"

fold1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
INPRAG(fold1M)
fold1M = foldl1M

foldlx :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
INPRAG(foldlx)
foldlx f = foldlMx (\a b -> return (f a b))

foldlMx :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
INPRAG1(foldlMx)
foldlMx m z (Stream step s _) = foldlMx_loop SPEC z s
  where
    foldlMx_loop SPEC z s
      = z `seq`
        do
          r <- step s
          case r of
            Yield x s' -> do { z' <- m z x; foldlMx_loop SPEC z' s' }
            Skip    s' -> foldlMx_loop SPEC z s'
            Done       -> return z

foldMx :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
INPRAG(foldMx)
foldMx = foldlMx

foldl1x :: Monad m => (a -> a -> a) -> Stream m a -> m a
INPRAG(foldl1x)
foldl1x f = foldl1Mx (\a b -> return (f a b))

foldl1Mx :: Monad m => (a -> a -> m a) -> Stream m a -> m a
INPRAG1(foldl1Mx)
foldl1Mx f (Stream step s sz) = foldl1Mx_loop SPEC s
  where
    foldl1Mx_loop SPEC s
      = do
          r <- step s
          case r of
            Yield x s' -> foldlMx f x (Stream step s' (sz - 1))
            Skip    s' -> foldl1Mx_loop SPEC s'
            Done       -> error "foldl1Mx"

fold1Mx :: Monad m => (a -> a -> m a) -> Stream m a -> m a
INPRAG(fold1Mx)
fold1Mx = foldl1Mx

foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
INPRAG(foldr)
foldr f = foldrM (\a b -> return (f a b))

foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m a -> m b
INPRAG1(foldrM)
foldrM f z (Stream step s _) = foldrM_loop SPEC s
  where
    foldrM_loop SPEC s
      = do
          r <- step s
          case r of
            Yield x s' -> f x =<< foldrM_loop SPEC s'
            Skip    s' -> foldrM_loop SPEC s'
            Done       -> return z

foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m a
INPRAG(foldr1)
foldr1 f = foldr1M (\a b -> return (f a b))

foldr1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
INPRAG1(foldr1M)
foldr1M f (Stream step s _) = foldr1M_loop0 SPEC s
  where
    foldr1M_loop0 SPEC s
      = do
          r <- step s
          case r of
            Yield x s' -> foldr1M_loop1 SPEC x s'
            Skip    s' -> foldr1M_loop0 SPEC   s'
            Done       -> error "foldr1M"

    foldr1M_loop1 SPEC x s
      = do
          r <- step s
          case r of
            Yield y s' -> f x =<< foldr1M_loop1 SPEC y s'
            Skip    s' -> foldr1M_loop1 SPEC x s'
            Done       -> return x

and :: Monad m => Stream m Bool -> m Bool
INPRAG1(and)
and (Stream step s _) = and_loop SPEC s
  where
    and_loop SPEC s
      = do
          r <- step s
          case r of
            Yield False _  -> return False
            Yield True  s' -> and_loop SPEC s'
            Skip        s' -> and_loop SPEC s'
            Done           -> return True

or :: Monad m => Stream m Bool -> m Bool
INPRAG1(or)
or (Stream step s _) = or_loop SPEC s
  where
    or_loop SPEC s
      = do
          r <- step s
          case r of
            Yield False s' -> or_loop SPEC s'
            Yield True  _  -> return True
            Skip        s' -> or_loop SPEC s'
            Done           -> return False

concatMap :: Monad m => (a -> Stream m b) -> Stream m a -> Stream m b
INPRAG(concatMap)
concatMap f = concatMapM (return . f)

concatMapM :: Monad m => (a -> m (Stream m b)) -> Stream m a -> Stream m b
INPRAG1(concatMapM)
concatMapM f (Stream step s _) = Stream concatMap_go (Left s) Unknown
  where
    concatMap_go (Left s) = do
        r <- step s
        case r of
            Yield a s' -> do
                b_stream <- f a
                return $ Skip (Right (b_stream, s'))
            Skip    s' -> return $ Skip (Left s')
            Done       -> return Done
    concatMap_go (Right (Stream inner_step inner_s sz, s)) = do
        r <- inner_step inner_s
        case r of
            Yield b inner_s' -> return $ Yield b (Right (Stream inner_step inner_s' sz, s))
            Skip    inner_s' -> return $ Skip (Right (Stream inner_step inner_s' sz, s))
            Done             -> return $ Skip (Left s)

unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
INPRAG1(unfoldr)
unfoldr f = unfoldrM (return . f)

unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
INPRAG1(unfoldrM)
unfoldrM f s = Stream step s Unknown
  where
    INPRAG0(step)
    step s = liftM (\r ->
               case r of
                 Just (x, s') -> Yield x s'
                 Nothing      -> Done
             ) (f s)

unfoldrN :: Monad m => Int -> (s -> Maybe (a, s)) -> s -> Stream m a
INPRAG1(unfoldrN)
unfoldrN n f = unfoldrNM n (return . f)

unfoldrNM :: Monad m => Int -> (s -> m (Maybe (a, s))) -> s -> Stream m a
INPRAG1(unfoldrNM)
unfoldrNM n f s = Stream step (s,n) (Max (delay_inline max n 0))
  where
    INPRAG0(step)
    step (s,n) | n <= 0    = return Done
               | otherwise = liftM (\r ->
                               case r of
                                 Just (x,s') -> Yield x (s',n-1)
                                 Nothing     -> Done
                             ) (f s)

prescanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
INPRAG(prescanl)
prescanl f = prescanlM (\a b -> return (f a b))

prescanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
INPRAG1(prescanlM)
prescanlM f z (Stream step s sz) = Stream stepx (s,z) sz
  where
    INPRAG0(stepx)
    stepx (s,x) = do
                    r <- step s
                    case r of
                      Yield y s' -> do
                                      z <- f x y
                                      return $ Yield x (s', z)
                      Skip    s' -> return $ Skip (s', x)
                      Done       -> return Done

prescanlx :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
INPRAG(prescanlx)
prescanlx f = prescanlMx (\a b -> return (f a b))

prescanlMx :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
INPRAG1(prescanlMx)
prescanlMx f z (Stream step s sz) = Stream stepx (s,z) sz
  where
    INPRAG0(stepx)
    stepx (s,x) = x `seq`
                  do
                    r <- step s
                    case r of
                      Yield y s' -> do
                                      z <- f x y
                                      return $ Yield x (s', z)
                      Skip    s' -> return $ Skip (s', x)
                      Done       -> return Done

postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
INPRAG(postscanl)
postscanl f = postscanlM (\a b -> return (f a b))

postscanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
INPRAG1(postscanlM)
postscanlM f z (Stream step s sz) = Stream stepx (s,z) sz
  where
    INPRAG0(stepx)
    stepx (s,x) = do
                    r <- step s
                    case r of
                      Yield y s' -> do
                                      z <- f x y
                                      return $ Yield z (s',z)
                      Skip    s' -> return $ Skip (s',x)
                      Done       -> return Done

postscanlx :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
INPRAG(postscanlx)
postscanlx f = postscanlMx (\a b -> return (f a b))

postscanlMx :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
INPRAG1(postscanlMx)
postscanlMx f z (Stream step s sz) = z `seq` Stream stepx (s,z) sz
  where
    INPRAG0(stepx)
    stepx (s,x) = x `seq`
                  do
                    r <- step s
                    case r of
                      Yield y s' -> do
                                      z <- f x y
                                      z `seq` return (Yield z (s',z))
                      Skip    s' -> return $ Skip (s',x)
                      Done       -> return Done

scanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
INPRAG(scanl)
scanl f = scanlM (\a b -> return (f a b))

scanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
INPRAG(scanlM)
scanlM f z s = z `cons` postscanlM f z s

scanlx :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
INPRAG(scanlx)
scanlx f = scanlMx (\a b -> return (f a b))

scanlMx :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
INPRAG(scanlMx)
scanlMx f z s = z `seq` (z `cons` postscanlM f z s)

scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
INPRAG(scanl1)
scanl1 f = scanl1M (\x y -> return (f x y))

scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
INPRAG1(scanl1M)
scanl1M f (Stream step s sz) = Stream stepx (s, Nothing) sz
  where
    INPRAG0(stepx)
    stepx (s, Nothing) = do
                           r <- step s
                           case r of
                             Yield x s' -> return $ Yield x (s', Just x)
                             Skip    s' -> return $ Skip (s', Nothing)
                             Done       -> error "scanl1M"

    stepx (s, Just x) = do
                          r <- step s
                          case r of
                            Yield y s' -> do
                                            z <- f x y
                                            return $ Yield z (s', Just z)
                            Skip    s' -> return $ Skip (s', Just x)
                            Done       -> return Done

scanl1x :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
INPRAG(scanl1x)
scanl1x f = scanl1Mx (\x y -> return (f x y))

scanl1Mx :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
INPRAG1(scanl1Mx)
scanl1Mx f (Stream step s sz) = Stream stepx (s, Nothing) sz
  where
    INPRAG0(stepx)
    stepx (s, Nothing) = do
                           r <- step s
                           case r of
                             Yield x s' -> x `seq` return (Yield x (s', Just x))
                             Skip    s' -> return $ Skip (s', Nothing)
                             Done       -> error "scanl1M"

    stepx (s, Just x) = x `seq`
                        do
                          r <- step s
                          case r of
                            Yield y s' -> do
                                            z <- f x y
                                            z `seq` return (Yield z (s', Just z))
                            Skip    s' -> return $ Skip (s', Just x)
                            Done       -> return Done

-- Enumerations
-- ------------

-- The Enum class is broken for this, there just doesn't seem to be a
-- way to implement this generically. We have to specialise for as many types
-- as we can but this doesn't help in polymorphic loops.

-- | Yield a 'Stream' of the given length containing the values @x@, @x+y@,
-- @x+y+y@ etc.
enumFromStepN :: (Num a, Monad m) => a -> a -> Int -> Stream m a
INPRAG1(enumFromStepN)
enumFromStepN x y n = n `seq` Stream step (x,n) (Exact (delay_inline max n 0))
  where
    INPRAG0(step)
    step (x,n) | n > 0     = return $ Yield x (x+y,n-1)
               | otherwise = return $ Done

-- | Enumerate values
--
-- /WARNING:/ This operation can be very inefficient. If at all possible, use
-- 'enumFromStepN' instead.
enumFromTo :: (Enum a, Monad m) => a -> a -> Stream m a
INPRAG1(enumFromTo)
enumFromTo x y = fromList [x .. y]

-- NOTE: We use (x+1) instead of (succ x) below because the latter checks for
-- overflow which can't happen here.

-- FIXME: add "too large" test for Int
enumFromTo_small :: (Integral a, Monad m) => a -> a -> Stream m a
INPRAG1(enumFromTo_small)
enumFromTo_small x y = Stream step x (Exact n)
  where
    n = delay_inline max (fromIntegral y - fromIntegral x + 1) 0

    INPRAG0(step)
    step x | x <= y    = return $ Yield x (x+1)
           | otherwise = return $ Done

enumFromTo_int :: (Integral a, Monad m) => a -> a -> Stream m a
INPRAG1(enumFromTo_int)
enumFromTo_int x y = Stream step x (Exact (len x y))
  where
    INPRAG0(len)
    len x y | x > y     = 0
            | otherwise = fromIntegral n
      where
        n = y-x+1

    INPRAG0(step)
    step x | x <= y    = return $ Yield x (x+1)
           | otherwise = return $ Done

enumFromTo_big_word :: (Integral a, Monad m) => a -> a -> Stream m a
INPRAG1(enumFromTo_big_word)
enumFromTo_big_word x y = Stream step x (Exact (len x y))
  where
    INPRAG0(len)
    len x y | x > y     = 0
            | otherwise = fromIntegral (n+1)
      where
        n = y-x

    INPRAG0(step)
    step x | x <= y    = return $ Yield x (x+1)
           | otherwise = return $ Done

-- FIXME: the "too large" test is totally wrong
enumFromTo_big_int :: (Integral a, Monad m) => a -> a -> Stream m a
INPRAG1(enumFromTo_big_int)
enumFromTo_big_int x y = Stream step x (Exact (len x y))
  where
    INPRAG0(len)
    len x y | x > y     = 0
            | otherwise = fromIntegral n
      where
        n = y-x+1

    INPRAG0(step)
    step x | x <= y    = return $ Yield x (x+1)
           | otherwise = return $ Done

enumFromTo_char :: Monad m => Char -> Char -> Stream m Char
INPRAG1(enumFromTo_char)
enumFromTo_char x y = Stream step xn (Exact n)
  where
    xn = ord x
    yn = ord y

    n = delay_inline max 0 (yn - xn + 1)

    INPRAG0(step)
    step xn | xn <= yn  = return $ Yield (unsafeChr xn) (xn+1)
            | otherwise = return $ Done

------------------------------------------------------------------------

-- Specialise enumFromTo for Float and Double.
-- Also, try to do something about pairs?

enumFromTo_double :: (Monad m, Ord a, RealFrac a) => a -> a -> Stream m a
INPRAG1(enumFromTo_double)
enumFromTo_double n m = Stream step n (Max (len n m))
  where
    lim = m + 1/2 -- important to float out

    INPRAG0(len)
    len x y | x > y     = 0
            | otherwise = fromInteger n
      where
        n = truncate (y-x)+2

    INPRAG0(step)
    step x | x <= lim  = return $ Yield x (x+1)
           | otherwise = return $ Done

------------------------------------------------------------------------

-- | Enumerate values with a given step.
--
-- /WARNING:/ This operation is very inefficient. If at all possible, use
-- 'enumFromStepN' instead.
enumFromThenTo :: (Enum a, Monad m) => a -> a -> a -> Stream m a
INPRAG1(enumFromThenTo)
enumFromThenTo x y z = fromList [x, y .. z]

-- FIXME: Specialise enumFromThenTo.

-- Conversions
-- -----------

-- | Convert a 'Stream' to a list
toList :: Monad m => Stream m a -> m [a]
INPRAG(toList)
toList = foldr (:) []

-- | Convert a list to a 'Stream'
fromList :: Monad m => [a] -> Stream m a
INPRAG(fromList)
fromList xs = unsafeFromList Unknown xs

-- | Convert the first @n@ elements of a list to a 'Stream'
fromListN :: Monad m => Int -> [a] -> Stream m a
INPRAG1(fromListN)
fromListN n xs = Stream step (xs,n) (Max (delay_inline max n 0))
  where
    INPRAG0(step)
    step (xs,n) | n <= 0 = return Done
    step (x:xs,n)        = return (Yield x (xs,n-1))
    step ([],n)          = return Done

-- | Convert a list to a 'Stream' with the given 'Size' hint. 
unsafeFromList :: Monad m => Size -> [a] -> Stream m a
INPRAG1(unsafeFromList)
unsafeFromList sz xs = Stream step xs sz
  where
    step (x:xs) = return (Yield x xs)
    step []     = return Done

delay_inline :: (a -> b) -> a -> b
INPRAG0(delay_inline)
delay_inline f = f

data Box a = Box a

-- | Size hint
data Size = Exact Int          -- ^ Exact size
          | Max   Int          -- ^ Upper bound on the size
          | Unknown            -- ^ Unknown size
        deriving( Eq, Show )

instance Num Size where
  Exact m + Exact n = Exact (m+n)
  Exact m + Max   n = Max   (m+n)

  Max   m + Exact n = Max   (m+n)
  Max   m + Max   n = Max   (m+n)

  _       + _       = Unknown


  Exact m - Exact n = Exact (m-n)
  Exact m - Max   _ = Max   m

  Max   m - Exact n = Max   (m-n)
  Max   m - Max   _ = Max   m
  Max   m - Unknown = Max   m

  _       - _       = Unknown


  fromInteger n     = Exact (fromInteger n)

  signum = undefined
  abs = undefined
  (*) = undefined

-- | Minimum of two size hints
smaller :: Size -> Size -> Size
INPRAG(smaller)
smaller (Exact m) (Exact n) = Exact (delay_inline min m n)
smaller (Exact m) (Max   n) = Max   (delay_inline min m n)
smaller (Exact m) Unknown   = Max   m
smaller (Max   m) (Exact n) = Max   (delay_inline min m n)
smaller (Max   m) (Max   n) = Max   (delay_inline min m n)
smaller (Max   m) Unknown   = Max   m
smaller Unknown   (Exact n) = Max   n
smaller Unknown   (Max   n) = Max   n
smaller Unknown   Unknown   = Unknown

-- | Convert a size hint to an upper bound
toMax :: Size -> Size
toMax (Exact n) = Max n
toMax (Max   n) = Max n
toMax Unknown   = Unknown

data Checks = Bounds | Unsafe | Internal deriving( Eq )