Ticket #4270: 4270.hs

File 4270.hs, 37.2 KB (added by igloo, 21 months ago)
Line 
1{-# LANGUAGE CPP, ExistentialQuantification, Rank2Types #-}
2
3#define USE_PRAGMAS
4
5#ifdef USE_PRAGMAS
6#define INPRAG(f)  {-# INLINE f #-}
7#define INPRAG0(f) {-# INLINE [0] f #-}
8#define INPRAG1(f) {-# INLINE [1] f #-}
9#else
10#define INPRAG(f)
11#define INPRAG0(f)
12#define INPRAG1(f)
13#endif
14
15module Data.Vector.Fusion.Stream.Monadic (
16  Stream(..), Step(..),
17  size, sized,
18  length, null,
19  empty, singleton, cons, snoc, replicate, replicateM, generate, generateM, (++),
20  head, last, (!!),
21  slice, init, tail, take, drop,
22  map, mapM, mapM_, trans, unbox, concatMap,
23  indexed, indexedR, zipWithM_,
24  zipWithM, zipWith3M, zipWith4M, zipWith5M, zipWith6M,
25  zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
26  zip, zip3, zip4, zip5, zip6,
27  filter, filterM, takeWhile, takeWhileM, dropWhile, dropWhileM,
28  elem, notElem, find, findM, findIndex, findIndexM,
29  foldl, foldlM, foldl1, foldl1M, foldM, fold1M,
30  foldlx, foldlMx, foldl1x, foldl1Mx, foldMx, fold1Mx,
31  foldr, foldrM, foldr1, foldr1M,
32  and, or, concatMapM,
33  unfoldr, unfoldrM,
34  unfoldrN, unfoldrNM,
35  prescanl, prescanlM, prescanlx, prescanlMx,
36  postscanl, postscanlM, postscanlx, postscanlMx,
37  scanl, scanlM, scanlx, scanlMx,
38  scanl1, scanl1M, scanl1x, scanl1Mx,
39  enumFromStepN, enumFromTo, enumFromThenTo,
40  toList, fromList, fromListN, unsafeFromList
41) where
42
43import qualified Prelude as P
44
45import Data.Char      ( ord )
46import GHC.Base       ( unsafeChr )
47import Control.Monad  ( liftM )
48import Prelude hiding ( length, null,
49                        replicate, (++),
50                        head, last, (!!),
51                        init, tail, take, drop,
52                        map, mapM, mapM_, concatMap,
53                        zipWith, zipWith3, zip, zip3,
54                        filter, takeWhile, dropWhile,
55                        elem, notElem,
56                        foldl, foldl1, foldr, foldr1,
57                        and, or,
58                        scanl, scanl1,
59                        enumFromTo, enumFromThenTo )
60
61data SPEC = SPEC
62
63data Step s a = Yield a s  -- ^ a new element and a new seed
64              | Skip    s  -- ^ just a new seed
65              | Done       -- ^ end of stream
66
67data Stream m a = forall s. Stream (s -> m (Step s a)) s Size
68
69size :: Stream m a -> Size
70INPRAG(size)
71size (Stream _ _ sz) = sz
72
73sized :: Stream m a -> Size -> Stream m a
74INPRAG1(sized)
75sized (Stream step s _) sz = Stream step s sz
76
77length :: Monad m => Stream m a -> m Int
78INPRAG1(length)
79length s = foldlx (\n _ -> n+1) 0 s
80
81null :: Monad m => Stream m a -> m Bool
82INPRAG1(null)
83null s = foldr (\_ _ -> False) True s
84
85empty :: Monad m => Stream m a
86INPRAG1(empty)
87empty = Stream (const (return Done)) () (Exact 0)
88
89singleton :: Monad m => a -> Stream m a
90INPRAG1(singleton)
91singleton x = Stream (return . step) True (Exact 1)
92  where
93    INPRAG0(step)
94    step True  = Yield x False
95    step False = Done
96
97replicate :: Monad m => Int -> a -> Stream m a
98INPRAG(replicate)
99replicate n x = replicateM n (return x)
100
101replicateM :: Monad m => Int -> m a -> Stream m a
102INPRAG1(replicateM)
103replicateM n p = Stream step n (Exact (delay_inline max n 0))
104  where
105    INPRAG0(step)
106    step i | i <= 0    = return Done
107           | otherwise = do { x <- p; return $ Yield x (i-1) }
108
109generate :: Monad m => Int -> (Int -> a) -> Stream m a
110INPRAG(generate)
111generate n f = generateM n (return . f)
112
113generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
114INPRAG1(generateM)
115generateM n f = n `seq` Stream step 0 (Exact (delay_inline max n 0))
116  where
117    INPRAG0(step)
118    step i | i < n     = do
119                           x <- f i
120                           return $ Yield x (i+1)
121           | otherwise = return Done
122
123cons :: Monad m => a -> Stream m a -> Stream m a
124INPRAG(cons)
125cons x s = singleton x ++ s
126
127snoc :: Monad m => Stream m a -> a -> Stream m a
128INPRAG(snoc)
129snoc s x = s ++ singleton x
130
131infixr 5 ++
132(++) :: Monad m => Stream m a -> Stream m a -> Stream m a
133INPRAG1((++))
134Stream stepa sa na ++ Stream stepb sb nb = Stream step (Left sa) (na + nb)
135  where
136    INPRAG0(step)
137    step (Left  sa) = do
138                        r <- stepa sa
139                        case r of
140                          Yield x sa' -> return $ Yield x (Left  sa')
141                          Skip    sa' -> return $ Skip    (Left  sa')
142                          Done        -> return $ Skip    (Right sb)
143    step (Right sb) = do
144                        r <- stepb sb
145                        case r of
146                          Yield x sb' -> return $ Yield x (Right sb')
147                          Skip    sb' -> return $ Skip    (Right sb')
148                          Done        -> return $ Done
149
150head :: Monad m => Stream m a -> m a
151INPRAG1(head)
152head (Stream step s _) = head_loop SPEC s
153  where
154    head_loop SPEC s
155      = do
156          r <- step s
157          case r of
158            Yield x _  -> return x
159            Skip    s' -> head_loop SPEC s'
160            Done       -> error "head"
161
162last :: Monad m => Stream m a -> m a
163INPRAG1(last)
164last (Stream step s _) = last_loop0 SPEC s
165  where
166    last_loop0 SPEC s
167      = do
168          r <- step s
169          case r of
170            Yield x s' -> last_loop1 SPEC x s'
171            Skip    s' -> last_loop0 SPEC   s'
172            Done       -> error "last"
173
174    last_loop1 SPEC x s
175      = do
176          r <- step s
177          case r of
178            Yield y s' -> last_loop1 SPEC y s'
179            Skip    s' -> last_loop1 SPEC x s'
180            Done       -> return x
181
182(!!) :: Monad m => Stream m a -> Int -> m a
183INPRAG((!!))
184Stream step s _ !! i | i < 0     = error "!! negative index"
185                     | otherwise = index_loop SPEC s i
186  where
187    index_loop SPEC s i
188      = i `seq`
189        do
190          r <- step s
191          case r of
192            Yield x s' | i == 0    -> return x
193                       | otherwise -> index_loop SPEC s' (i-1)
194            Skip    s'             -> index_loop SPEC s' i
195            Done                   -> error "!!"
196
197slice :: Monad m => Int
198                 -> Int
199                 -> Stream m a
200                 -> Stream m a
201INPRAG(slice)
202slice i n s = take n (drop i s)
203
204init :: Monad m => Stream m a -> Stream m a
205INPRAG1(init)
206init (Stream step s sz) = Stream stepx (Nothing, s) (sz - 1)
207  where
208    INPRAG0(stepx)
209    stepx (Nothing, s) = liftM (\r ->
210                           case r of
211                             Yield x s' -> Skip (Just x,  s')
212                             Skip    s' -> Skip (Nothing, s')
213                             Done       -> error "init"
214                         ) (step s)
215
216    stepx (Just x,  s) = liftM (\r ->
217                           case r of
218                             Yield y s' -> Yield x (Just y, s')
219                             Skip    s' -> Skip    (Just x, s')
220                             Done       -> Done
221                         ) (step s)
222
223tail :: Monad m => Stream m a -> Stream m a
224INPRAG1(tail)
225tail (Stream step s sz) = Stream stepx (Left s) (sz - 1)
226  where
227    INPRAG0(stepx)
228    stepx (Left  s) = liftM (\r ->
229                        case r of
230                          Yield x s' -> Skip (Right s')
231                          Skip    s' -> Skip (Left  s')
232                          Done       -> error "tail"
233                      ) (step s)
234
235    stepx (Right s) = liftM (\r ->
236                        case r of
237                          Yield x s' -> Yield x (Right s')
238                          Skip    s' -> Skip    (Right s')
239                          Done       -> Done
240                      ) (step s)
241
242take :: Monad m => Int -> Stream m a -> Stream m a
243INPRAG1(take)
244take n (Stream step s sz) = Stream stepx (s, 0) (smaller (Exact n) sz)
245  where
246    INPRAG0(stepx)
247    stepx (s, i) | i < n = liftM (\r ->
248                             case r of
249                               Yield x s' -> Yield x (s', i+1)
250                               Skip    s' -> Skip    (s', i)
251                               Done       -> Done
252                           ) (step s)
253    stepx (s, i) = return Done
254
255drop :: Monad m => Int -> Stream m a -> Stream m a
256INPRAG1(drop)
257drop n (Stream step s sz) = Stream stepx (s, Just n) (sz - Exact n)
258  where
259    INPRAG0(stepx)
260    stepx (s, Just i) | i > 0 = liftM (\r ->
261                                case r of
262                                   Yield x s' -> Skip (s', Just (i-1))
263                                   Skip    s' -> Skip (s', Just i)
264                                   Done       -> Done
265                                ) (step s)
266                      | otherwise = return $ Skip (s, Nothing)
267
268    stepx (s, Nothing) = liftM (\r ->
269                           case r of
270                             Yield x s' -> Yield x (s', Nothing)
271                             Skip    s' -> Skip    (s', Nothing)
272                             Done       -> Done
273                           ) (step s)
274
275instance Monad m => Functor (Stream m) where
276  INPRAG(fmap)
277  fmap = map
278
279map :: Monad m => (a -> b) -> Stream m a -> Stream m b
280INPRAG(map)
281map f = mapM (return . f)
282
283mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
284INPRAG1(mapM)
285mapM f (Stream step s n) = Stream stepx s n
286  where
287    INPRAG0(stepx)
288    stepx s = do
289                r <- step s
290                case r of
291                  Yield x s' -> liftM  (`Yield` s') (f x)
292                  Skip    s' -> return (Skip    s')
293                  Done       -> return Done
294
295consume :: Monad m => Stream m a -> m ()
296INPRAG1(consume)
297consume (Stream step s _) = consume_loop SPEC s
298  where
299    consume_loop SPEC s
300      = do
301          r <- step s
302          case r of
303            Yield _ s' -> consume_loop SPEC s'
304            Skip    s' -> consume_loop SPEC s'
305            Done       -> return ()
306
307mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
308INPRAG1(mapM_)
309mapM_ m = consume . mapM m
310
311trans :: (Monad m, Monad m') => (forall a. m a -> m' a)
312                             -> Stream m a -> Stream m' a
313INPRAG1(trans)
314trans f (Stream step s n) = Stream (f . step) s n
315
316unbox :: Monad m => Stream m (Box a) -> Stream m a
317INPRAG1(unbox)
318unbox (Stream step s n) = Stream stepx s n
319  where
320    INPRAG0(stepx)
321    stepx s = do
322                r <- step s
323                case r of
324                  Yield (Box x) s' -> return $ Yield x s'
325                  Skip          s' -> return $ Skip    s'
326                  Done             -> return $ Done
327
328indexed :: Monad m => Stream m a -> Stream m (Int,a)
329INPRAG1(indexed)
330indexed (Stream step s n) = Stream stepx (s,0) n
331  where
332    INPRAG0(stepx)
333    stepx (s,i) = i `seq`
334                  do
335                    r <- step s
336                    case r of
337                      Yield x s' -> return $ Yield (i,x) (s', i+1)
338                      Skip    s' -> return $ Skip        (s', i)
339                      Done       -> return Done
340
341indexedR :: Monad m => Int -> Stream m a -> Stream m (Int,a)
342INPRAG1(indexedR)
343indexedR m (Stream step s n) = Stream stepx (s,m) n
344  where
345    INPRAG0(stepx)
346    stepx (s,i) = i `seq`
347                  do
348                    r <- step s
349                    case r of
350                      Yield x s' -> let i' = i-1
351                                    in
352                                    return $ Yield (i',x) (s', i')
353                      Skip    s' -> return $ Skip         (s', i)
354                      Done       -> return Done
355
356zipWithM :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
357INPRAG1(zipWithM)
358zipWithM f (Stream stepa sa na) (Stream stepb sb nb)
359  = Stream step (sa, sb, Nothing) (smaller na nb)
360  where
361    INPRAG0(step)
362    step (sa, sb, Nothing) = liftM (\r ->
363                               case r of
364                                 Yield x sa' -> Skip (sa', sb, Just x)
365                                 Skip    sa' -> Skip (sa', sb, Nothing)
366                                 Done        -> Done
367                             ) (stepa sa)
368
369    step (sa, sb, Just x)  = do
370                               r <- stepb sb
371                               case r of
372                                 Yield y sb' ->
373                                   do
374                                     z <- f x y
375                                     return $ Yield z (sa, sb', Nothing)
376                                 Skip    sb' -> return $ Skip (sa, sb', Just x)
377                                 Done        -> return $ Done
378
379zipWithM_ :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> m ()
380INPRAG(zipWithM_)
381zipWithM_ f sa sb = consume (zipWithM f sa sb)
382
383zipWith3M :: Monad m => (a -> b -> c -> m d) -> Stream m a -> Stream m b -> Stream m c -> Stream m d
384INPRAG1(zipWith3M)
385zipWith3M f (Stream stepa sa na) (Stream stepb sb nb) (Stream stepc sc nc)
386  = Stream step (sa, sb, sc, Nothing) (smaller na (smaller nb nc))
387  where
388    INPRAG0(step)
389    step (sa, sb, sc, Nothing) = do
390        r <- stepa sa
391        return $ case r of
392            Yield x sa' -> Skip (sa', sb, sc, Just (x, Nothing))
393            Skip    sa' -> Skip (sa', sb, sc, Nothing)
394            Done        -> Done
395
396    step (sa, sb, sc, Just (x, Nothing)) = do
397        r <- stepb sb
398        return $ case r of
399            Yield y sb' -> Skip (sa, sb', sc, Just (x, Just y))
400            Skip    sb' -> Skip (sa, sb', sc, Just (x, Nothing))
401            Done        -> Done
402
403    step (sa, sb, sc, Just (x, Just y)) = do
404        r <- stepc sc
405        case r of
406            Yield z sc' -> f x y z >>= (\res -> return $ Yield res (sa, sb, sc', Nothing))
407            Skip    sc' -> return $ Skip (sa, sb, sc', Just (x, Just y))
408            Done        -> return $ Done
409
410zipWith4M :: Monad m => (a -> b -> c -> d -> m e)
411                     -> Stream m a -> Stream m b -> Stream m c -> Stream m d
412                     -> Stream m e
413INPRAG(zipWith4M)
414zipWith4M f sa sb sc sd
415  = zipWithM (\(a,b) (c,d) -> f a b c d) (zip sa sb) (zip sc sd)
416
417zipWith5M :: Monad m => (a -> b -> c -> d -> e -> m f)
418                     -> Stream m a -> Stream m b -> Stream m c -> Stream m d
419                     -> Stream m e -> Stream m f
420INPRAG(zipWith5M)
421zipWith5M f sa sb sc sd se
422  = zipWithM (\(a,b,c) (d,e) -> f a b c d e) (zip3 sa sb sc) (zip sd se)
423
424zipWith6M :: Monad m => (a -> b -> c -> d -> e -> f -> m g)
425                     -> Stream m a -> Stream m b -> Stream m c -> Stream m d
426                     -> Stream m e -> Stream m f -> Stream m g
427INPRAG(zipWith6M)
428zipWith6M fn sa sb sc sd se sf
429  = zipWithM (\(a,b,c) (d,e,f) -> fn a b c d e f) (zip3 sa sb sc)
430                                                  (zip3 sd se sf)
431
432zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
433INPRAG(zipWith)
434zipWith f = zipWithM (\a b -> return (f a b))
435
436zipWith3 :: Monad m => (a -> b -> c -> d)
437                    -> Stream m a -> Stream m b -> Stream m c -> Stream m d
438INPRAG(zipWith3)
439zipWith3 f = zipWith3M (\a b c -> return (f a b c))
440
441zipWith4 :: Monad m => (a -> b -> c -> d -> e)
442                    -> Stream m a -> Stream m b -> Stream m c -> Stream m d
443                    -> Stream m e
444INPRAG(zipWith4)
445zipWith4 f = zipWith4M (\a b c d -> return (f a b c d))
446
447zipWith5 :: Monad m => (a -> b -> c -> d -> e -> f)
448                    -> Stream m a -> Stream m b -> Stream m c -> Stream m d
449                    -> Stream m e -> Stream m f
450INPRAG(zipWith5)
451zipWith5 f = zipWith5M (\a b c d e -> return (f a b c d e))
452
453zipWith6 :: Monad m => (a -> b -> c -> d -> e -> f -> g)
454                    -> Stream m a -> Stream m b -> Stream m c -> Stream m d
455                    -> Stream m e -> Stream m f -> Stream m g
456INPRAG(zipWith6)
457zipWith6 fn = zipWith6M (\a b c d e f -> return (fn a b c d e f))
458
459zip :: Monad m => Stream m a -> Stream m b -> Stream m (a,b)
460INPRAG(zip)
461zip = zipWith (,)
462
463zip3 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m (a,b,c)
464INPRAG(zip3)
465zip3 = zipWith3 (,,)
466
467zip4 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m d
468                -> Stream m (a,b,c,d)
469INPRAG(zip4)
470zip4 = zipWith4 (,,,)
471
472zip5 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m d
473                -> Stream m e -> Stream m (a,b,c,d,e)
474INPRAG(zip5)
475zip5 = zipWith5 (,,,,)
476
477zip6 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m d
478                -> Stream m e -> Stream m f -> Stream m (a,b,c,d,e,f)
479INPRAG(zip6)
480zip6 = zipWith6 (,,,,,)
481
482filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
483INPRAG(filter)
484filter f = filterM (return . f)
485
486filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
487INPRAG1(filterM)
488filterM f (Stream step s n) = Stream stepx s (toMax n)
489  where
490    INPRAG0(stepx)
491    stepx s = do
492                r <- step s
493                case r of
494                  Yield x s' -> do
495                                  b <- f x
496                                  return $ if b then Yield x s'
497                                                else Skip    s'
498                  Skip    s' -> return $ Skip s'
499                  Done       -> return $ Done
500
501takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
502INPRAG(takeWhile)
503takeWhile f = takeWhileM (return . f)
504
505takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
506INPRAG1(takeWhileM)
507takeWhileM f (Stream step s n) = Stream stepx s (toMax n)
508  where
509    INPRAG0(stepx)
510    stepx s = do
511                r <- step s
512                case r of
513                  Yield x s' -> do
514                                  b <- f x
515                                  return $ if b then Yield x s' else Done
516                  Skip    s' -> return $ Skip s'
517                  Done       -> return $ Done
518
519dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
520INPRAG(dropWhile)
521dropWhile f = dropWhileM (return . f)
522
523data DropWhile s a = DropWhile_Drop s | DropWhile_Yield a s | DropWhile_Next s
524
525dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
526INPRAG1(dropWhileM)
527dropWhileM f (Stream step s n) = Stream stepx (DropWhile_Drop s) (toMax n)
528  where
529    INPRAG0(stepx)
530    stepx (DropWhile_Drop s)
531      = do
532          r <- step s
533          case r of
534            Yield x s' -> do
535                            b <- f x
536                            return $ if b then Skip (DropWhile_Drop    s')
537                                          else Skip (DropWhile_Yield x s')
538            Skip    s' -> return $ Skip (DropWhile_Drop    s')
539            Done       -> return $ Done
540
541    stepx (DropWhile_Yield x s) = return $ Yield x (DropWhile_Next s)
542
543    stepx (DropWhile_Next s)
544      = liftM (\r ->
545          case r of
546            Yield x s' -> Skip    (DropWhile_Yield x s')
547            Skip    s' -> Skip    (DropWhile_Next    s')
548            Done       -> Done
549        ) (step s)
550
551infix 4 `elem`
552elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
553INPRAG1(elem)
554elem x (Stream step s _) = elem_loop SPEC s
555  where
556    elem_loop SPEC s
557      = do
558          r <- step s
559          case r of
560            Yield y s' | x == y    -> return True
561                       | otherwise -> elem_loop SPEC s'
562            Skip    s'             -> elem_loop SPEC s'
563            Done                   -> return False
564
565infix 4 `notElem`
566notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
567INPRAG(notElem)
568notElem x s = liftM not (elem x s)
569
570find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
571INPRAG(find)
572find f = findM (return . f)
573
574findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
575INPRAG1(findM)
576findM f (Stream step s _) = find_loop SPEC s
577  where
578    find_loop SPEC s
579      = do
580          r <- step s
581          case r of
582            Yield x s' -> do
583                            b <- f x
584                            if b then return $ Just x
585                                 else find_loop SPEC s'
586            Skip    s' -> find_loop SPEC s'
587            Done       -> return Nothing
588
589findIndex :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe Int)
590INPRAG1(findIndex)
591findIndex f = findIndexM (return . f)
592
593findIndexM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe Int)
594INPRAG1(findIndexM)
595findIndexM f (Stream step s _) = findIndex_loop SPEC s 0
596  where
597    findIndex_loop SPEC s i
598      = do
599          r <- step s
600          case r of
601            Yield x s' -> do
602                            b <- f x
603                            if b then return $ Just i
604                                 else findIndex_loop SPEC s' (i+1)
605            Skip    s' -> findIndex_loop SPEC s' i
606            Done       -> return Nothing
607
608foldl :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
609INPRAG(foldl)
610foldl f = foldlM (\a b -> return (f a b))
611
612foldlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
613INPRAG1(foldlM)
614foldlM m z (Stream step s _) = foldlM_loop SPEC z s
615  where
616    foldlM_loop SPEC z s
617      = do
618          r <- step s
619          case r of
620            Yield x s' -> do { z' <- m z x; foldlM_loop SPEC z' s' }
621            Skip    s' -> foldlM_loop SPEC z s'
622            Done       -> return z
623
624foldM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
625INPRAG(foldM)
626foldM = foldlM
627
628foldl1 :: Monad m => (a -> a -> a) -> Stream m a -> m a
629INPRAG(foldl1)
630foldl1 f = foldl1M (\a b -> return (f a b))
631
632foldl1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
633INPRAG1(foldl1M)
634foldl1M f (Stream step s sz) = foldl1M_loop SPEC s
635  where
636    foldl1M_loop SPEC s
637      = do
638          r <- step s
639          case r of
640            Yield x s' -> foldlM f x (Stream step s' (sz - 1))
641            Skip    s' -> foldl1M_loop SPEC s'
642            Done       -> error "foldl1M"
643
644fold1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
645INPRAG(fold1M)
646fold1M = foldl1M
647
648foldlx :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
649INPRAG(foldlx)
650foldlx f = foldlMx (\a b -> return (f a b))
651
652foldlMx :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
653INPRAG1(foldlMx)
654foldlMx m z (Stream step s _) = foldlMx_loop SPEC z s
655  where
656    foldlMx_loop SPEC z s
657      = z `seq`
658        do
659          r <- step s
660          case r of
661            Yield x s' -> do { z' <- m z x; foldlMx_loop SPEC z' s' }
662            Skip    s' -> foldlMx_loop SPEC z s'
663            Done       -> return z
664
665foldMx :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
666INPRAG(foldMx)
667foldMx = foldlMx
668
669foldl1x :: Monad m => (a -> a -> a) -> Stream m a -> m a
670INPRAG(foldl1x)
671foldl1x f = foldl1Mx (\a b -> return (f a b))
672
673foldl1Mx :: Monad m => (a -> a -> m a) -> Stream m a -> m a
674INPRAG1(foldl1Mx)
675foldl1Mx f (Stream step s sz) = foldl1Mx_loop SPEC s
676  where
677    foldl1Mx_loop SPEC s
678      = do
679          r <- step s
680          case r of
681            Yield x s' -> foldlMx f x (Stream step s' (sz - 1))
682            Skip    s' -> foldl1Mx_loop SPEC s'
683            Done       -> error "foldl1Mx"
684
685fold1Mx :: Monad m => (a -> a -> m a) -> Stream m a -> m a
686INPRAG(fold1Mx)
687fold1Mx = foldl1Mx
688
689foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
690INPRAG(foldr)
691foldr f = foldrM (\a b -> return (f a b))
692
693foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m a -> m b
694INPRAG1(foldrM)
695foldrM f z (Stream step s _) = foldrM_loop SPEC s
696  where
697    foldrM_loop SPEC s
698      = do
699          r <- step s
700          case r of
701            Yield x s' -> f x =<< foldrM_loop SPEC s'
702            Skip    s' -> foldrM_loop SPEC s'
703            Done       -> return z
704
705foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m a
706INPRAG(foldr1)
707foldr1 f = foldr1M (\a b -> return (f a b))
708
709foldr1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
710INPRAG1(foldr1M)
711foldr1M f (Stream step s _) = foldr1M_loop0 SPEC s
712  where
713    foldr1M_loop0 SPEC s
714      = do
715          r <- step s
716          case r of
717            Yield x s' -> foldr1M_loop1 SPEC x s'
718            Skip    s' -> foldr1M_loop0 SPEC   s'
719            Done       -> error "foldr1M"
720
721    foldr1M_loop1 SPEC x s
722      = do
723          r <- step s
724          case r of
725            Yield y s' -> f x =<< foldr1M_loop1 SPEC y s'
726            Skip    s' -> foldr1M_loop1 SPEC x s'
727            Done       -> return x
728
729and :: Monad m => Stream m Bool -> m Bool
730INPRAG1(and)
731and (Stream step s _) = and_loop SPEC s
732  where
733    and_loop SPEC s
734      = do
735          r <- step s
736          case r of
737            Yield False _  -> return False
738            Yield True  s' -> and_loop SPEC s'
739            Skip        s' -> and_loop SPEC s'
740            Done           -> return True
741
742or :: Monad m => Stream m Bool -> m Bool
743INPRAG1(or)
744or (Stream step s _) = or_loop SPEC s
745  where
746    or_loop SPEC s
747      = do
748          r <- step s
749          case r of
750            Yield False s' -> or_loop SPEC s'
751            Yield True  _  -> return True
752            Skip        s' -> or_loop SPEC s'
753            Done           -> return False
754
755concatMap :: Monad m => (a -> Stream m b) -> Stream m a -> Stream m b
756INPRAG(concatMap)
757concatMap f = concatMapM (return . f)
758
759concatMapM :: Monad m => (a -> m (Stream m b)) -> Stream m a -> Stream m b
760INPRAG1(concatMapM)
761concatMapM f (Stream step s _) = Stream concatMap_go (Left s) Unknown
762  where
763    concatMap_go (Left s) = do
764        r <- step s
765        case r of
766            Yield a s' -> do
767                b_stream <- f a
768                return $ Skip (Right (b_stream, s'))
769            Skip    s' -> return $ Skip (Left s')
770            Done       -> return Done
771    concatMap_go (Right (Stream inner_step inner_s sz, s)) = do
772        r <- inner_step inner_s
773        case r of
774            Yield b inner_s' -> return $ Yield b (Right (Stream inner_step inner_s' sz, s))
775            Skip    inner_s' -> return $ Skip (Right (Stream inner_step inner_s' sz, s))
776            Done             -> return $ Skip (Left s)
777
778unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
779INPRAG1(unfoldr)
780unfoldr f = unfoldrM (return . f)
781
782unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
783INPRAG1(unfoldrM)
784unfoldrM f s = Stream step s Unknown
785  where
786    INPRAG0(step)
787    step s = liftM (\r ->
788               case r of
789                 Just (x, s') -> Yield x s'
790                 Nothing      -> Done
791             ) (f s)
792
793unfoldrN :: Monad m => Int -> (s -> Maybe (a, s)) -> s -> Stream m a
794INPRAG1(unfoldrN)
795unfoldrN n f = unfoldrNM n (return . f)
796
797unfoldrNM :: Monad m => Int -> (s -> m (Maybe (a, s))) -> s -> Stream m a
798INPRAG1(unfoldrNM)
799unfoldrNM n f s = Stream step (s,n) (Max (delay_inline max n 0))
800  where
801    INPRAG0(step)
802    step (s,n) | n <= 0    = return Done
803               | otherwise = liftM (\r ->
804                               case r of
805                                 Just (x,s') -> Yield x (s',n-1)
806                                 Nothing     -> Done
807                             ) (f s)
808
809prescanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
810INPRAG(prescanl)
811prescanl f = prescanlM (\a b -> return (f a b))
812
813prescanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
814INPRAG1(prescanlM)
815prescanlM f z (Stream step s sz) = Stream stepx (s,z) sz
816  where
817    INPRAG0(stepx)
818    stepx (s,x) = do
819                    r <- step s
820                    case r of
821                      Yield y s' -> do
822                                      z <- f x y
823                                      return $ Yield x (s', z)
824                      Skip    s' -> return $ Skip (s', x)
825                      Done       -> return Done
826
827prescanlx :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
828INPRAG(prescanlx)
829prescanlx f = prescanlMx (\a b -> return (f a b))
830
831prescanlMx :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
832INPRAG1(prescanlMx)
833prescanlMx f z (Stream step s sz) = Stream stepx (s,z) sz
834  where
835    INPRAG0(stepx)
836    stepx (s,x) = x `seq`
837                  do
838                    r <- step s
839                    case r of
840                      Yield y s' -> do
841                                      z <- f x y
842                                      return $ Yield x (s', z)
843                      Skip    s' -> return $ Skip (s', x)
844                      Done       -> return Done
845
846postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
847INPRAG(postscanl)
848postscanl f = postscanlM (\a b -> return (f a b))
849
850postscanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
851INPRAG1(postscanlM)
852postscanlM f z (Stream step s sz) = Stream stepx (s,z) sz
853  where
854    INPRAG0(stepx)
855    stepx (s,x) = do
856                    r <- step s
857                    case r of
858                      Yield y s' -> do
859                                      z <- f x y
860                                      return $ Yield z (s',z)
861                      Skip    s' -> return $ Skip (s',x)
862                      Done       -> return Done
863
864postscanlx :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
865INPRAG(postscanlx)
866postscanlx f = postscanlMx (\a b -> return (f a b))
867
868postscanlMx :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
869INPRAG1(postscanlMx)
870postscanlMx f z (Stream step s sz) = z `seq` Stream stepx (s,z) sz
871  where
872    INPRAG0(stepx)
873    stepx (s,x) = x `seq`
874                  do
875                    r <- step s
876                    case r of
877                      Yield y s' -> do
878                                      z <- f x y
879                                      z `seq` return (Yield z (s',z))
880                      Skip    s' -> return $ Skip (s',x)
881                      Done       -> return Done
882
883scanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
884INPRAG(scanl)
885scanl f = scanlM (\a b -> return (f a b))
886
887scanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
888INPRAG(scanlM)
889scanlM f z s = z `cons` postscanlM f z s
890
891scanlx :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
892INPRAG(scanlx)
893scanlx f = scanlMx (\a b -> return (f a b))
894
895scanlMx :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
896INPRAG(scanlMx)
897scanlMx f z s = z `seq` (z `cons` postscanlM f z s)
898
899scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
900INPRAG(scanl1)
901scanl1 f = scanl1M (\x y -> return (f x y))
902
903scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
904INPRAG1(scanl1M)
905scanl1M f (Stream step s sz) = Stream stepx (s, Nothing) sz
906  where
907    INPRAG0(stepx)
908    stepx (s, Nothing) = do
909                           r <- step s
910                           case r of
911                             Yield x s' -> return $ Yield x (s', Just x)
912                             Skip    s' -> return $ Skip (s', Nothing)
913                             Done       -> error "scanl1M"
914
915    stepx (s, Just x) = do
916                          r <- step s
917                          case r of
918                            Yield y s' -> do
919                                            z <- f x y
920                                            return $ Yield z (s', Just z)
921                            Skip    s' -> return $ Skip (s', Just x)
922                            Done       -> return Done
923
924scanl1x :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
925INPRAG(scanl1x)
926scanl1x f = scanl1Mx (\x y -> return (f x y))
927
928scanl1Mx :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
929INPRAG1(scanl1Mx)
930scanl1Mx f (Stream step s sz) = Stream stepx (s, Nothing) sz
931  where
932    INPRAG0(stepx)
933    stepx (s, Nothing) = do
934                           r <- step s
935                           case r of
936                             Yield x s' -> x `seq` return (Yield x (s', Just x))
937                             Skip    s' -> return $ Skip (s', Nothing)
938                             Done       -> error "scanl1M"
939
940    stepx (s, Just x) = x `seq`
941                        do
942                          r <- step s
943                          case r of
944                            Yield y s' -> do
945                                            z <- f x y
946                                            z `seq` return (Yield z (s', Just z))
947                            Skip    s' -> return $ Skip (s', Just x)
948                            Done       -> return Done
949
950-- Enumerations
951-- ------------
952
953-- The Enum class is broken for this, there just doesn't seem to be a
954-- way to implement this generically. We have to specialise for as many types
955-- as we can but this doesn't help in polymorphic loops.
956
957-- | Yield a 'Stream' of the given length containing the values @x@, @x+y@,
958-- @x+y+y@ etc.
959enumFromStepN :: (Num a, Monad m) => a -> a -> Int -> Stream m a
960INPRAG1(enumFromStepN)
961enumFromStepN x y n = n `seq` Stream step (x,n) (Exact (delay_inline max n 0))
962  where
963    INPRAG0(step)
964    step (x,n) | n > 0     = return $ Yield x (x+y,n-1)
965               | otherwise = return $ Done
966
967-- | Enumerate values
968--
969-- /WARNING:/ This operation can be very inefficient. If at all possible, use
970-- 'enumFromStepN' instead.
971enumFromTo :: (Enum a, Monad m) => a -> a -> Stream m a
972INPRAG1(enumFromTo)
973enumFromTo x y = fromList [x .. y]
974
975-- NOTE: We use (x+1) instead of (succ x) below because the latter checks for
976-- overflow which can't happen here.
977
978-- FIXME: add "too large" test for Int
979enumFromTo_small :: (Integral a, Monad m) => a -> a -> Stream m a
980INPRAG1(enumFromTo_small)
981enumFromTo_small x y = Stream step x (Exact n)
982  where
983    n = delay_inline max (fromIntegral y - fromIntegral x + 1) 0
984
985    INPRAG0(step)
986    step x | x <= y    = return $ Yield x (x+1)
987           | otherwise = return $ Done
988
989enumFromTo_int :: (Integral a, Monad m) => a -> a -> Stream m a
990INPRAG1(enumFromTo_int)
991enumFromTo_int x y = Stream step x (Exact (len x y))
992  where
993    INPRAG0(len)
994    len x y | x > y     = 0
995            | otherwise = fromIntegral n
996      where
997        n = y-x+1
998
999    INPRAG0(step)
1000    step x | x <= y    = return $ Yield x (x+1)
1001           | otherwise = return $ Done
1002
1003enumFromTo_big_word :: (Integral a, Monad m) => a -> a -> Stream m a
1004INPRAG1(enumFromTo_big_word)
1005enumFromTo_big_word x y = Stream step x (Exact (len x y))
1006  where
1007    INPRAG0(len)
1008    len x y | x > y     = 0
1009            | otherwise = fromIntegral (n+1)
1010      where
1011        n = y-x
1012
1013    INPRAG0(step)
1014    step x | x <= y    = return $ Yield x (x+1)
1015           | otherwise = return $ Done
1016
1017-- FIXME: the "too large" test is totally wrong
1018enumFromTo_big_int :: (Integral a, Monad m) => a -> a -> Stream m a
1019INPRAG1(enumFromTo_big_int)
1020enumFromTo_big_int x y = Stream step x (Exact (len x y))
1021  where
1022    INPRAG0(len)
1023    len x y | x > y     = 0
1024            | otherwise = fromIntegral n
1025      where
1026        n = y-x+1
1027
1028    INPRAG0(step)
1029    step x | x <= y    = return $ Yield x (x+1)
1030           | otherwise = return $ Done
1031
1032enumFromTo_char :: Monad m => Char -> Char -> Stream m Char
1033INPRAG1(enumFromTo_char)
1034enumFromTo_char x y = Stream step xn (Exact n)
1035  where
1036    xn = ord x
1037    yn = ord y
1038
1039    n = delay_inline max 0 (yn - xn + 1)
1040
1041    INPRAG0(step)
1042    step xn | xn <= yn  = return $ Yield (unsafeChr xn) (xn+1)
1043            | otherwise = return $ Done
1044
1045------------------------------------------------------------------------
1046
1047-- Specialise enumFromTo for Float and Double.
1048-- Also, try to do something about pairs?
1049
1050enumFromTo_double :: (Monad m, Ord a, RealFrac a) => a -> a -> Stream m a
1051INPRAG1(enumFromTo_double)
1052enumFromTo_double n m = Stream step n (Max (len n m))
1053  where
1054    lim = m + 1/2 -- important to float out
1055
1056    INPRAG0(len)
1057    len x y | x > y     = 0
1058            | otherwise = fromInteger n
1059      where
1060        n = truncate (y-x)+2
1061
1062    INPRAG0(step)
1063    step x | x <= lim  = return $ Yield x (x+1)
1064           | otherwise = return $ Done
1065
1066------------------------------------------------------------------------
1067
1068-- | Enumerate values with a given step.
1069--
1070-- /WARNING:/ This operation is very inefficient. If at all possible, use
1071-- 'enumFromStepN' instead.
1072enumFromThenTo :: (Enum a, Monad m) => a -> a -> a -> Stream m a
1073INPRAG1(enumFromThenTo)
1074enumFromThenTo x y z = fromList [x, y .. z]
1075
1076-- FIXME: Specialise enumFromThenTo.
1077
1078-- Conversions
1079-- -----------
1080
1081-- | Convert a 'Stream' to a list
1082toList :: Monad m => Stream m a -> m [a]
1083INPRAG(toList)
1084toList = foldr (:) []
1085
1086-- | Convert a list to a 'Stream'
1087fromList :: Monad m => [a] -> Stream m a
1088INPRAG(fromList)
1089fromList xs = unsafeFromList Unknown xs
1090
1091-- | Convert the first @n@ elements of a list to a 'Stream'
1092fromListN :: Monad m => Int -> [a] -> Stream m a
1093INPRAG1(fromListN)
1094fromListN n xs = Stream step (xs,n) (Max (delay_inline max n 0))
1095  where
1096    INPRAG0(step)
1097    step (xs,n) | n <= 0 = return Done
1098    step (x:xs,n)        = return (Yield x (xs,n-1))
1099    step ([],n)          = return Done
1100
1101-- | Convert a list to a 'Stream' with the given 'Size' hint.
1102unsafeFromList :: Monad m => Size -> [a] -> Stream m a
1103INPRAG1(unsafeFromList)
1104unsafeFromList sz xs = Stream step xs sz
1105  where
1106    step (x:xs) = return (Yield x xs)
1107    step []     = return Done
1108
1109delay_inline :: (a -> b) -> a -> b
1110INPRAG0(delay_inline)
1111delay_inline f = f
1112
1113data Box a = Box a
1114
1115-- | Size hint
1116data Size = Exact Int          -- ^ Exact size
1117          | Max   Int          -- ^ Upper bound on the size
1118          | Unknown            -- ^ Unknown size
1119        deriving( Eq, Show )
1120
1121instance Num Size where
1122  Exact m + Exact n = Exact (m+n)
1123  Exact m + Max   n = Max   (m+n)
1124
1125  Max   m + Exact n = Max   (m+n)
1126  Max   m + Max   n = Max   (m+n)
1127
1128  _       + _       = Unknown
1129
1130
1131  Exact m - Exact n = Exact (m-n)
1132  Exact m - Max   _ = Max   m
1133
1134  Max   m - Exact n = Max   (m-n)
1135  Max   m - Max   _ = Max   m
1136  Max   m - Unknown = Max   m
1137
1138  _       - _       = Unknown
1139
1140
1141  fromInteger n     = Exact (fromInteger n)
1142
1143  signum = undefined
1144  abs = undefined
1145  (*) = undefined
1146
1147-- | Minimum of two size hints
1148smaller :: Size -> Size -> Size
1149INPRAG(smaller)
1150smaller (Exact m) (Exact n) = Exact (delay_inline min m n)
1151smaller (Exact m) (Max   n) = Max   (delay_inline min m n)
1152smaller (Exact m) Unknown   = Max   m
1153smaller (Max   m) (Exact n) = Max   (delay_inline min m n)
1154smaller (Max   m) (Max   n) = Max   (delay_inline min m n)
1155smaller (Max   m) Unknown   = Max   m
1156smaller Unknown   (Exact n) = Max   n
1157smaller Unknown   (Max   n) = Max   n
1158smaller Unknown   Unknown   = Unknown
1159
1160-- | Convert a size hint to an upper bound
1161toMax :: Size -> Size
1162toMax (Exact n) = Max n
1163toMax (Max   n) = Max n
1164toMax Unknown   = Unknown
1165
1166data Checks = Bounds | Unsafe | Internal deriving( Eq )