INLINE_STREAM -> INLINE_FUSED
[darcs-mirrors/vector.git] / Data / Vector / Fusion / Stream / Monadic.hs
1 {-# LANGUAGE ExistentialQuantification, MultiParamTypeClasses, FlexibleInstances, Rank2Types, BangPatterns, KindSignatures, GADTs, ScopedTypeVariables #-}
2
3 -- |
4 -- Module : Data.Vector.Fusion.Stream.Monadic
5 -- Copyright : (c) Roman Leshchinskiy 2008-2010
6 -- License : BSD-style
7 --
8 -- Maintainer : Roman Leshchinskiy <rl@cse.unsw.edu.au>
9 -- Stability : experimental
10 -- Portability : non-portable
11 --
12 -- Monadic stream combinators.
13 --
14
15 module Data.Vector.Fusion.Stream.Monadic (
16 Facets(..), Unf(..), Step(..), Chunk(..), SPEC(..),
17
18 simple,
19
20 -- * Size hints
21 size, sized,
22
23 -- * Length
24 length, null,
25
26 -- * Construction
27 empty, singleton, cons, snoc, replicate, replicateM, generate, generateM, (++),
28
29 -- * Accessing elements
30 head, last, (!!), (!?),
31
32 -- * Substreams
33 slice, init, tail, take, drop,
34
35 -- * Mapping
36 map, mapM, mapM_, trans, unbox, concatMap, flatten,
37
38 -- * Zipping
39 indexed, indexedR, zipWithM_,
40 zipWithM, zipWith3M, zipWith4M, zipWith5M, zipWith6M,
41 zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
42 zip, zip3, zip4, zip5, zip6,
43
44 -- * Comparisons
45 eq, cmp,
46
47 -- * Filtering
48 filter, filterM, takeWhile, takeWhileM, dropWhile, dropWhileM,
49
50 -- * Searching
51 elem, notElem, find, findM, findIndex, findIndexM,
52
53 -- * Folding
54 foldl, foldlM, foldl1, foldl1M, foldM, fold1M,
55 foldl', foldlM', foldl1', foldl1M', foldM', fold1M',
56 foldr, foldrM, foldr1, foldr1M,
57
58 vfoldlM, vfoldlM',
59
60 -- * Specialised folds
61 and, or, concatMapM,
62
63 -- * Unfolding
64 unfoldr, unfoldrM,
65 unfoldrN, unfoldrNM,
66 iterateN, iterateNM,
67
68 -- * Scans
69 prescanl, prescanlM, prescanl', prescanlM',
70 postscanl, postscanlM, postscanl', postscanlM',
71 scanl, scanlM, scanl', scanlM',
72 scanl1, scanl1M, scanl1', scanl1M',
73
74 -- * Enumerations
75 enumFromStepN, enumFromTo, enumFromThenTo,
76
77 -- * Conversions
78 toList, fromList, fromListN, unsafeFromList,
79 fromVector, reVector, fromVectors, concatVectors
80 ) where
81
82 import Data.Vector.Generic.Base
83 import qualified Data.Vector.Generic.Mutable.Base as M
84 import Data.Vector.Fusion.Stream.Size
85 import Data.Vector.Fusion.Util ( Box(..), delay_inline )
86 import Control.Monad.Primitive
87
88 import qualified Data.List as List
89 import Data.Char ( ord )
90 import GHC.Base ( unsafeChr )
91 import Control.Monad ( liftM )
92 import Prelude hiding ( length, null,
93 replicate, (++),
94 head, last, (!!),
95 init, tail, take, drop,
96 map, mapM, mapM_, concatMap,
97 zipWith, zipWith3, zip, zip3,
98 filter, takeWhile, dropWhile,
99 elem, notElem,
100 foldl, foldl1, foldr, foldr1,
101 and, or,
102 scanl, scanl1,
103 enumFromTo, enumFromThenTo )
104
105 import Data.Int ( Int8, Int16, Int32, Int64 )
106 import Data.Word ( Word8, Word16, Word32, Word, Word64 )
107
108 #if __GLASGOW_HASKELL__ >= 700
109 import GHC.Exts ( SpecConstrAnnotation(..) )
110 #endif
111
112 #include "vector.h"
113
114 data SPEC = SPEC | SPEC2
115 #if __GLASGOW_HASKELL__ >= 700
116 {-# ANN type SPEC ForceSpecConstr #-}
117 #endif
118
119 emptyStream :: String
120 {-# NOINLINE emptyStream #-}
121 emptyStream = "empty stream"
122
123 #define EMPTY_STREAM (\s -> ERROR s emptyStream)
124
125 -- | Result of taking a single step in a stream
126 data Step s a where
127 Yield :: a -> s -> Step s a
128 Skip :: s -> Step s a
129 Done :: Step s a
130
131 instance Functor (Step s) where
132 {-# INLINE fmap #-}
133 fmap f (Yield x s) = Yield (f x) s
134 fmap f (Skip s) = Skip s
135 fmap f Done = Done
136
137 data Chunk v a = Chunk Int (forall m. (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> m ())
138
139 data Unf m a = forall s. Unf (s -> m (Step s a)) s
140
141 instance Monad m => Functor (Unf m) where
142 {-# INLINE fmap #-}
143 fmap f (Unf step s) = Unf step' s
144 where
145 step' s = do r <- step s ; return (fmap f r)
146
147 -- | Monadic streams
148 data Facets m v a = Facets { sElems :: Unf m a
149 , sChunks :: Unf m (Chunk v a)
150 , sVector :: Maybe (v a)
151 , sSize :: Size
152 }
153
154 simple :: Monad m => (s -> m (Step s a)) -> s -> Size -> Facets m v a
155 {-# INLINE simple #-}
156 simple step s sz = Facets (Unf step s) (Unf step' s) Nothing sz
157 where
158 step' s = do r <- step s
159 return $ fmap (\x -> Chunk 1 (\v -> M.basicUnsafeWrite v 0 x)) r
160
161 -- | 'Size' hint of a 'Facets'
162 size :: Facets m v a -> Size
163 {-# INLINE size #-}
164 size = sSize
165
166 -- | Attach a 'Size' hint to a 'Facets'
167 sized :: Facets m v a -> Size -> Facets m v a
168 {-# INLINE_FUSED sized #-}
169 sized s sz = s { sSize = sz }
170
171 -- Length
172 -- ------
173
174 -- | Length of a 'Facets'
175 length :: Monad m => Facets m v a -> m Int
176 {-# INLINE_FUSED length #-}
177 length Facets{sSize = Exact n} = return n
178 length s = vfoldl' (\n (Chunk k _) -> n+k) 0 s
179
180 -- | Check if a 'Facets' is empty
181 null :: Monad m => Facets m v a -> m Bool
182 {-# INLINE_FUSED null #-}
183 null Facets{sSize = Exact n} = return (n == 0)
184 null Facets{sChunks = Unf step s} = null_loop s
185 where
186 null_loop s = do
187 r <- step s
188 case r of
189 Yield (Chunk n _) s'
190 | n /= 0 -> return False
191 | otherwise -> null_loop s'
192 Skip s' -> null_loop s'
193 Done -> return True
194
195 -- Construction
196 -- ------------
197
198 -- | Empty 'Facets'
199 empty :: Monad m => Facets m v a
200 {-# INLINE_FUSED empty #-}
201 empty = simple (const (return Done)) () (Exact 0)
202
203 -- | Singleton 'Facets'
204 singleton :: Monad m => a -> Facets m v a
205 {-# INLINE_FUSED singleton #-}
206 singleton x = simple (return . step) True (Exact 1)
207 where
208 {-# INLINE_INNER step #-}
209 step True = Yield x False
210 step False = Done
211
212 -- | Replicate a value to a given length
213 replicate :: Monad m => Int -> a -> Facets m v a
214 {-# INLINE_FUSED replicate #-}
215 replicate n x = Facets (Unf pstep n) (Unf vstep True) Nothing (Exact len)
216 where
217 len = delay_inline max n 0
218
219 {-# INLINE_INNER pstep #-}
220 pstep i | i <= 0 = return Done
221 | otherwise = return $ Yield x (i-1)
222
223 {-# INLINE_INNER vstep #-}
224 vstep True = return $ Yield (Chunk len (\v -> M.basicSet v x)) False
225 vstep False = return Done
226
227 --replicate n x = replicateM n (return x)
228
229 -- | Yield a 'Facets' of values obtained by performing the monadic action the
230 -- given number of times
231 replicateM :: Monad m => Int -> m a -> Facets m v a
232 {-# INLINE_FUSED replicateM #-}
233 -- NOTE: We delay inlining max here because GHC will create a join point for
234 -- the call to newArray# otherwise which is not really nice.
235 replicateM n p = simple step n (Exact (delay_inline max n 0))
236 where
237 {-# INLINE_INNER step #-}
238 step i | i <= 0 = return Done
239 | otherwise = do { x <- p; return $ Yield x (i-1) }
240
241 generate :: Monad m => Int -> (Int -> a) -> Facets m v a
242 {-# INLINE generate #-}
243 generate n f = generateM n (return . f)
244
245 -- | Generate a stream from its indices
246 generateM :: Monad m => Int -> (Int -> m a) -> Facets m v a
247 {-# INLINE_FUSED generateM #-}
248 generateM n f = n `seq` simple step 0 (Exact (delay_inline max n 0))
249 where
250 {-# INLINE_INNER step #-}
251 step i | i < n = do
252 x <- f i
253 return $ Yield x (i+1)
254 | otherwise = return Done
255
256 -- | Prepend an element
257 cons :: Monad m => a -> Facets m v a -> Facets m v a
258 {-# INLINE cons #-}
259 cons x s = singleton x ++ s
260
261 -- | Append an element
262 snoc :: Monad m => Facets m v a -> a -> Facets m v a
263 {-# INLINE snoc #-}
264 snoc s x = s ++ singleton x
265
266 infixr 5 ++
267 -- | Concatenate two 'Facets's
268 (++) :: Monad m => Facets m v a -> Facets m v a -> Facets m v a
269 {-# INLINE_FUSED (++) #-}
270 Facets{sElems = Unf stepa sa, sChunks = Unf vstepa vsa, sSize = na}
271 ++ Facets{sElems = Unf stepb sb, sChunks = Unf vstepb vsb, sSize = nb}
272 = Facets (Unf step (Left sa)) (Unf vstep (Left vsa)) Nothing (na + nb)
273 where
274 {-# INLINE_INNER step #-}
275 step (Left sa) = do
276 r <- stepa sa
277 case r of
278 Yield x sa' -> return $ Yield x (Left sa')
279 Skip sa' -> return $ Skip (Left sa')
280 Done -> return $ Skip (Right sb)
281 step (Right sb) = do
282 r <- stepb sb
283 case r of
284 Yield x sb' -> return $ Yield x (Right sb')
285 Skip sb' -> return $ Skip (Right sb')
286 Done -> return $ Done
287
288 {-# INLINE_INNER vstep #-}
289 vstep (Left vsa) = do
290 r <- vstepa vsa
291 case r of
292 Yield x vsa' -> return $ Yield x (Left vsa')
293 Skip vsa' -> return $ Skip (Left vsa')
294 Done -> return $ Skip (Right vsb)
295 vstep (Right vsb) = do
296 r <- vstepb vsb
297 case r of
298 Yield x vsb' -> return $ Yield x (Right vsb')
299 Skip vsb' -> return $ Skip (Right vsb')
300 Done -> return $ Done
301
302 -- Accessing elements
303 -- ------------------
304
305 -- | First element of the 'Facets' or error if empty
306 head :: Monad m => Facets m v a -> m a
307 {-# INLINE_FUSED head #-}
308 head Facets{sElems = Unf step s} = head_loop SPEC s
309 where
310 head_loop !sPEC s
311 = do
312 r <- step s
313 case r of
314 Yield x _ -> return x
315 Skip s' -> head_loop SPEC s'
316 Done -> EMPTY_STREAM "head"
317
318
319
320 -- | Last element of the 'Facets' or error if empty
321 last :: Monad m => Facets m v a -> m a
322 {-# INLINE_FUSED last #-}
323 last Facets{sElems = Unf step s} = last_loop0 SPEC s
324 where
325 last_loop0 !sPEC s
326 = do
327 r <- step s
328 case r of
329 Yield x s' -> last_loop1 SPEC x s'
330 Skip s' -> last_loop0 SPEC s'
331 Done -> EMPTY_STREAM "last"
332
333 last_loop1 !sPEC x s
334 = do
335 r <- step s
336 case r of
337 Yield y s' -> last_loop1 SPEC y s'
338 Skip s' -> last_loop1 SPEC x s'
339 Done -> return x
340
341 infixl 9 !!
342 -- | Element at the given position
343 (!!) :: Monad m => Facets m v a -> Int -> m a
344 {-# INLINE (!!) #-}
345 Facets{sElems = Unf step s} !! i | i < 0 = ERROR "!!" "negative index"
346 | otherwise = index_loop SPEC s i
347 where
348 index_loop !sPEC s i
349 = i `seq`
350 do
351 r <- step s
352 case r of
353 Yield x s' | i == 0 -> return x
354 | otherwise -> index_loop SPEC s' (i-1)
355 Skip s' -> index_loop SPEC s' i
356 Done -> EMPTY_STREAM "!!"
357
358 infixl 9 !?
359 -- | Element at the given position or 'Nothing' if out of bounds
360 (!?) :: Monad m => Facets m v a -> Int -> m (Maybe a)
361 {-# INLINE (!?) #-}
362 Facets{sElems = Unf step s} !? i = index_loop SPEC s i
363 where
364 index_loop !sPEC s i
365 = i `seq`
366 do
367 r <- step s
368 case r of
369 Yield x s' | i == 0 -> return (Just x)
370 | otherwise -> index_loop SPEC s' (i-1)
371 Skip s' -> index_loop SPEC s' i
372 Done -> return Nothing
373
374 -- Substreams
375 -- ----------
376
377 -- | Extract a substream of the given length starting at the given position.
378 slice :: Monad m => Int -- ^ starting index
379 -> Int -- ^ length
380 -> Facets m v a
381 -> Facets m v a
382 {-# INLINE slice #-}
383 slice i n s = take n (drop i s)
384
385 -- | All but the last element
386 init :: Monad m => Facets m v a -> Facets m v a
387 {-# INLINE_FUSED init #-}
388 init Facets{sElems = Unf step s, sSize = sz} = simple step' (Nothing, s) (sz - 1)
389 where
390 {-# INLINE_INNER step' #-}
391 step' (Nothing, s) = liftM (\r ->
392 case r of
393 Yield x s' -> Skip (Just x, s')
394 Skip s' -> Skip (Nothing, s')
395 Done -> EMPTY_STREAM "init"
396 ) (step s)
397
398 step' (Just x, s) = liftM (\r ->
399 case r of
400 Yield y s' -> Yield x (Just y, s')
401 Skip s' -> Skip (Just x, s')
402 Done -> Done
403 ) (step s)
404
405 -- | All but the first element
406 tail :: Monad m => Facets m v a -> Facets m v a
407 {-# INLINE_FUSED tail #-}
408 tail Facets{sElems = Unf step s, sSize = sz} = simple step' (Left s) (sz - 1)
409 where
410 {-# INLINE_INNER step' #-}
411 step' (Left s) = liftM (\r ->
412 case r of
413 Yield x s' -> Skip (Right s')
414 Skip s' -> Skip (Left s')
415 Done -> EMPTY_STREAM "tail"
416 ) (step s)
417
418 step' (Right s) = liftM (\r ->
419 case r of
420 Yield x s' -> Yield x (Right s')
421 Skip s' -> Skip (Right s')
422 Done -> Done
423 ) (step s)
424
425 -- | The first @n@ elements
426 take :: Monad m => Int -> Facets m v a -> Facets m v a
427 {-# INLINE_FUSED take #-}
428 take n Facets{sElems = Unf step s, sSize = sz}
429 = simple step' (s, 0) (smaller (Exact n) sz)
430 where
431 {-# INLINE_INNER step' #-}
432 step' (s, i) | i < n = liftM (\r ->
433 case r of
434 Yield x s' -> Yield x (s', i+1)
435 Skip s' -> Skip (s', i)
436 Done -> Done
437 ) (step s)
438 step' (s, i) = return Done
439
440 -- | All but the first @n@ elements
441 drop :: Monad m => Int -> Facets m v a -> Facets m v a
442 {-# INLINE_FUSED drop #-}
443 drop n Facets{sElems = Unf step s, sSize = sz}
444 = simple step' (s, Just n) (sz - Exact n)
445 where
446 {-# INLINE_INNER step' #-}
447 step' (s, Just i) | i > 0 = liftM (\r ->
448 case r of
449 Yield x s' -> Skip (s', Just (i-1))
450 Skip s' -> Skip (s', Just i)
451 Done -> Done
452 ) (step s)
453 | otherwise = return $ Skip (s, Nothing)
454
455 step' (s, Nothing) = liftM (\r ->
456 case r of
457 Yield x s' -> Yield x (s', Nothing)
458 Skip s' -> Skip (s', Nothing)
459 Done -> Done
460 ) (step s)
461
462 -- Mapping
463 -- -------
464
465 instance Monad m => Functor (Facets m v) where
466 {-# INLINE fmap #-}
467 fmap = map
468
469 -- | Map a function over a 'Facets'
470 map :: Monad m => (a -> b) -> Facets m v a -> Facets m v b
471 {-# INLINE map #-}
472 map f = mapM (return . f)
473
474
475 -- | Map a monadic function over a 'Facets'
476 mapM :: Monad m => (a -> m b) -> Facets m v a -> Facets m v b
477 {-# INLINE_FUSED mapM #-}
478 mapM f Facets{sElems = Unf step s, sSize = n} = simple step' s n
479 where
480 {-# INLINE_INNER step' #-}
481 step' s = do
482 r <- step s
483 case r of
484 Yield x s' -> liftM (`Yield` s') (f x)
485 Skip s' -> return (Skip s')
486 Done -> return Done
487
488 consume :: Monad m => Facets m v a -> m ()
489 {-# INLINE_FUSED consume #-}
490 consume Facets {sChunks = Unf step s} = consume_loop SPEC s
491 where
492 consume_loop !sPEC s
493 = do
494 r <- step s
495 case r of
496 Yield _ s' -> consume_loop SPEC s'
497 Skip s' -> consume_loop SPEC s'
498 Done -> return ()
499
500 -- | Execute a monadic action for each element of the 'Facets'
501 mapM_ :: Monad m => (a -> m b) -> Facets m v a -> m ()
502 {-# INLINE_FUSED mapM_ #-}
503 mapM_ m = consume . mapM m
504
505 -- | Transform a 'Facets' to use a different monad
506 trans :: (Monad m, Monad m') => (forall a. m a -> m' a)
507 -> Facets m v a -> Facets m' v a
508 {-# INLINE_FUSED trans #-}
509 trans f Facets{sElems = Unf step s, sSize = n} = simple (f . step) s n
510
511 unbox :: Monad m => Facets m v (Box a) -> Facets m v a
512 {-# INLINE_FUSED unbox #-}
513 unbox Facets{sElems = Unf step s, sSize = n} = simple step' s n
514 where
515 {-# INLINE_INNER step' #-}
516 step' s = do
517 r <- step s
518 case r of
519 Yield (Box x) s' -> return $ Yield x s'
520 Skip s' -> return $ Skip s'
521 Done -> return $ Done
522
523 -- Zipping
524 -- -------
525
526 -- | Pair each element in a 'Facets' with its index
527 indexed :: Monad m => Facets m v a -> Facets m v (Int,a)
528 {-# INLINE_FUSED indexed #-}
529 indexed Facets{sElems = Unf step s, sSize = n} = simple step' (s,0) n
530 where
531 {-# INLINE_INNER step' #-}
532 step' (s,i) = i `seq`
533 do
534 r <- step s
535 case r of
536 Yield x s' -> return $ Yield (i,x) (s', i+1)
537 Skip s' -> return $ Skip (s', i)
538 Done -> return Done
539
540 -- | Pair each element in a 'Facets' with its index, starting from the right
541 -- and counting down
542 indexedR :: Monad m => Int -> Facets m v a -> Facets m v (Int,a)
543 {-# INLINE_FUSED indexedR #-}
544 indexedR m Facets{sElems = Unf step s, sSize = n} = simple step' (s,m) n
545 where
546 {-# INLINE_INNER step' #-}
547 step' (s,i) = i `seq`
548 do
549 r <- step s
550 case r of
551 Yield x s' -> let i' = i-1
552 in
553 return $ Yield (i',x) (s', i')
554 Skip s' -> return $ Skip (s', i)
555 Done -> return Done
556
557 -- | Zip two 'Facets's with the given monadic function
558 zipWithM :: Monad m => (a -> b -> m c) -> Facets m v a -> Facets m v b -> Facets m v c
559 {-# INLINE_FUSED zipWithM #-}
560 zipWithM f Facets{sElems = Unf stepa sa, sSize = na}
561 Facets{sElems = Unf stepb sb, sSize = nb}
562 = simple step (sa, sb, Nothing) (smaller na nb)
563 where
564 {-# INLINE_INNER step #-}
565 step (sa, sb, Nothing) = liftM (\r ->
566 case r of
567 Yield x sa' -> Skip (sa', sb, Just x)
568 Skip sa' -> Skip (sa', sb, Nothing)
569 Done -> Done
570 ) (stepa sa)
571
572 step (sa, sb, Just x) = do
573 r <- stepb sb
574 case r of
575 Yield y sb' ->
576 do
577 z <- f x y
578 return $ Yield z (sa, sb', Nothing)
579 Skip sb' -> return $ Skip (sa, sb', Just x)
580 Done -> return $ Done
581
582 -- FIXME: This might expose an opportunity for inplace execution.
583 {-# RULES
584
585 "zipWithM xs xs [Vector.Facets]" forall f xs.
586 zipWithM f xs xs = mapM (\x -> f x x) xs
587
588 #-}
589
590 zipWithM_ :: Monad m => (a -> b -> m c) -> Facets m v a -> Facets m v b -> m ()
591 {-# INLINE zipWithM_ #-}
592 zipWithM_ f sa sb = consume (zipWithM f sa sb)
593
594 zipWith3M :: Monad m => (a -> b -> c -> m d) -> Facets m v a -> Facets m v b -> Facets m v c -> Facets m v d
595 {-# INLINE_FUSED zipWith3M #-}
596 zipWith3M f Facets{sElems = Unf stepa sa, sSize = na}
597 Facets{sElems = Unf stepb sb, sSize = nb}
598 Facets{sElems = Unf stepc sc, sSize = nc}
599 = simple step (sa, sb, sc, Nothing) (smaller na (smaller nb nc))
600 where
601 {-# INLINE_INNER step #-}
602 step (sa, sb, sc, Nothing) = do
603 r <- stepa sa
604 return $ case r of
605 Yield x sa' -> Skip (sa', sb, sc, Just (x, Nothing))
606 Skip sa' -> Skip (sa', sb, sc, Nothing)
607 Done -> Done
608
609 step (sa, sb, sc, Just (x, Nothing)) = do
610 r <- stepb sb
611 return $ case r of
612 Yield y sb' -> Skip (sa, sb', sc, Just (x, Just y))
613 Skip sb' -> Skip (sa, sb', sc, Just (x, Nothing))
614 Done -> Done
615
616 step (sa, sb, sc, Just (x, Just y)) = do
617 r <- stepc sc
618 case r of
619 Yield z sc' -> f x y z >>= (\res -> return $ Yield res (sa, sb, sc', Nothing))
620 Skip sc' -> return $ Skip (sa, sb, sc', Just (x, Just y))
621 Done -> return $ Done
622
623 zipWith4M :: Monad m => (a -> b -> c -> d -> m e)
624 -> Facets m v a -> Facets m v b -> Facets m v c -> Facets m v d
625 -> Facets m v e
626 {-# INLINE zipWith4M #-}
627 zipWith4M f sa sb sc sd
628 = zipWithM (\(a,b) (c,d) -> f a b c d) (zip sa sb) (zip sc sd)
629
630 zipWith5M :: Monad m => (a -> b -> c -> d -> e -> m f)
631 -> Facets m v a -> Facets m v b -> Facets m v c -> Facets m v d
632 -> Facets m v e -> Facets m v f
633 {-# INLINE zipWith5M #-}
634 zipWith5M f sa sb sc sd se
635 = zipWithM (\(a,b,c) (d,e) -> f a b c d e) (zip3 sa sb sc) (zip sd se)
636
637 zipWith6M :: Monad m => (a -> b -> c -> d -> e -> f -> m g)
638 -> Facets m v a -> Facets m v b -> Facets m v c -> Facets m v d
639 -> Facets m v e -> Facets m v f -> Facets m v g
640 {-# INLINE zipWith6M #-}
641 zipWith6M fn sa sb sc sd se sf
642 = zipWithM (\(a,b,c) (d,e,f) -> fn a b c d e f) (zip3 sa sb sc)
643 (zip3 sd se sf)
644
645 zipWith :: Monad m => (a -> b -> c) -> Facets m v a -> Facets m v b -> Facets m v c
646 {-# INLINE zipWith #-}
647 zipWith f = zipWithM (\a b -> return (f a b))
648
649 zipWith3 :: Monad m => (a -> b -> c -> d)
650 -> Facets m v a -> Facets m v b -> Facets m v c -> Facets m v d
651 {-# INLINE zipWith3 #-}
652 zipWith3 f = zipWith3M (\a b c -> return (f a b c))
653
654 zipWith4 :: Monad m => (a -> b -> c -> d -> e)
655 -> Facets m v a -> Facets m v b -> Facets m v c -> Facets m v d
656 -> Facets m v e
657 {-# INLINE zipWith4 #-}
658 zipWith4 f = zipWith4M (\a b c d -> return (f a b c d))
659
660 zipWith5 :: Monad m => (a -> b -> c -> d -> e -> f)
661 -> Facets m v a -> Facets m v b -> Facets m v c -> Facets m v d
662 -> Facets m v e -> Facets m v f
663 {-# INLINE zipWith5 #-}
664 zipWith5 f = zipWith5M (\a b c d e -> return (f a b c d e))
665
666 zipWith6 :: Monad m => (a -> b -> c -> d -> e -> f -> g)
667 -> Facets m v a -> Facets m v b -> Facets m v c -> Facets m v d
668 -> Facets m v e -> Facets m v f -> Facets m v g
669 {-# INLINE zipWith6 #-}
670 zipWith6 fn = zipWith6M (\a b c d e f -> return (fn a b c d e f))
671
672 zip :: Monad m => Facets m v a -> Facets m v b -> Facets m v (a,b)
673 {-# INLINE zip #-}
674 zip = zipWith (,)
675
676 zip3 :: Monad m => Facets m v a -> Facets m v b -> Facets m v c -> Facets m v (a,b,c)
677 {-# INLINE zip3 #-}
678 zip3 = zipWith3 (,,)
679
680 zip4 :: Monad m => Facets m v a -> Facets m v b -> Facets m v c -> Facets m v d
681 -> Facets m v (a,b,c,d)
682 {-# INLINE zip4 #-}
683 zip4 = zipWith4 (,,,)
684
685 zip5 :: Monad m => Facets m v a -> Facets m v b -> Facets m v c -> Facets m v d
686 -> Facets m v e -> Facets m v (a,b,c,d,e)
687 {-# INLINE zip5 #-}
688 zip5 = zipWith5 (,,,,)
689
690 zip6 :: Monad m => Facets m v a -> Facets m v b -> Facets m v c -> Facets m v d
691 -> Facets m v e -> Facets m v f -> Facets m v (a,b,c,d,e,f)
692 {-# INLINE zip6 #-}
693 zip6 = zipWith6 (,,,,,)
694
695 -- Comparisons
696 -- -----------
697
698 -- | Check if two 'Facets's are equal
699 eq :: (Monad m, Eq a) => Facets m v a -> Facets m v a -> m Bool
700 {-# INLINE_FUSED eq #-}
701 eq Facets{sElems = Unf step1 s1}
702 Facets{sElems = Unf step2 s2} = eq_loop0 SPEC s1 s2
703 where
704 eq_loop0 !sPEC s1 s2 = do
705 r <- step1 s1
706 case r of
707 Yield x s1' -> eq_loop1 SPEC x s1' s2
708 Skip s1' -> eq_loop0 SPEC s1' s2
709 Done -> eq_null s2
710
711 eq_loop1 !sPEC x s1 s2 = do
712 r <- step2 s2
713 case r of
714 Yield y s2'
715 | x == y -> eq_loop0 SPEC s1 s2'
716 | otherwise -> return False
717 Skip s2' -> eq_loop1 SPEC x s1 s2'
718 Done -> return False
719
720 eq_null s2 = do
721 r <- step2 s2
722 case r of
723 Yield _ _ -> return False
724 Skip s2' -> eq_null s2'
725 Done -> return True
726
727 -- | Lexicographically compare two 'Facets's
728 cmp :: (Monad m, Ord a) => Facets m v a -> Facets m v a -> m Ordering
729 {-# INLINE_FUSED cmp #-}
730 cmp Facets{sElems = Unf step1 s1}
731 Facets{sElems = Unf step2 s2} = cmp_loop0 SPEC s1 s2
732 where
733 cmp_loop0 !sPEC s1 s2 = do
734 r <- step1 s1
735 case r of
736 Yield x s1' -> cmp_loop1 SPEC x s1' s2
737 Skip s1' -> cmp_loop0 SPEC s1' s2
738 Done -> cmp_null s2
739
740 cmp_loop1 !sPEC x s1 s2 = do
741 r <- step2 s2
742 case r of
743 Yield y s2' -> case x `compare` y of
744 EQ -> cmp_loop0 SPEC s1 s2'
745 c -> return c
746 Skip s2' -> cmp_loop1 SPEC x s1 s2'
747 Done -> return GT
748
749 cmp_null s2 = do
750 r <- step2 s2
751 case r of
752 Yield _ _ -> return LT
753 Skip s2' -> cmp_null s2'
754 Done -> return EQ
755
756 -- Filtering
757 -- ---------
758
759 -- | Drop elements which do not satisfy the predicate
760 filter :: Monad m => (a -> Bool) -> Facets m v a -> Facets m v a
761 {-# INLINE filter #-}
762 filter f = filterM (return . f)
763
764 -- | Drop elements which do not satisfy the monadic predicate
765 filterM :: Monad m => (a -> m Bool) -> Facets m v a -> Facets m v a
766 {-# INLINE_FUSED filterM #-}
767 filterM f Facets{sElems = Unf step s, sSize = n} = simple step' s (toMax n)
768 where
769 {-# INLINE_INNER step' #-}
770 step' s = do
771 r <- step s
772 case r of
773 Yield x s' -> do
774 b <- f x
775 return $ if b then Yield x s'
776 else Skip s'
777 Skip s' -> return $ Skip s'
778 Done -> return $ Done
779
780 -- | Longest prefix of elements that satisfy the predicate
781 takeWhile :: Monad m => (a -> Bool) -> Facets m v a -> Facets m v a
782 {-# INLINE takeWhile #-}
783 takeWhile f = takeWhileM (return . f)
784
785 -- | Longest prefix of elements that satisfy the monadic predicate
786 takeWhileM :: Monad m => (a -> m Bool) -> Facets m v a -> Facets m v a
787 {-# INLINE_FUSED takeWhileM #-}
788 takeWhileM f Facets{sElems = Unf step s, sSize = n} = simple step' s (toMax n)
789 where
790 {-# INLINE_INNER step' #-}
791 step' s = do
792 r <- step s
793 case r of
794 Yield x s' -> do
795 b <- f x
796 return $ if b then Yield x s' else Done
797 Skip s' -> return $ Skip s'
798 Done -> return $ Done
799
800 -- | Drop the longest prefix of elements that satisfy the predicate
801 dropWhile :: Monad m => (a -> Bool) -> Facets m v a -> Facets m v a
802 {-# INLINE dropWhile #-}
803 dropWhile f = dropWhileM (return . f)
804
805 data DropWhile s a = DropWhile_Drop s | DropWhile_Yield a s | DropWhile_Next s
806
807 -- | Drop the longest prefix of elements that satisfy the monadic predicate
808 dropWhileM :: Monad m => (a -> m Bool) -> Facets m v a -> Facets m v a
809 {-# INLINE_FUSED dropWhileM #-}
810 dropWhileM f Facets{sElems = Unf step s, sSize = n}
811 = simple step' (DropWhile_Drop s) (toMax n)
812 where
813 -- NOTE: we jump through hoops here to have only one Yield; local data
814 -- declarations would be nice!
815
816 {-# INLINE_INNER step' #-}
817 step' (DropWhile_Drop s)
818 = do
819 r <- step s
820 case r of
821 Yield x s' -> do
822 b <- f x
823 return $ if b then Skip (DropWhile_Drop s')
824 else Skip (DropWhile_Yield x s')
825 Skip s' -> return $ Skip (DropWhile_Drop s')
826 Done -> return $ Done
827
828 step' (DropWhile_Yield x s) = return $ Yield x (DropWhile_Next s)
829
830 step' (DropWhile_Next s)
831 = liftM (\r ->
832 case r of
833 Yield x s' -> Skip (DropWhile_Yield x s')
834 Skip s' -> Skip (DropWhile_Next s')
835 Done -> Done
836 ) (step s)
837
838 -- Searching
839 -- ---------
840
841 infix 4 `elem`
842 -- | Check whether the 'Facets' contains an element
843 elem :: (Monad m, Eq a) => a -> Facets m v a -> m Bool
844 {-# INLINE_FUSED elem #-}
845 elem x Facets{sElems = Unf step s} = elem_loop SPEC s
846 where
847 elem_loop !sPEC s
848 = do
849 r <- step s
850 case r of
851 Yield y s' | x == y -> return True
852 | otherwise -> elem_loop SPEC s'
853 Skip s' -> elem_loop SPEC s'
854 Done -> return False
855
856 infix 4 `notElem`
857 -- | Inverse of `elem`
858 notElem :: (Monad m, Eq a) => a -> Facets m v a -> m Bool
859 {-# INLINE notElem #-}
860 notElem x s = liftM not (elem x s)
861
862 -- | Yield 'Just' the first element that satisfies the predicate or 'Nothing'
863 -- if no such element exists.
864 find :: Monad m => (a -> Bool) -> Facets m v a -> m (Maybe a)
865 {-# INLINE find #-}
866 find f = findM (return . f)
867
868 -- | Yield 'Just' the first element that satisfies the monadic predicate or
869 -- 'Nothing' if no such element exists.
870 findM :: Monad m => (a -> m Bool) -> Facets m v a -> m (Maybe a)
871 {-# INLINE_FUSED findM #-}
872 findM f Facets{sElems = Unf step s} = find_loop SPEC s
873 where
874 find_loop !sPEC s
875 = do
876 r <- step s
877 case r of
878 Yield x s' -> do
879 b <- f x
880 if b then return $ Just x
881 else find_loop SPEC s'
882 Skip s' -> find_loop SPEC s'
883 Done -> return Nothing
884
885 -- | Yield 'Just' the index of the first element that satisfies the predicate
886 -- or 'Nothing' if no such element exists.
887 findIndex :: Monad m => (a -> Bool) -> Facets m v a -> m (Maybe Int)
888 {-# INLINE_FUSED findIndex #-}
889 findIndex f = findIndexM (return . f)
890
891 -- | Yield 'Just' the index of the first element that satisfies the monadic
892 -- predicate or 'Nothing' if no such element exists.
893 findIndexM :: Monad m => (a -> m Bool) -> Facets m v a -> m (Maybe Int)
894 {-# INLINE_FUSED findIndexM #-}
895 findIndexM f Facets{sElems = Unf step s} = findIndex_loop SPEC s 0
896 where
897 findIndex_loop !sPEC s i
898 = do
899 r <- step s
900 case r of
901 Yield x s' -> do
902 b <- f x
903 if b then return $ Just i
904 else findIndex_loop SPEC s' (i+1)
905 Skip s' -> findIndex_loop SPEC s' i
906 Done -> return Nothing
907
908 -- Folding
909 -- -------
910
911 -- | Left fold
912 foldl :: Monad m => (a -> b -> a) -> a -> Facets m v b -> m a
913 {-# INLINE foldl #-}
914 foldl f = foldlM (\a b -> return (f a b))
915
916 -- | Left fold with a monadic operator
917 foldlM :: Monad m => (a -> b -> m a) -> a -> Facets m v b -> m a
918 {-# INLINE_FUSED foldlM #-}
919 foldlM m z Facets{sElems = Unf step s} = foldlM_loop SPEC z s
920 where
921 foldlM_loop !sPEC z s
922 = do
923 r <- step s
924 case r of
925 Yield x s' -> do { z' <- m z x; foldlM_loop SPEC z' s' }
926 Skip s' -> foldlM_loop SPEC z s'
927 Done -> return z
928
929 -- | Same as 'foldlM'
930 foldM :: Monad m => (a -> b -> m a) -> a -> Facets m v b -> m a
931 {-# INLINE foldM #-}
932 foldM = foldlM
933
934 vfoldlM :: Monad m => (a -> Chunk v b -> m a) -> a -> Facets m v b -> m a
935 {-# INLINE_FUSED vfoldlM #-}
936 vfoldlM f z Facets{sChunks = Unf step s} = vfoldlM_loop SPEC z s
937 where
938 vfoldlM_loop !sPEC z s
939 = do
940 r <- step s
941 case r of
942 Yield x s' -> do { z' <- f z x; vfoldlM_loop SPEC z' s' }
943 Skip s' -> vfoldlM_loop SPEC z s'
944 Done -> return z
945
946 -- | Left fold over a non-empty 'Facets'
947 foldl1 :: Monad m => (a -> a -> a) -> Facets m v a -> m a
948 {-# INLINE foldl1 #-}
949 foldl1 f = foldl1M (\a b -> return (f a b))
950
951 -- | Left fold over a non-empty 'Facets' with a monadic operator
952 foldl1M :: Monad m => (a -> a -> m a) -> Facets m v a -> m a
953 {-# INLINE_FUSED foldl1M #-}
954 foldl1M f Facets{sElems = Unf step s, sSize = sz} = foldl1M_loop SPEC s
955 where
956 foldl1M_loop !sPEC s
957 = do
958 r <- step s
959 case r of
960 Yield x s' -> foldlM f x (simple step s' (sz - 1))
961 Skip s' -> foldl1M_loop SPEC s'
962 Done -> EMPTY_STREAM "foldl1M"
963
964 -- | Same as 'foldl1M'
965 fold1M :: Monad m => (a -> a -> m a) -> Facets m v a -> m a
966 {-# INLINE fold1M #-}
967 fold1M = foldl1M
968
969 -- | Left fold with a strict accumulator
970 foldl' :: Monad m => (a -> b -> a) -> a -> Facets m v b -> m a
971 {-# INLINE foldl' #-}
972 foldl' f = foldlM' (\a b -> return (f a b))
973
974 vfoldl' :: Monad m => (a -> Chunk v b -> a) -> a -> Facets m v b -> m a
975 {-# INLINE vfoldl' #-}
976 vfoldl' f = vfoldlM' (\a b -> return (f a b))
977
978 -- | Left fold with a strict accumulator and a monadic operator
979 foldlM' :: Monad m => (a -> b -> m a) -> a -> Facets m v b -> m a
980 {-# INLINE_FUSED foldlM' #-}
981 foldlM' m z Facets{sElems = Unf step s} = foldlM'_loop SPEC z s
982 where
983 foldlM'_loop !sPEC z s
984 = z `seq`
985 do
986 r <- step s
987 case r of
988 Yield x s' -> do { z' <- m z x; foldlM'_loop SPEC z' s' }
989 Skip s' -> foldlM'_loop SPEC z s'
990 Done -> return z
991
992 -- | Same as 'foldlM''
993 foldM' :: Monad m => (a -> b -> m a) -> a -> Facets m v b -> m a
994 {-# INLINE foldM' #-}
995 foldM' = foldlM'
996
997 vfoldlM' :: Monad m => (a -> Chunk v b -> m a) -> a -> Facets m v b -> m a
998 {-# INLINE_FUSED vfoldlM' #-}
999 vfoldlM' f z Facets{sChunks = Unf step s} = vfoldlM'_loop SPEC z s
1000 where
1001 vfoldlM'_loop !sPEC z s
1002 = z `seq` do
1003 r <- step s
1004 case r of
1005 Yield x s' -> do { z' <- f z x; vfoldlM'_loop SPEC z' s' }
1006 Skip s' -> vfoldlM'_loop SPEC z s'
1007 Done -> return z
1008
1009 -- | Left fold over a non-empty 'Facets' with a strict accumulator
1010 foldl1' :: Monad m => (a -> a -> a) -> Facets m v a -> m a
1011 {-# INLINE foldl1' #-}
1012 foldl1' f = foldl1M' (\a b -> return (f a b))
1013
1014 -- | Left fold over a non-empty 'Facets' with a strict accumulator and a
1015 -- monadic operator
1016 foldl1M' :: Monad m => (a -> a -> m a) -> Facets m v a -> m a
1017 {-# INLINE_FUSED foldl1M' #-}
1018 foldl1M' f Facets{sElems = Unf step s, sSize = sz} = foldl1M'_loop SPEC s
1019 where
1020 foldl1M'_loop !sPEC s
1021 = do
1022 r <- step s
1023 case r of
1024 Yield x s' -> foldlM' f x (simple step s' (sz - 1))
1025 Skip s' -> foldl1M'_loop SPEC s'
1026 Done -> EMPTY_STREAM "foldl1M'"
1027
1028 -- | Same as 'foldl1M''
1029 fold1M' :: Monad m => (a -> a -> m a) -> Facets m v a -> m a
1030 {-# INLINE fold1M' #-}
1031 fold1M' = foldl1M'
1032
1033 -- | Right fold
1034 foldr :: Monad m => (a -> b -> b) -> b -> Facets m v a -> m b
1035 {-# INLINE foldr #-}
1036 foldr f = foldrM (\a b -> return (f a b))
1037
1038 -- | Right fold with a monadic operator
1039 foldrM :: Monad m => (a -> b -> m b) -> b -> Facets m v a -> m b
1040 {-# INLINE_FUSED foldrM #-}
1041 foldrM f z Facets{sElems = Unf step s} = foldrM_loop SPEC s
1042 where
1043 foldrM_loop !sPEC s
1044 = do
1045 r <- step s
1046 case r of
1047 Yield x s' -> f x =<< foldrM_loop SPEC s'
1048 Skip s' -> foldrM_loop SPEC s'
1049 Done -> return z
1050
1051 -- | Right fold over a non-empty stream
1052 foldr1 :: Monad m => (a -> a -> a) -> Facets m v a -> m a
1053 {-# INLINE foldr1 #-}
1054 foldr1 f = foldr1M (\a b -> return (f a b))
1055
1056 -- | Right fold over a non-empty stream with a monadic operator
1057 foldr1M :: Monad m => (a -> a -> m a) -> Facets m v a -> m a
1058 {-# INLINE_FUSED foldr1M #-}
1059 foldr1M f Facets{sElems = Unf step s} = foldr1M_loop0 SPEC s
1060 where
1061 foldr1M_loop0 !sPEC s
1062 = do
1063 r <- step s
1064 case r of
1065 Yield x s' -> foldr1M_loop1 SPEC x s'
1066 Skip s' -> foldr1M_loop0 SPEC s'
1067 Done -> EMPTY_STREAM "foldr1M"
1068
1069 foldr1M_loop1 !sPEC x s
1070 = do
1071 r <- step s
1072 case r of
1073 Yield y s' -> f x =<< foldr1M_loop1 SPEC y s'
1074 Skip s' -> foldr1M_loop1 SPEC x s'
1075 Done -> return x
1076
1077 -- Specialised folds
1078 -- -----------------
1079
1080 and :: Monad m => Facets m v Bool -> m Bool
1081 {-# INLINE_FUSED and #-}
1082 and Facets{sElems = Unf step s} = and_loop SPEC s
1083 where
1084 and_loop !sPEC s
1085 = do
1086 r <- step s
1087 case r of
1088 Yield False _ -> return False
1089 Yield True s' -> and_loop SPEC s'
1090 Skip s' -> and_loop SPEC s'
1091 Done -> return True
1092
1093 or :: Monad m => Facets m v Bool -> m Bool
1094 {-# INLINE_FUSED or #-}
1095 or Facets{sElems = Unf step s} = or_loop SPEC s
1096 where
1097 or_loop !sPEC s
1098 = do
1099 r <- step s
1100 case r of
1101 Yield False s' -> or_loop SPEC s'
1102 Yield True _ -> return True
1103 Skip s' -> or_loop SPEC s'
1104 Done -> return False
1105
1106 concatMap :: Monad m => (a -> Facets m v b) -> Facets m v a -> Facets m v b
1107 {-# INLINE concatMap #-}
1108 concatMap f = concatMapM (return . f)
1109
1110 concatMapM :: Monad m => (a -> m (Facets m v b)) -> Facets m v a -> Facets m v b
1111 {-# INLINE_FUSED concatMapM #-}
1112 concatMapM f Facets{sElems = Unf step s} = simple concatMap_go (Left s) Unknown
1113 where
1114 concatMap_go (Left s) = do
1115 r <- step s
1116 case r of
1117 Yield a s' -> do
1118 b_stream <- f a
1119 return $ Skip (Right (b_stream, s'))
1120 Skip s' -> return $ Skip (Left s')
1121 Done -> return Done
1122 concatMap_go (Right (Facets{sElems = Unf inner_step inner_s, sSize = sz}, s)) = do
1123 r <- inner_step inner_s
1124 case r of
1125 Yield b inner_s' -> return $ Yield b (Right (simple inner_step inner_s' sz, s))
1126 Skip inner_s' -> return $ Skip (Right (simple inner_step inner_s' sz, s))
1127 Done -> return $ Skip (Left s)
1128
1129 -- | Create a 'Facets' of values from a 'Facets' of streamable things
1130 flatten :: Monad m => (a -> m s) -> (s -> m (Step s b)) -> Size
1131 -> Facets m v a -> Facets m v b
1132 {-# INLINE_FUSED flatten #-}
1133 flatten mk istep sz Facets{sElems = Unf ostep t} = simple step (Left t) sz
1134 where
1135 {-# INLINE_INNER step #-}
1136 step (Left t) = do
1137 r <- ostep t
1138 case r of
1139 Yield a t' -> do
1140 s <- mk a
1141 s `seq` return (Skip (Right (s,t')))
1142 Skip t' -> return $ Skip (Left t')
1143 Done -> return $ Done
1144
1145
1146 step (Right (s,t)) = do
1147 r <- istep s
1148 case r of
1149 Yield x s' -> return $ Yield x (Right (s',t))
1150 Skip s' -> return $ Skip (Right (s',t))
1151 Done -> return $ Skip (Left t)
1152
1153 -- Unfolding
1154 -- ---------
1155
1156 -- | Unfold
1157 unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Facets m u a
1158 {-# INLINE_FUSED unfoldr #-}
1159 unfoldr f = unfoldrM (return . f)
1160
1161 -- | Unfold with a monadic function
1162 unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Facets m u a
1163 {-# INLINE_FUSED unfoldrM #-}
1164 unfoldrM f s = simple step s Unknown
1165 where
1166 {-# INLINE_INNER step #-}
1167 step s = liftM (\r ->
1168 case r of
1169 Just (x, s') -> Yield x s'
1170 Nothing -> Done
1171 ) (f s)
1172
1173 -- | Unfold at most @n@ elements
1174 unfoldrN :: Monad m => Int -> (s -> Maybe (a, s)) -> s -> Facets m u a
1175 {-# INLINE_FUSED unfoldrN #-}
1176 unfoldrN n f = unfoldrNM n (return . f)
1177
1178 -- | Unfold at most @n@ elements with a monadic functions
1179 unfoldrNM :: Monad m => Int -> (s -> m (Maybe (a, s))) -> s -> Facets m u a
1180 {-# INLINE_FUSED unfoldrNM #-}
1181 unfoldrNM n f s = simple step (s,n) (Max (delay_inline max n 0))
1182 where
1183 {-# INLINE_INNER step #-}
1184 step (s,n) | n <= 0 = return Done
1185 | otherwise = liftM (\r ->
1186 case r of
1187 Just (x,s') -> Yield x (s',n-1)
1188 Nothing -> Done
1189 ) (f s)
1190
1191 -- | Apply monadic function n times to value. Zeroth element is original value.
1192 iterateNM :: Monad m => Int -> (a -> m a) -> a -> Facets m u a
1193 {-# INLINE_FUSED iterateNM #-}
1194 iterateNM n f x0 = simple step (x0,n) (Exact (delay_inline max n 0))
1195 where
1196 {-# INLINE_INNER step #-}
1197 step (x,i) | i <= 0 = return Done
1198 | i == n = return $ Yield x (x,i-1)
1199 | otherwise = do a <- f x
1200 return $ Yield a (a,i-1)
1201
1202 -- | Apply function n times to value. Zeroth element is original value.
1203 iterateN :: Monad m => Int -> (a -> a) -> a -> Facets m u a
1204 {-# INLINE_FUSED iterateN #-}
1205 iterateN n f x0 = iterateNM n (return . f) x0
1206
1207 -- Scans
1208 -- -----
1209
1210 -- | Prefix scan
1211 prescanl :: Monad m => (a -> b -> a) -> a -> Facets m v b -> Facets m v a
1212 {-# INLINE prescanl #-}
1213 prescanl f = prescanlM (\a b -> return (f a b))
1214
1215 -- | Prefix scan with a monadic operator
1216 prescanlM :: Monad m => (a -> b -> m a) -> a -> Facets m v b -> Facets m v a
1217 {-# INLINE_FUSED prescanlM #-}
1218 prescanlM f z Facets{sElems = Unf step s, sSize = sz} = simple step' (s,z) sz
1219 where
1220 {-# INLINE_INNER step' #-}
1221 step' (s,x) = do
1222 r <- step s
1223 case r of
1224 Yield y s' -> do
1225 z <- f x y
1226 return $ Yield x (s', z)
1227 Skip s' -> return $ Skip (s', x)
1228 Done -> return Done
1229
1230 -- | Prefix scan with strict accumulator
1231 prescanl' :: Monad m => (a -> b -> a) -> a -> Facets m v b -> Facets m v a
1232 {-# INLINE prescanl' #-}
1233 prescanl' f = prescanlM' (\a b -> return (f a b))
1234
1235 -- | Prefix scan with strict accumulator and a monadic operator
1236 prescanlM' :: Monad m => (a -> b -> m a) -> a -> Facets m v b -> Facets m v a
1237 {-# INLINE_FUSED prescanlM' #-}
1238 prescanlM' f z Facets{sElems = Unf step s, sSize = sz} = simple step' (s,z) sz
1239 where
1240 {-# INLINE_INNER step' #-}
1241 step' (s,x) = x `seq`
1242 do
1243 r <- step s
1244 case r of
1245 Yield y s' -> do
1246 z <- f x y
1247 return $ Yield x (s', z)
1248 Skip s' -> return $ Skip (s', x)
1249 Done -> return Done
1250
1251 -- | Suffix scan
1252 postscanl :: Monad m => (a -> b -> a) -> a -> Facets m v b -> Facets m v a
1253 {-# INLINE postscanl #-}
1254 postscanl f = postscanlM (\a b -> return (f a b))
1255
1256 -- | Suffix scan with a monadic operator
1257 postscanlM :: Monad m => (a -> b -> m a) -> a -> Facets m v b -> Facets m v a
1258 {-# INLINE_FUSED postscanlM #-}
1259 postscanlM f z Facets{sElems = Unf step s, sSize = sz} = simple step' (s,z) sz
1260 where
1261 {-# INLINE_INNER step' #-}
1262 step' (s,x) = do
1263 r <- step s
1264 case r of
1265 Yield y s' -> do
1266 z <- f x y
1267 return $ Yield z (s',z)
1268 Skip s' -> return $ Skip (s',x)
1269 Done -> return Done
1270
1271 -- | Suffix scan with strict accumulator
1272 postscanl' :: Monad m => (a -> b -> a) -> a -> Facets m v b -> Facets m v a
1273 {-# INLINE postscanl' #-}
1274 postscanl' f = postscanlM' (\a b -> return (f a b))
1275
1276 -- | Suffix scan with strict acccumulator and a monadic operator
1277 postscanlM' :: Monad m => (a -> b -> m a) -> a -> Facets m v b -> Facets m v a
1278 {-# INLINE_FUSED postscanlM' #-}
1279 postscanlM' f z Facets{sElems = Unf step s, sSize = sz}
1280 = z `seq` simple step' (s,z) sz
1281 where
1282 {-# INLINE_INNER step' #-}
1283 step' (s,x) = x `seq`
1284 do
1285 r <- step s
1286 case r of
1287 Yield y s' -> do
1288 z <- f x y
1289 z `seq` return (Yield z (s',z))
1290 Skip s' -> return $ Skip (s',x)
1291 Done -> return Done
1292
1293 -- | Haskell-style scan
1294 scanl :: Monad m => (a -> b -> a) -> a -> Facets m v b -> Facets m v a
1295 {-# INLINE scanl #-}
1296 scanl f = scanlM (\a b -> return (f a b))
1297
1298 -- | Haskell-style scan with a monadic operator
1299 scanlM :: Monad m => (a -> b -> m a) -> a -> Facets m v b -> Facets m v a
1300 {-# INLINE scanlM #-}
1301 scanlM f z s = z `cons` postscanlM f z s
1302
1303 -- | Haskell-style scan with strict accumulator
1304 scanl' :: Monad m => (a -> b -> a) -> a -> Facets m v b -> Facets m v a
1305 {-# INLINE scanl' #-}
1306 scanl' f = scanlM' (\a b -> return (f a b))
1307
1308 -- | Haskell-style scan with strict accumulator and a monadic operator
1309 scanlM' :: Monad m => (a -> b -> m a) -> a -> Facets m v b -> Facets m v a
1310 {-# INLINE scanlM' #-}
1311 scanlM' f z s = z `seq` (z `cons` postscanlM f z s)
1312
1313 -- | Scan over a non-empty 'Facets'
1314 scanl1 :: Monad m => (a -> a -> a) -> Facets m v a -> Facets m v a
1315 {-# INLINE scanl1 #-}
1316 scanl1 f = scanl1M (\x y -> return (f x y))
1317
1318 -- | Scan over a non-empty 'Facets' with a monadic operator
1319 scanl1M :: Monad m => (a -> a -> m a) -> Facets m v a -> Facets m v a
1320 {-# INLINE_FUSED scanl1M #-}
1321 scanl1M f Facets{sElems = Unf step s, sSize = sz} = simple step' (s, Nothing) sz
1322 where
1323 {-# INLINE_INNER step' #-}
1324 step' (s, Nothing) = do
1325 r <- step s
1326 case r of
1327 Yield x s' -> return $ Yield x (s', Just x)
1328 Skip s' -> return $ Skip (s', Nothing)
1329 Done -> EMPTY_STREAM "scanl1M"
1330
1331 step' (s, Just x) = do
1332 r <- step s
1333 case r of
1334 Yield y s' -> do
1335 z <- f x y
1336 return $ Yield z (s', Just z)
1337 Skip s' -> return $ Skip (s', Just x)
1338 Done -> return Done
1339
1340 -- | Scan over a non-empty 'Facets' with a strict accumulator
1341 scanl1' :: Monad m => (a -> a -> a) -> Facets m v a -> Facets m v a
1342 {-# INLINE scanl1' #-}
1343 scanl1' f = scanl1M' (\x y -> return (f x y))
1344
1345 -- | Scan over a non-empty 'Facets' with a strict accumulator and a monadic
1346 -- operator
1347 scanl1M' :: Monad m => (a -> a -> m a) -> Facets m v a -> Facets m v a
1348 {-# INLINE_FUSED scanl1M' #-}
1349 scanl1M' f Facets{sElems = Unf step s, sSize = sz}
1350 = simple step' (s, Nothing) sz
1351 where
1352 {-# INLINE_INNER step' #-}
1353 step' (s, Nothing) = do
1354 r <- step s
1355 case r of
1356 Yield x s' -> x `seq` return (Yield x (s', Just x))
1357 Skip s' -> return $ Skip (s', Nothing)
1358 Done -> EMPTY_STREAM "scanl1M"
1359
1360 step' (s, Just x) = x `seq`
1361 do
1362 r <- step s
1363 case r of
1364 Yield y s' -> do
1365 z <- f x y
1366 z `seq` return (Yield z (s', Just z))
1367 Skip s' -> return $ Skip (s', Just x)
1368 Done -> return Done
1369
1370 -- Enumerations
1371 -- ------------
1372
1373 -- The Enum class is broken for this, there just doesn't seem to be a
1374 -- way to implement this generically. We have to specialise for as many types
1375 -- as we can but this doesn't help in polymorphic loops.
1376
1377 -- | Yield a 'Facets' of the given length containing the values @x@, @x+y@,
1378 -- @x+y+y@ etc.
1379 enumFromStepN :: (Num a, Monad m) => a -> a -> Int -> Facets m v a
1380 {-# INLINE_FUSED enumFromStepN #-}
1381 enumFromStepN x y n = x `seq` y `seq` n `seq`
1382 simple step (x,n) (Exact (delay_inline max n 0))
1383 where
1384 {-# INLINE_INNER step #-}
1385 step (x,n) | n > 0 = return $ Yield x (x+y,n-1)
1386 | otherwise = return $ Done
1387
1388 -- | Enumerate values
1389 --
1390 -- /WARNING:/ This operation can be very inefficient. If at all possible, use
1391 -- 'enumFromStepN' instead.
1392 enumFromTo :: (Enum a, Monad m) => a -> a -> Facets m v a
1393 {-# INLINE_FUSED enumFromTo #-}
1394 enumFromTo x y = fromList [x .. y]
1395
1396 -- NOTE: We use (x+1) instead of (succ x) below because the latter checks for
1397 -- overflow which can't happen here.
1398
1399 -- FIXME: add "too large" test for Int
1400 enumFromTo_small :: (Integral a, Monad m) => a -> a -> Facets m v a
1401 {-# INLINE_FUSED enumFromTo_small #-}
1402 enumFromTo_small x y = x `seq` y `seq` simple step x (Exact n)
1403 where
1404 n = delay_inline max (fromIntegral y - fromIntegral x + 1) 0
1405
1406 {-# INLINE_INNER step #-}
1407 step x | x <= y = return $ Yield x (x+1)
1408 | otherwise = return $ Done
1409
1410 {-# RULES
1411
1412 "enumFromTo<Int8> [Facets]"
1413 enumFromTo = enumFromTo_small :: Monad m => Int8 -> Int8 -> Facets m v Int8
1414
1415 "enumFromTo<Int16> [Facets]"
1416 enumFromTo = enumFromTo_small :: Monad m => Int16 -> Int16 -> Facets m v Int16
1417
1418 "enumFromTo<Word8> [Facets]"
1419 enumFromTo = enumFromTo_small :: Monad m => Word8 -> Word8 -> Facets m v Word8
1420
1421 "enumFromTo<Word16> [Facets]"
1422 enumFromTo = enumFromTo_small :: Monad m => Word16 -> Word16 -> Facets m v Word16
1423
1424 #-}
1425
1426 #if WORD_SIZE_IN_BITS > 32
1427
1428 {-# RULES
1429
1430 "enumFromTo<Int32> [Facets]"
1431 enumFromTo = enumFromTo_small :: Monad m => Int32 -> Int32 -> Facets m v Int32
1432
1433 "enumFromTo<Word32> [Facets]"
1434 enumFromTo = enumFromTo_small :: Monad m => Word32 -> Word32 -> Facets m v Word32
1435
1436 #-}
1437
1438 #endif
1439
1440 -- NOTE: We could implement a generic "too large" test:
1441 --
1442 -- len x y | x > y = 0
1443 -- | n > 0 && n <= fromIntegral (maxBound :: Int) = fromIntegral n
1444 -- | otherwise = error
1445 -- where
1446 -- n = y-x+1
1447 --
1448 -- Alas, GHC won't eliminate unnecessary comparisons (such as n >= 0 for
1449 -- unsigned types). See http://hackage.haskell.org/trac/ghc/ticket/3744
1450 --
1451
1452 enumFromTo_int :: forall m v. Monad m => Int -> Int -> Facets m v Int
1453 {-# INLINE_FUSED enumFromTo_int #-}
1454 enumFromTo_int x y = x `seq` y `seq` simple step x (Exact (len x y))
1455 where
1456 {-# INLINE [0] len #-}
1457 len :: Int -> Int -> Int
1458 len x y | x > y = 0
1459 | otherwise = BOUNDS_CHECK(check) "enumFromTo" "vector too large"
1460 (n > 0)
1461 $ n
1462 where
1463 n = y-x+1
1464
1465 {-# INLINE_INNER step #-}
1466 step x | x <= y = return $ Yield x (x+1)
1467 | otherwise = return $ Done
1468
1469 enumFromTo_intlike :: (Integral a, Monad m) => a -> a -> Facets m v a
1470 {-# INLINE_FUSED enumFromTo_intlike #-}
1471 enumFromTo_intlike x y = x `seq` y `seq` simple step x (Exact (len x y))
1472 where
1473 {-# INLINE [0] len #-}
1474 len x y | x > y = 0
1475 | otherwise = BOUNDS_CHECK(check) "enumFromTo" "vector too large"
1476 (n > 0)
1477 $ fromIntegral n
1478 where
1479 n = y-x+1
1480
1481 {-# INLINE_INNER step #-}
1482 step x | x <= y = return $ Yield x (x+1)
1483 | otherwise = return $ Done
1484
1485 {-# RULES
1486
1487 "enumFromTo<Int> [Facets]"
1488 enumFromTo = enumFromTo_int :: Monad m => Int -> Int -> Facets m v Int
1489
1490 #if WORD_SIZE_IN_BITS > 32
1491
1492 "enumFromTo<Int64> [Facets]"
1493 enumFromTo = enumFromTo_intlike :: Monad m => Int64 -> Int64 -> Facets m v Int64
1494
1495 #else
1496
1497 "enumFromTo<Int32> [Facets]"
1498 enumFromTo = enumFromTo_intlike :: Monad m => Int32 -> Int32 -> Facets m v Int32
1499
1500 #endif
1501
1502 #-}
1503
1504 enumFromTo_big_word :: (Integral a, Monad m) => a -> a -> Facets m v a
1505 {-# INLINE_FUSED enumFromTo_big_word #-}
1506 enumFromTo_big_word x y = x `seq` y `seq` simple step x (Exact (len x y))
1507 where
1508 {-# INLINE [0] len #-}
1509 len x y | x > y = 0
1510 | otherwise = BOUNDS_CHECK(check) "enumFromTo" "vector too large"
1511 (n < fromIntegral (maxBound :: Int))
1512 $ fromIntegral (n+1)
1513 where
1514 n = y-x
1515
1516 {-# INLINE_INNER step #-}
1517 step x | x <= y = return $ Yield x (x+1)
1518 | otherwise = return $ Done
1519
1520 {-# RULES
1521
1522 "enumFromTo<Word> [Facets]"
1523 enumFromTo = enumFromTo_big_word :: Monad m => Word -> Word -> Facets m v Word
1524
1525 "enumFromTo<Word64> [Facets]"
1526 enumFromTo = enumFromTo_big_word
1527 :: Monad m => Word64 -> Word64 -> Facets m v Word64
1528
1529 #if WORD_SIZE_IN_BITS == 32
1530
1531 "enumFromTo<Word32> [Facets]"
1532 enumFromTo = enumFromTo_big_word
1533 :: Monad m => Word32 -> Word32 -> Facets m v Word32
1534
1535 #endif
1536
1537 "enumFromTo<Integer> [Facets]"
1538 enumFromTo = enumFromTo_big_word
1539 :: Monad m => Integer -> Integer -> Facets m v Integer
1540
1541 #-}
1542
1543 -- FIXME: the "too large" test is totally wrong
1544 enumFromTo_big_int :: (Integral a, Monad m) => a -> a -> Facets m v a
1545 {-# INLINE_FUSED enumFromTo_big_int #-}
1546 enumFromTo_big_int x y = x `seq` y `seq` simple step x (Exact (len x y))
1547 where
1548 {-# INLINE [0] len #-}
1549 len x y | x > y = 0
1550 | otherwise = BOUNDS_CHECK(check) "enumFromTo" "vector too large"
1551 (n > 0 && n <= fromIntegral (maxBound :: Int))
1552 $ fromIntegral n
1553 where
1554 n = y-x+1
1555
1556 {-# INLINE_INNER step #-}
1557 step x | x <= y = return $ Yield x (x+1)
1558 | otherwise = return $ Done
1559
1560 #if WORD_SIZE_IN_BITS > 32
1561
1562 {-# RULES
1563
1564 "enumFromTo<Int64> [Facets]"
1565 enumFromTo = enumFromTo_big :: Monad m => Int64 -> Int64 -> Facets m v Int64
1566
1567 #-}
1568
1569 #endif
1570
1571 enumFromTo_char :: Monad m => Char -> Char -> Facets m v Char
1572 {-# INLINE_FUSED enumFromTo_char #-}
1573 enumFromTo_char x y = x `seq` y `seq` simple step xn (Exact n)
1574 where
1575 xn = ord x
1576 yn = ord y
1577
1578 n = delay_inline max 0 (yn - xn + 1)
1579
1580 {-# INLINE_INNER step #-}
1581 step xn | xn <= yn = return $ Yield (unsafeChr xn) (xn+1)
1582 | otherwise = return $ Done
1583
1584 {-# RULES
1585
1586 "enumFromTo<Char> [Facets]"
1587 enumFromTo = enumFromTo_char
1588
1589 #-}
1590
1591 ------------------------------------------------------------------------
1592
1593 -- Specialise enumFromTo for Float and Double.
1594 -- Also, try to do something about pairs?
1595
1596 enumFromTo_double :: (Monad m, Ord a, RealFrac a) => a -> a -> Facets m v a
1597 {-# INLINE_FUSED enumFromTo_double #-}
1598 enumFromTo_double n m = n `seq` m `seq` simple step n (Max (len n m))
1599 where
1600 lim = m + 1/2 -- important to float out
1601
1602 {-# INLINE [0] len #-}
1603 len x y | x > y = 0
1604 | otherwise = BOUNDS_CHECK(check) "enumFromTo" "vector too large"
1605 (n > 0)
1606 $ fromIntegral n
1607 where
1608 n = truncate (y-x)+2
1609
1610 {-# INLINE_INNER step #-}
1611 step x | x <= lim = return $ Yield x (x+1)
1612 | otherwise = return $ Done
1613
1614 {-# RULES
1615
1616 "enumFromTo<Double> [Facets]"
1617 enumFromTo = enumFromTo_double :: Monad m => Double -> Double -> Facets m v Double
1618
1619 "enumFromTo<Float> [Facets]"
1620 enumFromTo = enumFromTo_double :: Monad m => Float -> Float -> Facets m v Float
1621
1622 #-}
1623
1624 ------------------------------------------------------------------------
1625
1626 -- | Enumerate values with a given step.
1627 --
1628 -- /WARNING:/ This operation is very inefficient. If at all possible, use
1629 -- 'enumFromStepN' instead.
1630 enumFromThenTo :: (Enum a, Monad m) => a -> a -> a -> Facets m v a
1631 {-# INLINE_FUSED enumFromThenTo #-}
1632 enumFromThenTo x y z = fromList [x, y .. z]
1633
1634 -- FIXME: Specialise enumFromThenTo.
1635
1636 -- Conversions
1637 -- -----------
1638
1639 -- | Convert a 'Facets' to a list
1640 toList :: Monad m => Facets m v a -> m [a]
1641 {-# INLINE toList #-}
1642 toList = foldr (:) []
1643
1644 -- | Convert a list to a 'Facets'
1645 fromList :: Monad m => [a] -> Facets m v a
1646 {-# INLINE fromList #-}
1647 fromList xs = unsafeFromList Unknown xs
1648
1649 -- | Convert the first @n@ elements of a list to a 'Facets'
1650 fromListN :: Monad m => Int -> [a] -> Facets m v a
1651 {-# INLINE_FUSED fromListN #-}
1652 fromListN n xs = simple step (xs,n) (Max (delay_inline max n 0))
1653 where
1654 {-# INLINE_INNER step #-}
1655 step (xs,n) | n <= 0 = return Done
1656 step (x:xs,n) = return (Yield x (xs,n-1))
1657 step ([],n) = return Done
1658
1659 -- | Convert a list to a 'Facets' with the given 'Size' hint.
1660 unsafeFromList :: Monad m => Size -> [a] -> Facets m v a
1661 {-# INLINE_FUSED unsafeFromList #-}
1662 unsafeFromList sz xs = simple step xs sz
1663 where
1664 step (x:xs) = return (Yield x xs)
1665 step [] = return Done
1666
1667 fromVector :: (Monad m, Vector v a) => v a -> Facets m v a
1668 {-# INLINE_FUSED fromVector #-}
1669 fromVector v = v `seq` n `seq` Facets (Unf step 0)
1670 (Unf vstep True)
1671 (Just v)
1672 (Exact n)
1673 where
1674 n = basicLength v
1675
1676 {-# INLINE step #-}
1677 step i | i >= n = return Done
1678 | otherwise = case basicUnsafeIndexM v i of
1679 Box x -> return $ Yield x (i+1)
1680
1681
1682 {-# INLINE vstep #-}
1683 vstep True = return (Yield (Chunk (basicLength v) (\mv -> basicUnsafeCopy mv v)) False)
1684 vstep False = return Done
1685
1686 fromVectors :: (Monad m, Vector v a) => [v a] -> Facets m v a
1687 {-# INLINE_FUSED fromVectors #-}
1688 fromVectors vs = Facets (Unf pstep (Left vs))
1689 (Unf vstep vs)
1690 Nothing
1691 (Exact n)
1692 where
1693 n = List.foldl' (\k v -> k + basicLength v) 0 vs
1694
1695 pstep (Left []) = return Done
1696 pstep (Left (v:vs)) = basicLength v `seq` return (Skip (Right (v,0,vs)))
1697
1698 pstep (Right (v,i,vs))
1699 | i >= basicLength v = return $ Skip (Left vs)
1700 | otherwise = case basicUnsafeIndexM v i of
1701 Box x -> return $ Yield x (Right (v,i+1,vs))
1702
1703
1704 vstep [] = return Done
1705 vstep (v:vs) = return $ Yield (Chunk (basicLength v)
1706 (\mv -> basicUnsafeCopy mv v)) vs
1707
1708
1709 concatVectors :: (Monad m, Vector v a) => Facets m u (v a) -> Facets m v a
1710 {-# INLINE_FUSED concatVectors #-}
1711 concatVectors Facets{sElems = Unf step s}
1712 = Facets (Unf pstep (Left s))
1713 (Unf vstep s)
1714 Nothing
1715 Unknown
1716 where
1717 pstep (Left s) = do
1718 r <- step s
1719 case r of
1720 Yield v s' -> basicLength v `seq` return (Skip (Right (v,0,s')))
1721 Skip s' -> return (Skip (Left s'))
1722 Done -> return Done
1723
1724 pstep (Right (v,i,s))
1725 | i >= basicLength v = return (Skip (Left s))
1726 | otherwise = case basicUnsafeIndexM v i of
1727 Box x -> return (Yield x (Right (v,i+1,s)))
1728
1729
1730 vstep s = do
1731 r <- step s
1732 case r of
1733 Yield v s' -> return (Yield (Chunk (basicLength v)
1734 (\mv -> basicUnsafeCopy mv v)) s')
1735 Skip s' -> return (Skip s')
1736 Done -> return Done
1737
1738 reVector :: Monad m => Facets m u a -> Facets m v a
1739 {-# INLINE_FUSED reVector #-}
1740 reVector Facets{sElems = Unf step s, sSize = n} = simple step s n
1741
1742 {-# RULES
1743
1744 "reVector [Vector]"
1745 reVector = id
1746
1747 "reVector/reVector [Vector]" forall s.
1748 reVector (reVector s) = s
1749
1750 #-}
1751