diff options
Diffstat (limited to 'third_party/bazel/rules_haskell/examples/vector/Data/Vector/Fusion')
5 files changed, 3581 insertions, 0 deletions
diff --git a/third_party/bazel/rules_haskell/examples/vector/Data/Vector/Fusion/Bundle.hs b/third_party/bazel/rules_haskell/examples/vector/Data/Vector/Fusion/Bundle.hs new file mode 100644 index 0000000000..6b6b6236d7 --- /dev/null +++ b/third_party/bazel/rules_haskell/examples/vector/Data/Vector/Fusion/Bundle.hs @@ -0,0 +1,655 @@ +{-# LANGUAGE CPP, FlexibleInstances, Rank2Types, BangPatterns #-} + +-- | +-- Module : Data.Vector.Fusion.Bundle +-- Copyright : (c) Roman Leshchinskiy 2008-2010 +-- License : BSD-style +-- +-- Maintainer : Roman Leshchinskiy <rl@cse.unsw.edu.au> +-- Stability : experimental +-- Portability : non-portable +-- +-- Bundles for stream fusion +-- + +module Data.Vector.Fusion.Bundle ( + -- * Types + Step(..), Chunk(..), Bundle, MBundle, + + -- * In-place markers + inplace, + + -- * Size hints + size, sized, + + -- * Length information + length, null, + + -- * Construction + empty, singleton, cons, snoc, replicate, generate, (++), + + -- * Accessing individual elements + head, last, (!!), (!?), + + -- * Substreams + slice, init, tail, take, drop, + + -- * Mapping + map, concatMap, flatten, unbox, + + -- * Zipping + indexed, indexedR, + zipWith, zipWith3, zipWith4, zipWith5, zipWith6, + zip, zip3, zip4, zip5, zip6, + + -- * Filtering + filter, takeWhile, dropWhile, + + -- * Searching + elem, notElem, find, findIndex, + + -- * Folding + foldl, foldl1, foldl', foldl1', foldr, foldr1, + + -- * Specialised folds + and, or, + + -- * Unfolding + unfoldr, unfoldrN, iterateN, + + -- * Scans + prescanl, prescanl', + postscanl, postscanl', + scanl, scanl', + scanl1, scanl1', + + -- * Enumerations + enumFromStepN, enumFromTo, enumFromThenTo, + + -- * Conversions + toList, fromList, fromListN, unsafeFromList, lift, + fromVector, reVector, fromVectors, concatVectors, + + -- * Monadic combinators + mapM, mapM_, zipWithM, zipWithM_, filterM, foldM, fold1M, foldM', fold1M', + + eq, cmp, eqBy, cmpBy +) where + +import Data.Vector.Generic.Base ( Vector ) +import Data.Vector.Fusion.Bundle.Size +import Data.Vector.Fusion.Util +import Data.Vector.Fusion.Stream.Monadic ( Stream(..), Step(..) ) +import Data.Vector.Fusion.Bundle.Monadic ( Chunk(..) ) +import qualified Data.Vector.Fusion.Bundle.Monadic as M +import qualified Data.Vector.Fusion.Stream.Monadic as S + +import Prelude hiding ( length, null, + replicate, (++), + head, last, (!!), + init, tail, take, drop, + map, concatMap, + zipWith, zipWith3, zip, zip3, + filter, takeWhile, dropWhile, + elem, notElem, + foldl, foldl1, foldr, foldr1, + and, or, + scanl, scanl1, + enumFromTo, enumFromThenTo, + mapM, mapM_ ) + +#if MIN_VERSION_base(4,9,0) +import Data.Functor.Classes (Eq1 (..), Ord1 (..)) +#endif + +import GHC.Base ( build ) + +-- Data.Vector.Internal.Check is unused +#define NOT_VECTOR_MODULE +#include "vector.h" + +-- | The type of pure streams +type Bundle = M.Bundle Id + +-- | Alternative name for monadic streams +type MBundle = M.Bundle + +inplace :: (forall m. Monad m => S.Stream m a -> S.Stream m b) + -> (Size -> Size) -> Bundle v a -> Bundle v b +{-# INLINE_FUSED inplace #-} +inplace f g b = b `seq` M.fromStream (f (M.elements b)) (g (M.size b)) + +{-# RULES + +"inplace/inplace [Vector]" + forall (f1 :: forall m. Monad m => S.Stream m a -> S.Stream m a) + (f2 :: forall m. Monad m => S.Stream m a -> S.Stream m a) + g1 g2 s. + inplace f1 g1 (inplace f2 g2 s) = inplace (f1 . f2) (g1 . g2) s #-} + + + +-- | Convert a pure stream to a monadic stream +lift :: Monad m => Bundle v a -> M.Bundle m v a +{-# INLINE_FUSED lift #-} +lift (M.Bundle (Stream step s) (Stream vstep t) v sz) + = M.Bundle (Stream (return . unId . step) s) + (Stream (return . unId . vstep) t) v sz + +-- | 'Size' hint of a 'Bundle' +size :: Bundle v a -> Size +{-# INLINE size #-} +size = M.size + +-- | Attach a 'Size' hint to a 'Bundle' +sized :: Bundle v a -> Size -> Bundle v a +{-# INLINE sized #-} +sized = M.sized + +-- Length +-- ------ + +-- | Length of a 'Bundle' +length :: Bundle v a -> Int +{-# INLINE length #-} +length = unId . M.length + +-- | Check if a 'Bundle' is empty +null :: Bundle v a -> Bool +{-# INLINE null #-} +null = unId . M.null + +-- Construction +-- ------------ + +-- | Empty 'Bundle' +empty :: Bundle v a +{-# INLINE empty #-} +empty = M.empty + +-- | Singleton 'Bundle' +singleton :: a -> Bundle v a +{-# INLINE singleton #-} +singleton = M.singleton + +-- | Replicate a value to a given length +replicate :: Int -> a -> Bundle v a +{-# INLINE replicate #-} +replicate = M.replicate + +-- | Generate a stream from its indices +generate :: Int -> (Int -> a) -> Bundle v a +{-# INLINE generate #-} +generate = M.generate + +-- | Prepend an element +cons :: a -> Bundle v a -> Bundle v a +{-# INLINE cons #-} +cons = M.cons + +-- | Append an element +snoc :: Bundle v a -> a -> Bundle v a +{-# INLINE snoc #-} +snoc = M.snoc + +infixr 5 ++ +-- | Concatenate two 'Bundle's +(++) :: Bundle v a -> Bundle v a -> Bundle v a +{-# INLINE (++) #-} +(++) = (M.++) + +-- Accessing elements +-- ------------------ + +-- | First element of the 'Bundle' or error if empty +head :: Bundle v a -> a +{-# INLINE head #-} +head = unId . M.head + +-- | Last element of the 'Bundle' or error if empty +last :: Bundle v a -> a +{-# INLINE last #-} +last = unId . M.last + +infixl 9 !! +-- | Element at the given position +(!!) :: Bundle v a -> Int -> a +{-# INLINE (!!) #-} +s !! i = unId (s M.!! i) + +infixl 9 !? +-- | Element at the given position or 'Nothing' if out of bounds +(!?) :: Bundle v a -> Int -> Maybe a +{-# INLINE (!?) #-} +s !? i = unId (s M.!? i) + +-- Substreams +-- ---------- + +-- | Extract a substream of the given length starting at the given position. +slice :: Int -- ^ starting index + -> Int -- ^ length + -> Bundle v a + -> Bundle v a +{-# INLINE slice #-} +slice = M.slice + +-- | All but the last element +init :: Bundle v a -> Bundle v a +{-# INLINE init #-} +init = M.init + +-- | All but the first element +tail :: Bundle v a -> Bundle v a +{-# INLINE tail #-} +tail = M.tail + +-- | The first @n@ elements +take :: Int -> Bundle v a -> Bundle v a +{-# INLINE take #-} +take = M.take + +-- | All but the first @n@ elements +drop :: Int -> Bundle v a -> Bundle v a +{-# INLINE drop #-} +drop = M.drop + +-- Mapping +-- --------------- + +-- | Map a function over a 'Bundle' +map :: (a -> b) -> Bundle v a -> Bundle v b +{-# INLINE map #-} +map = M.map + +unbox :: Bundle v (Box a) -> Bundle v a +{-# INLINE unbox #-} +unbox = M.unbox + +concatMap :: (a -> Bundle v b) -> Bundle v a -> Bundle v b +{-# INLINE concatMap #-} +concatMap = M.concatMap + +-- Zipping +-- ------- + +-- | Pair each element in a 'Bundle' with its index +indexed :: Bundle v a -> Bundle v (Int,a) +{-# INLINE indexed #-} +indexed = M.indexed + +-- | Pair each element in a 'Bundle' with its index, starting from the right +-- and counting down +indexedR :: Int -> Bundle v a -> Bundle v (Int,a) +{-# INLINE_FUSED indexedR #-} +indexedR = M.indexedR + +-- | Zip two 'Bundle's with the given function +zipWith :: (a -> b -> c) -> Bundle v a -> Bundle v b -> Bundle v c +{-# INLINE zipWith #-} +zipWith = M.zipWith + +-- | Zip three 'Bundle's with the given function +zipWith3 :: (a -> b -> c -> d) -> Bundle v a -> Bundle v b -> Bundle v c -> Bundle v d +{-# INLINE zipWith3 #-} +zipWith3 = M.zipWith3 + +zipWith4 :: (a -> b -> c -> d -> e) + -> Bundle v a -> Bundle v b -> Bundle v c -> Bundle v d + -> Bundle v e +{-# INLINE zipWith4 #-} +zipWith4 = M.zipWith4 + +zipWith5 :: (a -> b -> c -> d -> e -> f) + -> Bundle v a -> Bundle v b -> Bundle v c -> Bundle v d + -> Bundle v e -> Bundle v f +{-# INLINE zipWith5 #-} +zipWith5 = M.zipWith5 + +zipWith6 :: (a -> b -> c -> d -> e -> f -> g) + -> Bundle v a -> Bundle v b -> Bundle v c -> Bundle v d + -> Bundle v e -> Bundle v f -> Bundle v g +{-# INLINE zipWith6 #-} +zipWith6 = M.zipWith6 + +zip :: Bundle v a -> Bundle v b -> Bundle v (a,b) +{-# INLINE zip #-} +zip = M.zip + +zip3 :: Bundle v a -> Bundle v b -> Bundle v c -> Bundle v (a,b,c) +{-# INLINE zip3 #-} +zip3 = M.zip3 + +zip4 :: Bundle v a -> Bundle v b -> Bundle v c -> Bundle v d + -> Bundle v (a,b,c,d) +{-# INLINE zip4 #-} +zip4 = M.zip4 + +zip5 :: Bundle v a -> Bundle v b -> Bundle v c -> Bundle v d + -> Bundle v e -> Bundle v (a,b,c,d,e) +{-# INLINE zip5 #-} +zip5 = M.zip5 + +zip6 :: Bundle v a -> Bundle v b -> Bundle v c -> Bundle v d + -> Bundle v e -> Bundle v f -> Bundle v (a,b,c,d,e,f) +{-# INLINE zip6 #-} +zip6 = M.zip6 + +-- Filtering +-- --------- + +-- | Drop elements which do not satisfy the predicate +filter :: (a -> Bool) -> Bundle v a -> Bundle v a +{-# INLINE filter #-} +filter = M.filter + +-- | Longest prefix of elements that satisfy the predicate +takeWhile :: (a -> Bool) -> Bundle v a -> Bundle v a +{-# INLINE takeWhile #-} +takeWhile = M.takeWhile + +-- | Drop the longest prefix of elements that satisfy the predicate +dropWhile :: (a -> Bool) -> Bundle v a -> Bundle v a +{-# INLINE dropWhile #-} +dropWhile = M.dropWhile + +-- Searching +-- --------- + +infix 4 `elem` +-- | Check whether the 'Bundle' contains an element +elem :: Eq a => a -> Bundle v a -> Bool +{-# INLINE elem #-} +elem x = unId . M.elem x + +infix 4 `notElem` +-- | Inverse of `elem` +notElem :: Eq a => a -> Bundle v a -> Bool +{-# INLINE notElem #-} +notElem x = unId . M.notElem x + +-- | Yield 'Just' the first element matching the predicate or 'Nothing' if no +-- such element exists. +find :: (a -> Bool) -> Bundle v a -> Maybe a +{-# INLINE find #-} +find f = unId . M.find f + +-- | Yield 'Just' the index of the first element matching the predicate or +-- 'Nothing' if no such element exists. +findIndex :: (a -> Bool) -> Bundle v a -> Maybe Int +{-# INLINE findIndex #-} +findIndex f = unId . M.findIndex f + +-- Folding +-- ------- + +-- | Left fold +foldl :: (a -> b -> a) -> a -> Bundle v b -> a +{-# INLINE foldl #-} +foldl f z = unId . M.foldl f z + +-- | Left fold on non-empty 'Bundle's +foldl1 :: (a -> a -> a) -> Bundle v a -> a +{-# INLINE foldl1 #-} +foldl1 f = unId . M.foldl1 f + +-- | Left fold with strict accumulator +foldl' :: (a -> b -> a) -> a -> Bundle v b -> a +{-# INLINE foldl' #-} +foldl' f z = unId . M.foldl' f z + +-- | Left fold on non-empty 'Bundle's with strict accumulator +foldl1' :: (a -> a -> a) -> Bundle v a -> a +{-# INLINE foldl1' #-} +foldl1' f = unId . M.foldl1' f + +-- | Right fold +foldr :: (a -> b -> b) -> b -> Bundle v a -> b +{-# INLINE foldr #-} +foldr f z = unId . M.foldr f z + +-- | Right fold on non-empty 'Bundle's +foldr1 :: (a -> a -> a) -> Bundle v a -> a +{-# INLINE foldr1 #-} +foldr1 f = unId . M.foldr1 f + +-- Specialised folds +-- ----------------- + +and :: Bundle v Bool -> Bool +{-# INLINE and #-} +and = unId . M.and + +or :: Bundle v Bool -> Bool +{-# INLINE or #-} +or = unId . M.or + +-- Unfolding +-- --------- + +-- | Unfold +unfoldr :: (s -> Maybe (a, s)) -> s -> Bundle v a +{-# INLINE unfoldr #-} +unfoldr = M.unfoldr + +-- | Unfold at most @n@ elements +unfoldrN :: Int -> (s -> Maybe (a, s)) -> s -> Bundle v a +{-# INLINE unfoldrN #-} +unfoldrN = M.unfoldrN + +-- | Apply function n-1 times to value. Zeroth element is original value. +iterateN :: Int -> (a -> a) -> a -> Bundle v a +{-# INLINE iterateN #-} +iterateN = M.iterateN + +-- Scans +-- ----- + +-- | Prefix scan +prescanl :: (a -> b -> a) -> a -> Bundle v b -> Bundle v a +{-# INLINE prescanl #-} +prescanl = M.prescanl + +-- | Prefix scan with strict accumulator +prescanl' :: (a -> b -> a) -> a -> Bundle v b -> Bundle v a +{-# INLINE prescanl' #-} +prescanl' = M.prescanl' + +-- | Suffix scan +postscanl :: (a -> b -> a) -> a -> Bundle v b -> Bundle v a +{-# INLINE postscanl #-} +postscanl = M.postscanl + +-- | Suffix scan with strict accumulator +postscanl' :: (a -> b -> a) -> a -> Bundle v b -> Bundle v a +{-# INLINE postscanl' #-} +postscanl' = M.postscanl' + +-- | Haskell-style scan +scanl :: (a -> b -> a) -> a -> Bundle v b -> Bundle v a +{-# INLINE scanl #-} +scanl = M.scanl + +-- | Haskell-style scan with strict accumulator +scanl' :: (a -> b -> a) -> a -> Bundle v b -> Bundle v a +{-# INLINE scanl' #-} +scanl' = M.scanl' + +-- | Scan over a non-empty 'Bundle' +scanl1 :: (a -> a -> a) -> Bundle v a -> Bundle v a +{-# INLINE scanl1 #-} +scanl1 = M.scanl1 + +-- | Scan over a non-empty 'Bundle' with a strict accumulator +scanl1' :: (a -> a -> a) -> Bundle v a -> Bundle v a +{-# INLINE scanl1' #-} +scanl1' = M.scanl1' + + +-- Comparisons +-- ----------- + +-- | Check if two 'Bundle's are equal +eq :: (Eq a) => Bundle v a -> Bundle v a -> Bool +{-# INLINE eq #-} +eq = eqBy (==) + +eqBy :: (a -> b -> Bool) -> Bundle v a -> Bundle v b -> Bool +{-# INLINE eqBy #-} +eqBy e x y = unId (M.eqBy e x y) + +-- | Lexicographically compare two 'Bundle's +cmp :: (Ord a) => Bundle v a -> Bundle v a -> Ordering +{-# INLINE cmp #-} +cmp = cmpBy compare + +cmpBy :: (a -> b -> Ordering) -> Bundle v a -> Bundle v b -> Ordering +{-# INLINE cmpBy #-} +cmpBy c x y = unId (M.cmpBy c x y) + +instance Eq a => Eq (M.Bundle Id v a) where + {-# INLINE (==) #-} + (==) = eq + +instance Ord a => Ord (M.Bundle Id v a) where + {-# INLINE compare #-} + compare = cmp + +#if MIN_VERSION_base(4,9,0) +instance Eq1 (M.Bundle Id v) where + {-# INLINE liftEq #-} + liftEq = eqBy + +instance Ord1 (M.Bundle Id v) where + {-# INLINE liftCompare #-} + liftCompare = cmpBy +#endif + +-- Monadic combinators +-- ------------------- + +-- | Apply a monadic action to each element of the stream, producing a monadic +-- stream of results +mapM :: Monad m => (a -> m b) -> Bundle v a -> M.Bundle m v b +{-# INLINE mapM #-} +mapM f = M.mapM f . lift + +-- | Apply a monadic action to each element of the stream +mapM_ :: Monad m => (a -> m b) -> Bundle v a -> m () +{-# INLINE mapM_ #-} +mapM_ f = M.mapM_ f . lift + +zipWithM :: Monad m => (a -> b -> m c) -> Bundle v a -> Bundle v b -> M.Bundle m v c +{-# INLINE zipWithM #-} +zipWithM f as bs = M.zipWithM f (lift as) (lift bs) + +zipWithM_ :: Monad m => (a -> b -> m c) -> Bundle v a -> Bundle v b -> m () +{-# INLINE zipWithM_ #-} +zipWithM_ f as bs = M.zipWithM_ f (lift as) (lift bs) + +-- | Yield a monadic stream of elements that satisfy the monadic predicate +filterM :: Monad m => (a -> m Bool) -> Bundle v a -> M.Bundle m v a +{-# INLINE filterM #-} +filterM f = M.filterM f . lift + +-- | Monadic fold +foldM :: Monad m => (a -> b -> m a) -> a -> Bundle v b -> m a +{-# INLINE foldM #-} +foldM m z = M.foldM m z . lift + +-- | Monadic fold over non-empty stream +fold1M :: Monad m => (a -> a -> m a) -> Bundle v a -> m a +{-# INLINE fold1M #-} +fold1M m = M.fold1M m . lift + +-- | Monadic fold with strict accumulator +foldM' :: Monad m => (a -> b -> m a) -> a -> Bundle v b -> m a +{-# INLINE foldM' #-} +foldM' m z = M.foldM' m z . lift + +-- | Monad fold over non-empty stream with strict accumulator +fold1M' :: Monad m => (a -> a -> m a) -> Bundle v a -> m a +{-# INLINE fold1M' #-} +fold1M' m = M.fold1M' m . lift + +-- Enumerations +-- ------------ + +-- | Yield a 'Bundle' of the given length containing the values @x@, @x+y@, +-- @x+y+y@ etc. +enumFromStepN :: Num a => a -> a -> Int -> Bundle v a +{-# INLINE enumFromStepN #-} +enumFromStepN = M.enumFromStepN + +-- | Enumerate values +-- +-- /WARNING:/ This operations can be very inefficient. If at all possible, use +-- 'enumFromStepN' instead. +enumFromTo :: Enum a => a -> a -> Bundle v a +{-# INLINE enumFromTo #-} +enumFromTo = M.enumFromTo + +-- | Enumerate values with a given step. +-- +-- /WARNING:/ This operations is very inefficient. If at all possible, use +-- 'enumFromStepN' instead. +enumFromThenTo :: Enum a => a -> a -> a -> Bundle v a +{-# INLINE enumFromThenTo #-} +enumFromThenTo = M.enumFromThenTo + +-- Conversions +-- ----------- + +-- | Convert a 'Bundle' to a list +toList :: Bundle v a -> [a] +{-# INLINE toList #-} +-- toList s = unId (M.toList s) +toList s = build (\c n -> toListFB c n s) + +-- This supports foldr/build list fusion that GHC implements +toListFB :: (a -> b -> b) -> b -> Bundle v a -> b +{-# INLINE [0] toListFB #-} +toListFB c n M.Bundle{M.sElems = Stream step t} = go t + where + go s = case unId (step s) of + Yield x s' -> x `c` go s' + Skip s' -> go s' + Done -> n + +-- | Create a 'Bundle' from a list +fromList :: [a] -> Bundle v a +{-# INLINE fromList #-} +fromList = M.fromList + +-- | Create a 'Bundle' from the first @n@ elements of a list +-- +-- > fromListN n xs = fromList (take n xs) +fromListN :: Int -> [a] -> Bundle v a +{-# INLINE fromListN #-} +fromListN = M.fromListN + +unsafeFromList :: Size -> [a] -> Bundle v a +{-# INLINE unsafeFromList #-} +unsafeFromList = M.unsafeFromList + +fromVector :: Vector v a => v a -> Bundle v a +{-# INLINE fromVector #-} +fromVector = M.fromVector + +reVector :: Bundle u a -> Bundle v a +{-# INLINE reVector #-} +reVector = M.reVector + +fromVectors :: Vector v a => [v a] -> Bundle v a +{-# INLINE fromVectors #-} +fromVectors = M.fromVectors + +concatVectors :: Vector v a => Bundle u (v a) -> Bundle v a +{-# INLINE concatVectors #-} +concatVectors = M.concatVectors + +-- | Create a 'Bundle' of values from a 'Bundle' of streamable things +flatten :: (a -> s) -> (s -> Step s b) -> Size -> Bundle v a -> Bundle v b +{-# INLINE_FUSED flatten #-} +flatten mk istep sz = M.flatten (return . mk) (return . istep) sz . lift + diff --git a/third_party/bazel/rules_haskell/examples/vector/Data/Vector/Fusion/Bundle/Monadic.hs b/third_party/bazel/rules_haskell/examples/vector/Data/Vector/Fusion/Bundle/Monadic.hs new file mode 100644 index 0000000000..46f4a165f8 --- /dev/null +++ b/third_party/bazel/rules_haskell/examples/vector/Data/Vector/Fusion/Bundle/Monadic.hs @@ -0,0 +1,1106 @@ +{-# LANGUAGE CPP, ExistentialQuantification, MultiParamTypeClasses, FlexibleInstances, Rank2Types, BangPatterns, KindSignatures, GADTs, ScopedTypeVariables #-} + +-- | +-- Module : Data.Vector.Fusion.Bundle.Monadic +-- Copyright : (c) Roman Leshchinskiy 2008-2010 +-- License : BSD-style +-- +-- Maintainer : Roman Leshchinskiy <rl@cse.unsw.edu.au> +-- Stability : experimental +-- Portability : non-portable +-- +-- Monadic bundles. +-- + +module Data.Vector.Fusion.Bundle.Monadic ( + Bundle(..), Chunk(..), + + -- * Size hints + size, sized, + + -- * Length + length, null, + + -- * Construction + empty, singleton, cons, snoc, replicate, replicateM, generate, generateM, (++), + + -- * Accessing elements + head, last, (!!), (!?), + + -- * Substreams + slice, init, tail, take, drop, + + -- * Mapping + map, mapM, mapM_, trans, unbox, concatMap, flatten, + + -- * Zipping + indexed, indexedR, zipWithM_, + zipWithM, zipWith3M, zipWith4M, zipWith5M, zipWith6M, + zipWith, zipWith3, zipWith4, zipWith5, zipWith6, + zip, zip3, zip4, zip5, zip6, + + -- * Comparisons + eqBy, cmpBy, + + -- * Filtering + filter, filterM, takeWhile, takeWhileM, dropWhile, dropWhileM, + + -- * Searching + elem, notElem, find, findM, findIndex, findIndexM, + + -- * Folding + foldl, foldlM, foldl1, foldl1M, foldM, fold1M, + foldl', foldlM', foldl1', foldl1M', foldM', fold1M', + foldr, foldrM, foldr1, foldr1M, + + -- * Specialised folds + and, or, concatMapM, + + -- * Unfolding + unfoldr, unfoldrM, + unfoldrN, unfoldrNM, + iterateN, iterateNM, + + -- * Scans + prescanl, prescanlM, prescanl', prescanlM', + postscanl, postscanlM, postscanl', postscanlM', + scanl, scanlM, scanl', scanlM', + scanl1, scanl1M, scanl1', scanl1M', + + -- * Enumerations + enumFromStepN, enumFromTo, enumFromThenTo, + + -- * Conversions + toList, fromList, fromListN, unsafeFromList, + fromVector, reVector, fromVectors, concatVectors, + fromStream, chunks, elements +) where + +import Data.Vector.Generic.Base +import qualified Data.Vector.Generic.Mutable.Base as M +import Data.Vector.Fusion.Bundle.Size +import Data.Vector.Fusion.Util ( Box(..), delay_inline ) +import Data.Vector.Fusion.Stream.Monadic ( Stream(..), Step(..) ) +import qualified Data.Vector.Fusion.Stream.Monadic as S +import Control.Monad.Primitive + +import qualified Data.List as List +import Data.Char ( ord ) +import GHC.Base ( unsafeChr ) +import Control.Monad ( liftM ) +import Prelude hiding ( length, null, + replicate, (++), + head, last, (!!), + init, tail, take, drop, + map, mapM, mapM_, concatMap, + zipWith, zipWith3, zip, zip3, + filter, takeWhile, dropWhile, + elem, notElem, + foldl, foldl1, foldr, foldr1, + and, or, + scanl, scanl1, + enumFromTo, enumFromThenTo ) + +import Data.Int ( Int8, Int16, Int32 ) +import Data.Word ( Word8, Word16, Word32, Word64 ) + +#if !MIN_VERSION_base(4,8,0) +import Data.Word ( Word ) +#endif + +#include "vector.h" +#include "MachDeps.h" + +#if WORD_SIZE_IN_BITS > 32 +import Data.Int ( Int64 ) +#endif + +data Chunk v a = Chunk Int (forall m. (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> m ()) + +-- | Monadic streams +data Bundle m v a = Bundle { sElems :: Stream m a + , sChunks :: Stream m (Chunk v a) + , sVector :: Maybe (v a) + , sSize :: Size + } + +fromStream :: Monad m => Stream m a -> Size -> Bundle m v a +{-# INLINE fromStream #-} +fromStream (Stream step t) sz = Bundle (Stream step t) (Stream step' t) Nothing sz + where + step' s = do r <- step s + return $ fmap (\x -> Chunk 1 (\v -> M.basicUnsafeWrite v 0 x)) r + +chunks :: Bundle m v a -> Stream m (Chunk v a) +{-# INLINE chunks #-} +chunks = sChunks + +elements :: Bundle m v a -> Stream m a +{-# INLINE elements #-} +elements = sElems + +-- | 'Size' hint of a 'Bundle' +size :: Bundle m v a -> Size +{-# INLINE size #-} +size = sSize + +-- | Attach a 'Size' hint to a 'Bundle' +sized :: Bundle m v a -> Size -> Bundle m v a +{-# INLINE_FUSED sized #-} +sized s sz = s { sSize = sz } + +-- Length +-- ------ + +-- | Length of a 'Bundle' +length :: Monad m => Bundle m v a -> m Int +{-# INLINE_FUSED length #-} +length Bundle{sSize = Exact n} = return n +length Bundle{sChunks = s} = S.foldl' (\n (Chunk k _) -> n+k) 0 s + +-- | Check if a 'Bundle' is empty +null :: Monad m => Bundle m v a -> m Bool +{-# INLINE_FUSED null #-} +null Bundle{sSize = Exact n} = return (n == 0) +null Bundle{sChunks = s} = S.foldr (\(Chunk n _) z -> n == 0 && z) True s + +-- Construction +-- ------------ + +-- | Empty 'Bundle' +empty :: Monad m => Bundle m v a +{-# INLINE_FUSED empty #-} +empty = fromStream S.empty (Exact 0) + +-- | Singleton 'Bundle' +singleton :: Monad m => a -> Bundle m v a +{-# INLINE_FUSED singleton #-} +singleton x = fromStream (S.singleton x) (Exact 1) + +-- | Replicate a value to a given length +replicate :: Monad m => Int -> a -> Bundle m v a +{-# INLINE_FUSED replicate #-} +replicate n x = Bundle (S.replicate n x) + (S.singleton $ Chunk len (\v -> M.basicSet v x)) + Nothing + (Exact len) + where + len = delay_inline max n 0 + +-- | Yield a 'Bundle' of values obtained by performing the monadic action the +-- given number of times +replicateM :: Monad m => Int -> m a -> Bundle m v a +{-# INLINE_FUSED replicateM #-} +-- NOTE: We delay inlining max here because GHC will create a join point for +-- the call to newArray# otherwise which is not really nice. +replicateM n p = fromStream (S.replicateM n p) (Exact (delay_inline max n 0)) + +generate :: Monad m => Int -> (Int -> a) -> Bundle m v a +{-# INLINE generate #-} +generate n f = generateM n (return . f) + +-- | Generate a stream from its indices +generateM :: Monad m => Int -> (Int -> m a) -> Bundle m v a +{-# INLINE_FUSED generateM #-} +generateM n f = fromStream (S.generateM n f) (Exact (delay_inline max n 0)) + +-- | Prepend an element +cons :: Monad m => a -> Bundle m v a -> Bundle m v a +{-# INLINE cons #-} +cons x s = singleton x ++ s + +-- | Append an element +snoc :: Monad m => Bundle m v a -> a -> Bundle m v a +{-# INLINE snoc #-} +snoc s x = s ++ singleton x + +infixr 5 ++ +-- | Concatenate two 'Bundle's +(++) :: Monad m => Bundle m v a -> Bundle m v a -> Bundle m v a +{-# INLINE_FUSED (++) #-} +Bundle sa ta _ na ++ Bundle sb tb _ nb = Bundle (sa S.++ sb) (ta S.++ tb) Nothing (na + nb) + +-- Accessing elements +-- ------------------ + +-- | First element of the 'Bundle' or error if empty +head :: Monad m => Bundle m v a -> m a +{-# INLINE_FUSED head #-} +head = S.head . sElems + +-- | Last element of the 'Bundle' or error if empty +last :: Monad m => Bundle m v a -> m a +{-# INLINE_FUSED last #-} +last = S.last . sElems + +infixl 9 !! +-- | Element at the given position +(!!) :: Monad m => Bundle m v a -> Int -> m a +{-# INLINE (!!) #-} +b !! i = sElems b S.!! i + +infixl 9 !? +-- | Element at the given position or 'Nothing' if out of bounds +(!?) :: Monad m => Bundle m v a -> Int -> m (Maybe a) +{-# INLINE (!?) #-} +b !? i = sElems b S.!? i + +-- Substreams +-- ---------- + +-- | Extract a substream of the given length starting at the given position. +slice :: Monad m => Int -- ^ starting index + -> Int -- ^ length + -> Bundle m v a + -> Bundle m v a +{-# INLINE slice #-} +slice i n s = take n (drop i s) + +-- | All but the last element +init :: Monad m => Bundle m v a -> Bundle m v a +{-# INLINE_FUSED init #-} +init Bundle{sElems = s, sSize = sz} = fromStream (S.init s) (sz-1) + +-- | All but the first element +tail :: Monad m => Bundle m v a -> Bundle m v a +{-# INLINE_FUSED tail #-} +tail Bundle{sElems = s, sSize = sz} = fromStream (S.tail s) (sz-1) + +-- | The first @n@ elements +take :: Monad m => Int -> Bundle m v a -> Bundle m v a +{-# INLINE_FUSED take #-} +take n Bundle{sElems = s, sSize = sz} = fromStream (S.take n s) (smaller (Exact n) sz) + +-- | All but the first @n@ elements +drop :: Monad m => Int -> Bundle m v a -> Bundle m v a +{-# INLINE_FUSED drop #-} +drop n Bundle{sElems = s, sSize = sz} = + fromStream (S.drop n s) (clampedSubtract sz (Exact n)) + +-- Mapping +-- ------- + +instance Monad m => Functor (Bundle m v) where + {-# INLINE fmap #-} + fmap = map + +-- | Map a function over a 'Bundle' +map :: Monad m => (a -> b) -> Bundle m v a -> Bundle m v b +{-# INLINE map #-} +map f = mapM (return . f) + +-- | Map a monadic function over a 'Bundle' +mapM :: Monad m => (a -> m b) -> Bundle m v a -> Bundle m v b +{-# INLINE_FUSED mapM #-} +mapM f Bundle{sElems = s, sSize = n} = fromStream (S.mapM f s) n + +-- | Execute a monadic action for each element of the 'Bundle' +mapM_ :: Monad m => (a -> m b) -> Bundle m v a -> m () +{-# INLINE_FUSED mapM_ #-} +mapM_ m = S.mapM_ m . sElems + +-- | Transform a 'Bundle' to use a different monad +trans :: (Monad m, Monad m') => (forall z. m z -> m' z) + -> Bundle m v a -> Bundle m' v a +{-# INLINE_FUSED trans #-} +trans f Bundle{sElems = s, sChunks = cs, sVector = v, sSize = n} + = Bundle { sElems = S.trans f s, sChunks = S.trans f cs, sVector = v, sSize = n } + +unbox :: Monad m => Bundle m v (Box a) -> Bundle m v a +{-# INLINE_FUSED unbox #-} +unbox Bundle{sElems = s, sSize = n} = fromStream (S.unbox s) n + +-- Zipping +-- ------- + +-- | Pair each element in a 'Bundle' with its index +indexed :: Monad m => Bundle m v a -> Bundle m v (Int,a) +{-# INLINE_FUSED indexed #-} +indexed Bundle{sElems = s, sSize = n} = fromStream (S.indexed s) n + +-- | Pair each element in a 'Bundle' with its index, starting from the right +-- and counting down +indexedR :: Monad m => Int -> Bundle m v a -> Bundle m v (Int,a) +{-# INLINE_FUSED indexedR #-} +indexedR m Bundle{sElems = s, sSize = n} = fromStream (S.indexedR m s) n + +-- | Zip two 'Bundle's with the given monadic function +zipWithM :: Monad m => (a -> b -> m c) -> Bundle m v a -> Bundle m v b -> Bundle m v c +{-# INLINE_FUSED zipWithM #-} +zipWithM f Bundle{sElems = sa, sSize = na} + Bundle{sElems = sb, sSize = nb} = fromStream (S.zipWithM f sa sb) (smaller na nb) + +-- FIXME: This might expose an opportunity for inplace execution. +{-# RULES + +"zipWithM xs xs [Vector.Bundle]" forall f xs. + zipWithM f xs xs = mapM (\x -> f x x) xs #-} + + +zipWithM_ :: Monad m => (a -> b -> m c) -> Bundle m v a -> Bundle m v b -> m () +{-# INLINE zipWithM_ #-} +zipWithM_ f sa sb = S.zipWithM_ f (sElems sa) (sElems sb) + +zipWith3M :: Monad m => (a -> b -> c -> m d) -> Bundle m v a -> Bundle m v b -> Bundle m v c -> Bundle m v d +{-# INLINE_FUSED zipWith3M #-} +zipWith3M f Bundle{sElems = sa, sSize = na} + Bundle{sElems = sb, sSize = nb} + Bundle{sElems = sc, sSize = nc} + = fromStream (S.zipWith3M f sa sb sc) (smaller na (smaller nb nc)) + +zipWith4M :: Monad m => (a -> b -> c -> d -> m e) + -> Bundle m v a -> Bundle m v b -> Bundle m v c -> Bundle m v d + -> Bundle m v e +{-# INLINE zipWith4M #-} +zipWith4M f sa sb sc sd + = zipWithM (\(a,b) (c,d) -> f a b c d) (zip sa sb) (zip sc sd) + +zipWith5M :: Monad m => (a -> b -> c -> d -> e -> m f) + -> Bundle m v a -> Bundle m v b -> Bundle m v c -> Bundle m v d + -> Bundle m v e -> Bundle m v f +{-# INLINE zipWith5M #-} +zipWith5M f sa sb sc sd se + = zipWithM (\(a,b,c) (d,e) -> f a b c d e) (zip3 sa sb sc) (zip sd se) + +zipWith6M :: Monad m => (a -> b -> c -> d -> e -> f -> m g) + -> Bundle m v a -> Bundle m v b -> Bundle m v c -> Bundle m v d + -> Bundle m v e -> Bundle m v f -> Bundle m v g +{-# INLINE zipWith6M #-} +zipWith6M fn sa sb sc sd se sf + = zipWithM (\(a,b,c) (d,e,f) -> fn a b c d e f) (zip3 sa sb sc) + (zip3 sd se sf) + +zipWith :: Monad m => (a -> b -> c) -> Bundle m v a -> Bundle m v b -> Bundle m v c +{-# INLINE zipWith #-} +zipWith f = zipWithM (\a b -> return (f a b)) + +zipWith3 :: Monad m => (a -> b -> c -> d) + -> Bundle m v a -> Bundle m v b -> Bundle m v c -> Bundle m v d +{-# INLINE zipWith3 #-} +zipWith3 f = zipWith3M (\a b c -> return (f a b c)) + +zipWith4 :: Monad m => (a -> b -> c -> d -> e) + -> Bundle m v a -> Bundle m v b -> Bundle m v c -> Bundle m v d + -> Bundle m v e +{-# INLINE zipWith4 #-} +zipWith4 f = zipWith4M (\a b c d -> return (f a b c d)) + +zipWith5 :: Monad m => (a -> b -> c -> d -> e -> f) + -> Bundle m v a -> Bundle m v b -> Bundle m v c -> Bundle m v d + -> Bundle m v e -> Bundle m v f +{-# INLINE zipWith5 #-} +zipWith5 f = zipWith5M (\a b c d e -> return (f a b c d e)) + +zipWith6 :: Monad m => (a -> b -> c -> d -> e -> f -> g) + -> Bundle m v a -> Bundle m v b -> Bundle m v c -> Bundle m v d + -> Bundle m v e -> Bundle m v f -> Bundle m v g +{-# INLINE zipWith6 #-} +zipWith6 fn = zipWith6M (\a b c d e f -> return (fn a b c d e f)) + +zip :: Monad m => Bundle m v a -> Bundle m v b -> Bundle m v (a,b) +{-# INLINE zip #-} +zip = zipWith (,) + +zip3 :: Monad m => Bundle m v a -> Bundle m v b -> Bundle m v c -> Bundle m v (a,b,c) +{-# INLINE zip3 #-} +zip3 = zipWith3 (,,) + +zip4 :: Monad m => Bundle m v a -> Bundle m v b -> Bundle m v c -> Bundle m v d + -> Bundle m v (a,b,c,d) +{-# INLINE zip4 #-} +zip4 = zipWith4 (,,,) + +zip5 :: Monad m => Bundle m v a -> Bundle m v b -> Bundle m v c -> Bundle m v d + -> Bundle m v e -> Bundle m v (a,b,c,d,e) +{-# INLINE zip5 #-} +zip5 = zipWith5 (,,,,) + +zip6 :: Monad m => Bundle m v a -> Bundle m v b -> Bundle m v c -> Bundle m v d + -> Bundle m v e -> Bundle m v f -> Bundle m v (a,b,c,d,e,f) +{-# INLINE zip6 #-} +zip6 = zipWith6 (,,,,,) + +-- Comparisons +-- ----------- + +-- | Check if two 'Bundle's are equal +eqBy :: (Monad m) => (a -> b -> Bool) -> Bundle m v a -> Bundle m v b -> m Bool +{-# INLINE_FUSED eqBy #-} +eqBy eq x y = S.eqBy eq (sElems x) (sElems y) + +-- | Lexicographically compare two 'Bundle's +cmpBy :: (Monad m) => (a -> b -> Ordering) -> Bundle m v a -> Bundle m v b -> m Ordering +{-# INLINE_FUSED cmpBy #-} +cmpBy cmp x y = S.cmpBy cmp (sElems x) (sElems y) + +-- Filtering +-- --------- + +-- | Drop elements which do not satisfy the predicate +filter :: Monad m => (a -> Bool) -> Bundle m v a -> Bundle m v a +{-# INLINE filter #-} +filter f = filterM (return . f) + +-- | Drop elements which do not satisfy the monadic predicate +filterM :: Monad m => (a -> m Bool) -> Bundle m v a -> Bundle m v a +{-# INLINE_FUSED filterM #-} +filterM f Bundle{sElems = s, sSize = n} = fromStream (S.filterM f s) (toMax n) + +-- | Longest prefix of elements that satisfy the predicate +takeWhile :: Monad m => (a -> Bool) -> Bundle m v a -> Bundle m v a +{-# INLINE takeWhile #-} +takeWhile f = takeWhileM (return . f) + +-- | Longest prefix of elements that satisfy the monadic predicate +takeWhileM :: Monad m => (a -> m Bool) -> Bundle m v a -> Bundle m v a +{-# INLINE_FUSED takeWhileM #-} +takeWhileM f Bundle{sElems = s, sSize = n} = fromStream (S.takeWhileM f s) (toMax n) + +-- | Drop the longest prefix of elements that satisfy the predicate +dropWhile :: Monad m => (a -> Bool) -> Bundle m v a -> Bundle m v a +{-# INLINE dropWhile #-} +dropWhile f = dropWhileM (return . f) + +-- | Drop the longest prefix of elements that satisfy the monadic predicate +dropWhileM :: Monad m => (a -> m Bool) -> Bundle m v a -> Bundle m v a +{-# INLINE_FUSED dropWhileM #-} +dropWhileM f Bundle{sElems = s, sSize = n} = fromStream (S.dropWhileM f s) (toMax n) + +-- Searching +-- --------- + +infix 4 `elem` +-- | Check whether the 'Bundle' contains an element +elem :: (Monad m, Eq a) => a -> Bundle m v a -> m Bool +{-# INLINE_FUSED elem #-} +elem x = S.elem x . sElems + +infix 4 `notElem` +-- | Inverse of `elem` +notElem :: (Monad m, Eq a) => a -> Bundle m v a -> m Bool +{-# INLINE notElem #-} +notElem x = S.notElem x . sElems + +-- | Yield 'Just' the first element that satisfies the predicate or 'Nothing' +-- if no such element exists. +find :: Monad m => (a -> Bool) -> Bundle m v a -> m (Maybe a) +{-# INLINE find #-} +find f = findM (return . f) + +-- | Yield 'Just' the first element that satisfies the monadic predicate or +-- 'Nothing' if no such element exists. +findM :: Monad m => (a -> m Bool) -> Bundle m v a -> m (Maybe a) +{-# INLINE_FUSED findM #-} +findM f = S.findM f . sElems + +-- | Yield 'Just' the index of the first element that satisfies the predicate +-- or 'Nothing' if no such element exists. +findIndex :: Monad m => (a -> Bool) -> Bundle m v a -> m (Maybe Int) +{-# INLINE_FUSED findIndex #-} +findIndex f = findIndexM (return . f) + +-- | Yield 'Just' the index of the first element that satisfies the monadic +-- predicate or 'Nothing' if no such element exists. +findIndexM :: Monad m => (a -> m Bool) -> Bundle m v a -> m (Maybe Int) +{-# INLINE_FUSED findIndexM #-} +findIndexM f = S.findIndexM f . sElems + +-- Folding +-- ------- + +-- | Left fold +foldl :: Monad m => (a -> b -> a) -> a -> Bundle m v b -> m a +{-# INLINE foldl #-} +foldl f = foldlM (\a b -> return (f a b)) + +-- | Left fold with a monadic operator +foldlM :: Monad m => (a -> b -> m a) -> a -> Bundle m v b -> m a +{-# INLINE_FUSED foldlM #-} +foldlM m z = S.foldlM m z . sElems + +-- | Same as 'foldlM' +foldM :: Monad m => (a -> b -> m a) -> a -> Bundle m v b -> m a +{-# INLINE foldM #-} +foldM = foldlM + +-- | Left fold over a non-empty 'Bundle' +foldl1 :: Monad m => (a -> a -> a) -> Bundle m v a -> m a +{-# INLINE foldl1 #-} +foldl1 f = foldl1M (\a b -> return (f a b)) + +-- | Left fold over a non-empty 'Bundle' with a monadic operator +foldl1M :: Monad m => (a -> a -> m a) -> Bundle m v a -> m a +{-# INLINE_FUSED foldl1M #-} +foldl1M f = S.foldl1M f . sElems + +-- | Same as 'foldl1M' +fold1M :: Monad m => (a -> a -> m a) -> Bundle m v a -> m a +{-# INLINE fold1M #-} +fold1M = foldl1M + +-- | Left fold with a strict accumulator +foldl' :: Monad m => (a -> b -> a) -> a -> Bundle m v b -> m a +{-# INLINE foldl' #-} +foldl' f = foldlM' (\a b -> return (f a b)) + +-- | Left fold with a strict accumulator and a monadic operator +foldlM' :: Monad m => (a -> b -> m a) -> a -> Bundle m v b -> m a +{-# INLINE_FUSED foldlM' #-} +foldlM' m z = S.foldlM' m z . sElems + +-- | Same as 'foldlM'' +foldM' :: Monad m => (a -> b -> m a) -> a -> Bundle m v b -> m a +{-# INLINE foldM' #-} +foldM' = foldlM' + +-- | Left fold over a non-empty 'Bundle' with a strict accumulator +foldl1' :: Monad m => (a -> a -> a) -> Bundle m v a -> m a +{-# INLINE foldl1' #-} +foldl1' f = foldl1M' (\a b -> return (f a b)) + +-- | Left fold over a non-empty 'Bundle' with a strict accumulator and a +-- monadic operator +foldl1M' :: Monad m => (a -> a -> m a) -> Bundle m v a -> m a +{-# INLINE_FUSED foldl1M' #-} +foldl1M' f = S.foldl1M' f . sElems + +-- | Same as 'foldl1M'' +fold1M' :: Monad m => (a -> a -> m a) -> Bundle m v a -> m a +{-# INLINE fold1M' #-} +fold1M' = foldl1M' + +-- | Right fold +foldr :: Monad m => (a -> b -> b) -> b -> Bundle m v a -> m b +{-# INLINE foldr #-} +foldr f = foldrM (\a b -> return (f a b)) + +-- | Right fold with a monadic operator +foldrM :: Monad m => (a -> b -> m b) -> b -> Bundle m v a -> m b +{-# INLINE_FUSED foldrM #-} +foldrM f z = S.foldrM f z . sElems + +-- | Right fold over a non-empty stream +foldr1 :: Monad m => (a -> a -> a) -> Bundle m v a -> m a +{-# INLINE foldr1 #-} +foldr1 f = foldr1M (\a b -> return (f a b)) + +-- | Right fold over a non-empty stream with a monadic operator +foldr1M :: Monad m => (a -> a -> m a) -> Bundle m v a -> m a +{-# INLINE_FUSED foldr1M #-} +foldr1M f = S.foldr1M f . sElems + +-- Specialised folds +-- ----------------- + +and :: Monad m => Bundle m v Bool -> m Bool +{-# INLINE_FUSED and #-} +and = S.and . sElems + +or :: Monad m => Bundle m v Bool -> m Bool +{-# INLINE_FUSED or #-} +or = S.or . sElems + +concatMap :: Monad m => (a -> Bundle m v b) -> Bundle m v a -> Bundle m v b +{-# INLINE concatMap #-} +concatMap f = concatMapM (return . f) + +concatMapM :: Monad m => (a -> m (Bundle m v b)) -> Bundle m v a -> Bundle m v b +{-# INLINE_FUSED concatMapM #-} +concatMapM f Bundle{sElems = s} = fromStream (S.concatMapM (liftM sElems . f) s) Unknown + +-- | Create a 'Bundle' of values from a 'Bundle' of streamable things +flatten :: Monad m => (a -> m s) -> (s -> m (Step s b)) -> Size + -> Bundle m v a -> Bundle m v b +{-# INLINE_FUSED flatten #-} +flatten mk istep sz Bundle{sElems = s} = fromStream (S.flatten mk istep s) sz + +-- Unfolding +-- --------- + +-- | Unfold +unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Bundle m u a +{-# INLINE_FUSED unfoldr #-} +unfoldr f = unfoldrM (return . f) + +-- | Unfold with a monadic function +unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Bundle m u a +{-# INLINE_FUSED unfoldrM #-} +unfoldrM f s = fromStream (S.unfoldrM f s) Unknown + +-- | Unfold at most @n@ elements +unfoldrN :: Monad m => Int -> (s -> Maybe (a, s)) -> s -> Bundle m u a +{-# INLINE_FUSED unfoldrN #-} +unfoldrN n f = unfoldrNM n (return . f) + +-- | Unfold at most @n@ elements with a monadic functions +unfoldrNM :: Monad m => Int -> (s -> m (Maybe (a, s))) -> s -> Bundle m u a +{-# INLINE_FUSED unfoldrNM #-} +unfoldrNM n f s = fromStream (S.unfoldrNM n f s) (Max (delay_inline max n 0)) + +-- | Apply monadic function n times to value. Zeroth element is original value. +iterateNM :: Monad m => Int -> (a -> m a) -> a -> Bundle m u a +{-# INLINE_FUSED iterateNM #-} +iterateNM n f x0 = fromStream (S.iterateNM n f x0) (Exact (delay_inline max n 0)) + +-- | Apply function n times to value. Zeroth element is original value. +iterateN :: Monad m => Int -> (a -> a) -> a -> Bundle m u a +{-# INLINE_FUSED iterateN #-} +iterateN n f x0 = iterateNM n (return . f) x0 + +-- Scans +-- ----- + +-- | Prefix scan +prescanl :: Monad m => (a -> b -> a) -> a -> Bundle m v b -> Bundle m v a +{-# INLINE prescanl #-} +prescanl f = prescanlM (\a b -> return (f a b)) + +-- | Prefix scan with a monadic operator +prescanlM :: Monad m => (a -> b -> m a) -> a -> Bundle m v b -> Bundle m v a +{-# INLINE_FUSED prescanlM #-} +prescanlM f z Bundle{sElems = s, sSize = sz} = fromStream (S.prescanlM f z s) sz + +-- | Prefix scan with strict accumulator +prescanl' :: Monad m => (a -> b -> a) -> a -> Bundle m v b -> Bundle m v a +{-# INLINE prescanl' #-} +prescanl' f = prescanlM' (\a b -> return (f a b)) + +-- | Prefix scan with strict accumulator and a monadic operator +prescanlM' :: Monad m => (a -> b -> m a) -> a -> Bundle m v b -> Bundle m v a +{-# INLINE_FUSED prescanlM' #-} +prescanlM' f z Bundle{sElems = s, sSize = sz} = fromStream (S.prescanlM' f z s) sz + +-- | Suffix scan +postscanl :: Monad m => (a -> b -> a) -> a -> Bundle m v b -> Bundle m v a +{-# INLINE postscanl #-} +postscanl f = postscanlM (\a b -> return (f a b)) + +-- | Suffix scan with a monadic operator +postscanlM :: Monad m => (a -> b -> m a) -> a -> Bundle m v b -> Bundle m v a +{-# INLINE_FUSED postscanlM #-} +postscanlM f z Bundle{sElems = s, sSize = sz} = fromStream (S.postscanlM f z s) sz + +-- | Suffix scan with strict accumulator +postscanl' :: Monad m => (a -> b -> a) -> a -> Bundle m v b -> Bundle m v a +{-# INLINE postscanl' #-} +postscanl' f = postscanlM' (\a b -> return (f a b)) + +-- | Suffix scan with strict acccumulator and a monadic operator +postscanlM' :: Monad m => (a -> b -> m a) -> a -> Bundle m v b -> Bundle m v a +{-# INLINE_FUSED postscanlM' #-} +postscanlM' f z Bundle{sElems = s, sSize = sz} = fromStream (S.postscanlM' f z s) sz + +-- | Haskell-style scan +scanl :: Monad m => (a -> b -> a) -> a -> Bundle m v b -> Bundle m v a +{-# INLINE scanl #-} +scanl f = scanlM (\a b -> return (f a b)) + +-- | Haskell-style scan with a monadic operator +scanlM :: Monad m => (a -> b -> m a) -> a -> Bundle m v b -> Bundle m v a +{-# INLINE scanlM #-} +scanlM f z s = z `cons` postscanlM f z s + +-- | Haskell-style scan with strict accumulator +scanl' :: Monad m => (a -> b -> a) -> a -> Bundle m v b -> Bundle m v a +{-# INLINE scanl' #-} +scanl' f = scanlM' (\a b -> return (f a b)) + +-- | Haskell-style scan with strict accumulator and a monadic operator +scanlM' :: Monad m => (a -> b -> m a) -> a -> Bundle m v b -> Bundle m v a +{-# INLINE scanlM' #-} +scanlM' f z s = z `seq` (z `cons` postscanlM f z s) + +-- | Scan over a non-empty 'Bundle' +scanl1 :: Monad m => (a -> a -> a) -> Bundle m v a -> Bundle m v a +{-# INLINE scanl1 #-} +scanl1 f = scanl1M (\x y -> return (f x y)) + +-- | Scan over a non-empty 'Bundle' with a monadic operator +scanl1M :: Monad m => (a -> a -> m a) -> Bundle m v a -> Bundle m v a +{-# INLINE_FUSED scanl1M #-} +scanl1M f Bundle{sElems = s, sSize = sz} = fromStream (S.scanl1M f s) sz + +-- | Scan over a non-empty 'Bundle' with a strict accumulator +scanl1' :: Monad m => (a -> a -> a) -> Bundle m v a -> Bundle m v a +{-# INLINE scanl1' #-} +scanl1' f = scanl1M' (\x y -> return (f x y)) + +-- | Scan over a non-empty 'Bundle' with a strict accumulator and a monadic +-- operator +scanl1M' :: Monad m => (a -> a -> m a) -> Bundle m v a -> Bundle m v a +{-# INLINE_FUSED scanl1M' #-} +scanl1M' f Bundle{sElems = s, sSize = sz} = fromStream (S.scanl1M' f s) sz + +-- Enumerations +-- ------------ + +-- The Enum class is broken for this, there just doesn't seem to be a +-- way to implement this generically. We have to specialise for as many types +-- as we can but this doesn't help in polymorphic loops. + +-- | Yield a 'Bundle' of the given length containing the values @x@, @x+y@, +-- @x+y+y@ etc. +enumFromStepN :: (Num a, Monad m) => a -> a -> Int -> Bundle m v a +{-# INLINE_FUSED enumFromStepN #-} +enumFromStepN x y n = fromStream (S.enumFromStepN x y n) (Exact (delay_inline max n 0)) + +-- | Enumerate values +-- +-- /WARNING:/ This operation can be very inefficient. If at all possible, use +-- 'enumFromStepN' instead. +enumFromTo :: (Enum a, Monad m) => a -> a -> Bundle m v a +{-# INLINE_FUSED enumFromTo #-} +enumFromTo x y = fromList [x .. y] + +-- NOTE: We use (x+1) instead of (succ x) below because the latter checks for +-- overflow which can't happen here. + +-- FIXME: add "too large" test for Int +enumFromTo_small :: (Integral a, Monad m) => a -> a -> Bundle m v a +{-# INLINE_FUSED enumFromTo_small #-} +enumFromTo_small x y = x `seq` y `seq` fromStream (Stream step x) (Exact n) + where + n = delay_inline max (fromIntegral y - fromIntegral x + 1) 0 + + {-# INLINE_INNER step #-} + step z | z <= y = return $ Yield z (z+1) + | otherwise = return $ Done + +{-# RULES + +"enumFromTo<Int8> [Bundle]" + enumFromTo = enumFromTo_small :: Monad m => Int8 -> Int8 -> Bundle m v Int8 + +"enumFromTo<Int16> [Bundle]" + enumFromTo = enumFromTo_small :: Monad m => Int16 -> Int16 -> Bundle m v Int16 + +"enumFromTo<Word8> [Bundle]" + enumFromTo = enumFromTo_small :: Monad m => Word8 -> Word8 -> Bundle m v Word8 + +"enumFromTo<Word16> [Bundle]" + enumFromTo = enumFromTo_small :: Monad m => Word16 -> Word16 -> Bundle m v Word16 #-} + + + +#if WORD_SIZE_IN_BITS > 32 + +{-# RULES + +"enumFromTo<Int32> [Bundle]" + enumFromTo = enumFromTo_small :: Monad m => Int32 -> Int32 -> Bundle m v Int32 + +"enumFromTo<Word32> [Bundle]" + enumFromTo = enumFromTo_small :: Monad m => Word32 -> Word32 -> Bundle m v Word32 #-} + +#endif + +-- NOTE: We could implement a generic "too large" test: +-- +-- len x y | x > y = 0 +-- | n > 0 && n <= fromIntegral (maxBound :: Int) = fromIntegral n +-- | otherwise = error +-- where +-- n = y-x+1 +-- +-- Alas, GHC won't eliminate unnecessary comparisons (such as n >= 0 for +-- unsigned types). See http://hackage.haskell.org/trac/ghc/ticket/3744 +-- + +enumFromTo_int :: forall m v. Monad m => Int -> Int -> Bundle m v Int +{-# INLINE_FUSED enumFromTo_int #-} +enumFromTo_int x y = x `seq` y `seq` fromStream (Stream step x) (Exact (len x y)) + where + {-# INLINE [0] len #-} + len :: Int -> Int -> Int + len u v | u > v = 0 + | otherwise = BOUNDS_CHECK(check) "enumFromTo" "vector too large" + (n > 0) + $ n + where + n = v-u+1 + + {-# INLINE_INNER step #-} + step z | z <= y = return $ Yield z (z+1) + | otherwise = return $ Done + +enumFromTo_intlike :: (Integral a, Monad m) => a -> a -> Bundle m v a +{-# INLINE_FUSED enumFromTo_intlike #-} +enumFromTo_intlike x y = x `seq` y `seq` fromStream (Stream step x) (Exact (len x y)) + where + {-# INLINE [0] len #-} + len u v | u > v = 0 + | otherwise = BOUNDS_CHECK(check) "enumFromTo" "vector too large" + (n > 0) + $ fromIntegral n + where + n = v-u+1 + + {-# INLINE_INNER step #-} + step z | z <= y = return $ Yield z (z+1) + | otherwise = return $ Done + +{-# RULES + +"enumFromTo<Int> [Bundle]" + enumFromTo = enumFromTo_int :: Monad m => Int -> Int -> Bundle m v Int + +#if WORD_SIZE_IN_BITS > 32 + +"enumFromTo<Int64> [Bundle]" + enumFromTo = enumFromTo_intlike :: Monad m => Int64 -> Int64 -> Bundle m v Int64 #-} + +#else + +"enumFromTo<Int32> [Bundle]" + enumFromTo = enumFromTo_intlike :: Monad m => Int32 -> Int32 -> Bundle m v Int32 #-} + +#endif + + + +enumFromTo_big_word :: (Integral a, Monad m) => a -> a -> Bundle m v a +{-# INLINE_FUSED enumFromTo_big_word #-} +enumFromTo_big_word x y = x `seq` y `seq` fromStream (Stream step x) (Exact (len x y)) + where + {-# INLINE [0] len #-} + len u v | u > v = 0 + | otherwise = BOUNDS_CHECK(check) "enumFromTo" "vector too large" + (n < fromIntegral (maxBound :: Int)) + $ fromIntegral (n+1) + where + n = v-u + + {-# INLINE_INNER step #-} + step z | z <= y = return $ Yield z (z+1) + | otherwise = return $ Done + +{-# RULES + +"enumFromTo<Word> [Bundle]" + enumFromTo = enumFromTo_big_word :: Monad m => Word -> Word -> Bundle m v Word + +"enumFromTo<Word64> [Bundle]" + enumFromTo = enumFromTo_big_word + :: Monad m => Word64 -> Word64 -> Bundle m v Word64 + +#if WORD_SIZE_IN_BITS == 32 + +"enumFromTo<Word32> [Bundle]" + enumFromTo = enumFromTo_big_word + :: Monad m => Word32 -> Word32 -> Bundle m v Word32 + +#endif + +"enumFromTo<Integer> [Bundle]" + enumFromTo = enumFromTo_big_word + :: Monad m => Integer -> Integer -> Bundle m v Integer #-} + + +#if WORD_SIZE_IN_BITS > 32 + +-- FIXME: the "too large" test is totally wrong +enumFromTo_big_int :: (Integral a, Monad m) => a -> a -> Bundle m v a +{-# INLINE_FUSED enumFromTo_big_int #-} +enumFromTo_big_int x y = x `seq` y `seq` fromStream (Stream step x) (Exact (len x y)) + where + {-# INLINE [0] len #-} + len u v | u > v = 0 + | otherwise = BOUNDS_CHECK(check) "enumFromTo" "vector too large" + (n > 0 && n <= fromIntegral (maxBound :: Int)) + $ fromIntegral n + where + n = v-u+1 + + {-# INLINE_INNER step #-} + step z | z <= y = return $ Yield z (z+1) + | otherwise = return $ Done + + +{-# RULES + +"enumFromTo<Int64> [Bundle]" + enumFromTo = enumFromTo_big_int :: Monad m => Int64 -> Int64 -> Bundle m v Int64 #-} + + + +#endif + +enumFromTo_char :: Monad m => Char -> Char -> Bundle m v Char +{-# INLINE_FUSED enumFromTo_char #-} +enumFromTo_char x y = x `seq` y `seq` fromStream (Stream step xn) (Exact n) + where + xn = ord x + yn = ord y + + n = delay_inline max 0 (yn - xn + 1) + + {-# INLINE_INNER step #-} + step zn | zn <= yn = return $ Yield (unsafeChr zn) (zn+1) + | otherwise = return $ Done + +{-# RULES + +"enumFromTo<Char> [Bundle]" + enumFromTo = enumFromTo_char #-} + + + +------------------------------------------------------------------------ + +-- Specialise enumFromTo for Float and Double. +-- Also, try to do something about pairs? + +enumFromTo_double :: (Monad m, Ord a, RealFrac a) => a -> a -> Bundle m v a +{-# INLINE_FUSED enumFromTo_double #-} +enumFromTo_double n m = n `seq` m `seq` fromStream (Stream step n) (Max (len n lim)) + where + lim = m + 1/2 -- important to float out + + {-# INLINE [0] len #-} + len x y | x > y = 0 + | otherwise = BOUNDS_CHECK(check) "enumFromTo" "vector too large" + (l > 0) + $ fromIntegral l + where + l :: Integer + l = truncate (y-x)+2 + + {-# INLINE_INNER step #-} + step x | x <= lim = return $ Yield x (x+1) + | otherwise = return $ Done + +{-# RULES + +"enumFromTo<Double> [Bundle]" + enumFromTo = enumFromTo_double :: Monad m => Double -> Double -> Bundle m v Double + +"enumFromTo<Float> [Bundle]" + enumFromTo = enumFromTo_double :: Monad m => Float -> Float -> Bundle m v Float #-} + + + +------------------------------------------------------------------------ + +-- | Enumerate values with a given step. +-- +-- /WARNING:/ This operation is very inefficient. If at all possible, use +-- 'enumFromStepN' instead. +enumFromThenTo :: (Enum a, Monad m) => a -> a -> a -> Bundle m v a +{-# INLINE_FUSED enumFromThenTo #-} +enumFromThenTo x y z = fromList [x, y .. z] + +-- FIXME: Specialise enumFromThenTo. + +-- Conversions +-- ----------- + +-- | Convert a 'Bundle' to a list +toList :: Monad m => Bundle m v a -> m [a] +{-# INLINE toList #-} +toList = foldr (:) [] + +-- | Convert a list to a 'Bundle' +fromList :: Monad m => [a] -> Bundle m v a +{-# INLINE fromList #-} +fromList xs = unsafeFromList Unknown xs + +-- | Convert the first @n@ elements of a list to a 'Bundle' +fromListN :: Monad m => Int -> [a] -> Bundle m v a +{-# INLINE_FUSED fromListN #-} +fromListN n xs = fromStream (S.fromListN n xs) (Max (delay_inline max n 0)) + +-- | Convert a list to a 'Bundle' with the given 'Size' hint. +unsafeFromList :: Monad m => Size -> [a] -> Bundle m v a +{-# INLINE_FUSED unsafeFromList #-} +unsafeFromList sz xs = fromStream (S.fromList xs) sz + +fromVector :: (Monad m, Vector v a) => v a -> Bundle m v a +{-# INLINE_FUSED fromVector #-} +fromVector v = v `seq` n `seq` Bundle (Stream step 0) + (Stream vstep True) + (Just v) + (Exact n) + where + n = basicLength v + + {-# INLINE step #-} + step i | i >= n = return Done + | otherwise = case basicUnsafeIndexM v i of + Box x -> return $ Yield x (i+1) + + + {-# INLINE vstep #-} + vstep True = return (Yield (Chunk (basicLength v) (\mv -> basicUnsafeCopy mv v)) False) + vstep False = return Done + +fromVectors :: forall m v a. (Monad m, Vector v a) => [v a] -> Bundle m v a +{-# INLINE_FUSED fromVectors #-} +fromVectors us = Bundle (Stream pstep (Left us)) + (Stream vstep us) + Nothing + (Exact n) + where + n = List.foldl' (\k v -> k + basicLength v) 0 us + + pstep (Left []) = return Done + pstep (Left (v:vs)) = basicLength v `seq` return (Skip (Right (v,0,vs))) + + pstep (Right (v,i,vs)) + | i >= basicLength v = return $ Skip (Left vs) + | otherwise = case basicUnsafeIndexM v i of + Box x -> return $ Yield x (Right (v,i+1,vs)) + + -- FIXME: work around bug in GHC 7.6.1 + vstep :: [v a] -> m (Step [v a] (Chunk v a)) + vstep [] = return Done + vstep (v:vs) = return $ Yield (Chunk (basicLength v) + (\mv -> INTERNAL_CHECK(check) "concatVectors" "length mismatch" + (M.basicLength mv == basicLength v) + $ basicUnsafeCopy mv v)) vs + + +concatVectors :: (Monad m, Vector v a) => Bundle m u (v a) -> Bundle m v a +{-# INLINE_FUSED concatVectors #-} +concatVectors Bundle{sElems = Stream step t} + = Bundle (Stream pstep (Left t)) + (Stream vstep t) + Nothing + Unknown + where + pstep (Left s) = do + r <- step s + case r of + Yield v s' -> basicLength v `seq` return (Skip (Right (v,0,s'))) + Skip s' -> return (Skip (Left s')) + Done -> return Done + + pstep (Right (v,i,s)) + | i >= basicLength v = return (Skip (Left s)) + | otherwise = case basicUnsafeIndexM v i of + Box x -> return (Yield x (Right (v,i+1,s))) + + + vstep s = do + r <- step s + case r of + Yield v s' -> return (Yield (Chunk (basicLength v) + (\mv -> INTERNAL_CHECK(check) "concatVectors" "length mismatch" + (M.basicLength mv == basicLength v) + $ basicUnsafeCopy mv v)) s') + Skip s' -> return (Skip s') + Done -> return Done + +reVector :: Monad m => Bundle m u a -> Bundle m v a +{-# INLINE_FUSED reVector #-} +reVector Bundle{sElems = s, sSize = n} = fromStream s n + +{-# RULES + +"reVector [Vector]" + reVector = id + +"reVector/reVector [Vector]" forall s. + reVector (reVector s) = s #-} + + + diff --git a/third_party/bazel/rules_haskell/examples/vector/Data/Vector/Fusion/Bundle/Size.hs b/third_party/bazel/rules_haskell/examples/vector/Data/Vector/Fusion/Bundle/Size.hs new file mode 100644 index 0000000000..e90cf37320 --- /dev/null +++ b/third_party/bazel/rules_haskell/examples/vector/Data/Vector/Fusion/Bundle/Size.hs @@ -0,0 +1,121 @@ +-- | +-- Module : Data.Vector.Fusion.Bundle.Size +-- Copyright : (c) Roman Leshchinskiy 2008-2010 +-- License : BSD-style +-- +-- Maintainer : Roman Leshchinskiy <rl@cse.unsw.edu.au> +-- Stability : experimental +-- Portability : portable +-- +-- Size hints for streams. +-- + +module Data.Vector.Fusion.Bundle.Size ( + Size(..), clampedSubtract, smaller, larger, toMax, upperBound, lowerBound +) where + +import Data.Vector.Fusion.Util ( delay_inline ) + +-- | Size hint +data Size = Exact Int -- ^ Exact size + | Max Int -- ^ Upper bound on the size + | Unknown -- ^ Unknown size + deriving( Eq, Show ) + +instance Num Size where + Exact m + Exact n = checkedAdd Exact m n + Exact m + Max n = checkedAdd Max m n + + Max m + Exact n = checkedAdd Max m n + Max m + Max n = checkedAdd Max m n + + _ + _ = Unknown + + + Exact m - Exact n = checkedSubtract Exact m n + Exact m - Max _ = Max m + + Max m - Exact n = checkedSubtract Max m n + Max m - Max _ = Max m + Max m - Unknown = Max m + + _ - _ = Unknown + + + fromInteger n = Exact (fromInteger n) + + (*) = error "vector: internal error * for Bundle.size isn't defined" + abs = error "vector: internal error abs for Bundle.size isn't defined" + signum = error "vector: internal error signum for Bundle.size isn't defined" + +{-# INLINE checkedAdd #-} +checkedAdd :: (Int -> Size) -> Int -> Int -> Size +checkedAdd con m n + -- Note: we assume m and n are >= 0. + | r < m || r < n = + error $ "Data.Vector.Fusion.Bundle.Size.checkedAdd: overflow: " ++ show r + | otherwise = con r + where + r = m + n + +{-# INLINE checkedSubtract #-} +checkedSubtract :: (Int -> Size) -> Int -> Int -> Size +checkedSubtract con m n + | r < 0 = + error $ "Data.Vector.Fusion.Bundle.Size.checkedSubtract: underflow: " ++ show r + | otherwise = con r + where + r = m - n + +-- | Subtract two sizes with clamping to 0, for drop-like things +{-# INLINE clampedSubtract #-} +clampedSubtract :: Size -> Size -> Size +clampedSubtract (Exact m) (Exact n) = Exact (max 0 (m - n)) +clampedSubtract (Max m) (Exact n) + | m <= n = Exact 0 + | otherwise = Max (m - n) +clampedSubtract (Exact m) (Max _) = Max m +clampedSubtract (Max m) (Max _) = Max m +clampedSubtract _ _ = Unknown + +-- | Minimum of two size hints +smaller :: Size -> Size -> Size +{-# INLINE smaller #-} +smaller (Exact m) (Exact n) = Exact (delay_inline min m n) +smaller (Exact m) (Max n) = Max (delay_inline min m n) +smaller (Exact m) Unknown = Max m +smaller (Max m) (Exact n) = Max (delay_inline min m n) +smaller (Max m) (Max n) = Max (delay_inline min m n) +smaller (Max m) Unknown = Max m +smaller Unknown (Exact n) = Max n +smaller Unknown (Max n) = Max n +smaller Unknown Unknown = Unknown + +-- | Maximum of two size hints +larger :: Size -> Size -> Size +{-# INLINE larger #-} +larger (Exact m) (Exact n) = Exact (delay_inline max m n) +larger (Exact m) (Max n) | m >= n = Exact m + | otherwise = Max n +larger (Max m) (Exact n) | n >= m = Exact n + | otherwise = Max m +larger (Max m) (Max n) = Max (delay_inline max m n) +larger _ _ = Unknown + +-- | Convert a size hint to an upper bound +toMax :: Size -> Size +toMax (Exact n) = Max n +toMax (Max n) = Max n +toMax Unknown = Unknown + +-- | Compute the minimum size from a size hint +lowerBound :: Size -> Int +lowerBound (Exact n) = n +lowerBound _ = 0 + +-- | Compute the maximum size from a size hint if possible +upperBound :: Size -> Maybe Int +upperBound (Exact n) = Just n +upperBound (Max n) = Just n +upperBound Unknown = Nothing + diff --git a/third_party/bazel/rules_haskell/examples/vector/Data/Vector/Fusion/Stream/Monadic.hs b/third_party/bazel/rules_haskell/examples/vector/Data/Vector/Fusion/Stream/Monadic.hs new file mode 100644 index 0000000000..cca002ca6f --- /dev/null +++ b/third_party/bazel/rules_haskell/examples/vector/Data/Vector/Fusion/Stream/Monadic.hs @@ -0,0 +1,1639 @@ +{-# LANGUAGE CPP, ExistentialQuantification, MultiParamTypeClasses, FlexibleInstances, Rank2Types, BangPatterns, KindSignatures, GADTs, ScopedTypeVariables #-} + +-- | +-- Module : Data.Vector.Fusion.Stream.Monadic +-- Copyright : (c) Roman Leshchinskiy 2008-2010 +-- License : BSD-style +-- +-- Maintainer : Roman Leshchinskiy <rl@cse.unsw.edu.au> +-- Stability : experimental +-- Portability : non-portable +-- +-- Monadic stream combinators. +-- + +module Data.Vector.Fusion.Stream.Monadic ( + Stream(..), Step(..), SPEC(..), + + -- * Length + length, null, + + -- * Construction + empty, singleton, cons, snoc, replicate, replicateM, generate, generateM, (++), + + -- * Accessing elements + head, last, (!!), (!?), + + -- * Substreams + slice, init, tail, take, drop, + + -- * Mapping + map, mapM, mapM_, trans, unbox, concatMap, flatten, + + -- * Zipping + indexed, indexedR, zipWithM_, + zipWithM, zipWith3M, zipWith4M, zipWith5M, zipWith6M, + zipWith, zipWith3, zipWith4, zipWith5, zipWith6, + zip, zip3, zip4, zip5, zip6, + + -- * Comparisons + eqBy, cmpBy, + + -- * Filtering + filter, filterM, uniq, mapMaybe, takeWhile, takeWhileM, dropWhile, dropWhileM, + + -- * Searching + elem, notElem, find, findM, findIndex, findIndexM, + + -- * Folding + foldl, foldlM, foldl1, foldl1M, foldM, fold1M, + foldl', foldlM', foldl1', foldl1M', foldM', fold1M', + foldr, foldrM, foldr1, foldr1M, + + -- * Specialised folds + and, or, concatMapM, + + -- * Unfolding + unfoldr, unfoldrM, + unfoldrN, unfoldrNM, + iterateN, iterateNM, + + -- * Scans + prescanl, prescanlM, prescanl', prescanlM', + postscanl, postscanlM, postscanl', postscanlM', + scanl, scanlM, scanl', scanlM', + scanl1, scanl1M, scanl1', scanl1M', + + -- * Enumerations + enumFromStepN, enumFromTo, enumFromThenTo, + + -- * Conversions + toList, fromList, fromListN +) where + +import Data.Vector.Fusion.Util ( Box(..) ) + +import Data.Char ( ord ) +import GHC.Base ( unsafeChr ) +import Control.Monad ( liftM ) +import Prelude hiding ( length, null, + replicate, (++), + head, last, (!!), + init, tail, take, drop, + map, mapM, mapM_, concatMap, + zipWith, zipWith3, zip, zip3, + filter, takeWhile, dropWhile, + elem, notElem, + foldl, foldl1, foldr, foldr1, + and, or, + scanl, scanl1, + enumFromTo, enumFromThenTo ) + +import Data.Int ( Int8, Int16, Int32 ) +import Data.Word ( Word8, Word16, Word32, Word64 ) + +#if !MIN_VERSION_base(4,8,0) +import Data.Word ( Word8, Word16, Word32, Word, Word64 ) +#endif + +#if __GLASGOW_HASKELL__ >= 708 +import GHC.Types ( SPEC(..) ) +#elif __GLASGOW_HASKELL__ >= 700 +import GHC.Exts ( SpecConstrAnnotation(..) ) +#endif + +#include "vector.h" +#include "MachDeps.h" + +#if WORD_SIZE_IN_BITS > 32 +import Data.Int ( Int64 ) +#endif + +#if __GLASGOW_HASKELL__ < 708 +data SPEC = SPEC | SPEC2 +#if __GLASGOW_HASKELL__ >= 700 +{-# ANN type SPEC ForceSpecConstr #-} +#endif +#endif + +emptyStream :: String +{-# NOINLINE emptyStream #-} +emptyStream = "empty stream" + +#define EMPTY_STREAM (\state -> ERROR state emptyStream) + +-- | Result of taking a single step in a stream +data Step s a where + Yield :: a -> s -> Step s a + Skip :: s -> Step s a + Done :: Step s a + +instance Functor (Step s) where + {-# INLINE fmap #-} + fmap f (Yield x s) = Yield (f x) s + fmap _ (Skip s) = Skip s + fmap _ Done = Done + +-- | Monadic streams +data Stream m a = forall s. Stream (s -> m (Step s a)) s + +-- Length +-- ------ + +-- | Length of a 'Stream' +length :: Monad m => Stream m a -> m Int +{-# INLINE_FUSED length #-} +length = foldl' (\n _ -> n+1) 0 + +-- | Check if a 'Stream' is empty +null :: Monad m => Stream m a -> m Bool +{-# INLINE_FUSED null #-} +null (Stream step t) = null_loop t + where + null_loop s = do + r <- step s + case r of + Yield _ _ -> return False + Skip s' -> null_loop s' + Done -> return True + +-- Construction +-- ------------ + +-- | Empty 'Stream' +empty :: Monad m => Stream m a +{-# INLINE_FUSED empty #-} +empty = Stream (const (return Done)) () + +-- | Singleton 'Stream' +singleton :: Monad m => a -> Stream m a +{-# INLINE_FUSED singleton #-} +singleton x = Stream (return . step) True + where + {-# INLINE_INNER step #-} + step True = Yield x False + step False = Done + +-- | Replicate a value to a given length +replicate :: Monad m => Int -> a -> Stream m a +{-# INLINE_FUSED replicate #-} +replicate n x = replicateM n (return x) + +-- | Yield a 'Stream' of values obtained by performing the monadic action the +-- given number of times +replicateM :: Monad m => Int -> m a -> Stream m a +{-# INLINE_FUSED replicateM #-} +replicateM n p = Stream step n + where + {-# INLINE_INNER step #-} + step i | i <= 0 = return Done + | otherwise = do { x <- p; return $ Yield x (i-1) } + +generate :: Monad m => Int -> (Int -> a) -> Stream m a +{-# INLINE generate #-} +generate n f = generateM n (return . f) + +-- | Generate a stream from its indices +generateM :: Monad m => Int -> (Int -> m a) -> Stream m a +{-# INLINE_FUSED generateM #-} +generateM n f = n `seq` Stream step 0 + where + {-# INLINE_INNER step #-} + step i | i < n = do + x <- f i + return $ Yield x (i+1) + | otherwise = return Done + +-- | Prepend an element +cons :: Monad m => a -> Stream m a -> Stream m a +{-# INLINE cons #-} +cons x s = singleton x ++ s + +-- | Append an element +snoc :: Monad m => Stream m a -> a -> Stream m a +{-# INLINE snoc #-} +snoc s x = s ++ singleton x + +infixr 5 ++ +-- | Concatenate two 'Stream's +(++) :: Monad m => Stream m a -> Stream m a -> Stream m a +{-# INLINE_FUSED (++) #-} +Stream stepa ta ++ Stream stepb tb = Stream step (Left ta) + where + {-# INLINE_INNER step #-} + step (Left sa) = do + r <- stepa sa + case r of + Yield x sa' -> return $ Yield x (Left sa') + Skip sa' -> return $ Skip (Left sa') + Done -> return $ Skip (Right tb) + step (Right sb) = do + r <- stepb sb + case r of + Yield x sb' -> return $ Yield x (Right sb') + Skip sb' -> return $ Skip (Right sb') + Done -> return $ Done + +-- Accessing elements +-- ------------------ + +-- | First element of the 'Stream' or error if empty +head :: Monad m => Stream m a -> m a +{-# INLINE_FUSED head #-} +head (Stream step t) = head_loop SPEC t + where + head_loop !_ s + = do + r <- step s + case r of + Yield x _ -> return x + Skip s' -> head_loop SPEC s' + Done -> EMPTY_STREAM "head" + + + +-- | Last element of the 'Stream' or error if empty +last :: Monad m => Stream m a -> m a +{-# INLINE_FUSED last #-} +last (Stream step t) = last_loop0 SPEC t + where + last_loop0 !_ s + = do + r <- step s + case r of + Yield x s' -> last_loop1 SPEC x s' + Skip s' -> last_loop0 SPEC s' + Done -> EMPTY_STREAM "last" + + last_loop1 !_ x s + = do + r <- step s + case r of + Yield y s' -> last_loop1 SPEC y s' + Skip s' -> last_loop1 SPEC x s' + Done -> return x + +infixl 9 !! +-- | Element at the given position +(!!) :: Monad m => Stream m a -> Int -> m a +{-# INLINE (!!) #-} +Stream step t !! j | j < 0 = ERROR "!!" "negative index" + | otherwise = index_loop SPEC t j + where + index_loop !_ s i + = i `seq` + do + r <- step s + case r of + Yield x s' | i == 0 -> return x + | otherwise -> index_loop SPEC s' (i-1) + Skip s' -> index_loop SPEC s' i + Done -> EMPTY_STREAM "!!" + +infixl 9 !? +-- | Element at the given position or 'Nothing' if out of bounds +(!?) :: Monad m => Stream m a -> Int -> m (Maybe a) +{-# INLINE (!?) #-} +Stream step t !? j = index_loop SPEC t j + where + index_loop !_ s i + = i `seq` + do + r <- step s + case r of + Yield x s' | i == 0 -> return (Just x) + | otherwise -> index_loop SPEC s' (i-1) + Skip s' -> index_loop SPEC s' i + Done -> return Nothing + +-- Substreams +-- ---------- + +-- | Extract a substream of the given length starting at the given position. +slice :: Monad m => Int -- ^ starting index + -> Int -- ^ length + -> Stream m a + -> Stream m a +{-# INLINE slice #-} +slice i n s = take n (drop i s) + +-- | All but the last element +init :: Monad m => Stream m a -> Stream m a +{-# INLINE_FUSED init #-} +init (Stream step t) = Stream step' (Nothing, t) + where + {-# INLINE_INNER step' #-} + step' (Nothing, s) = liftM (\r -> + case r of + Yield x s' -> Skip (Just x, s') + Skip s' -> Skip (Nothing, s') + Done -> EMPTY_STREAM "init" + ) (step s) + + step' (Just x, s) = liftM (\r -> + case r of + Yield y s' -> Yield x (Just y, s') + Skip s' -> Skip (Just x, s') + Done -> Done + ) (step s) + +-- | All but the first element +tail :: Monad m => Stream m a -> Stream m a +{-# INLINE_FUSED tail #-} +tail (Stream step t) = Stream step' (Left t) + where + {-# INLINE_INNER step' #-} + step' (Left s) = liftM (\r -> + case r of + Yield _ s' -> Skip (Right s') + Skip s' -> Skip (Left s') + Done -> EMPTY_STREAM "tail" + ) (step s) + + step' (Right s) = liftM (\r -> + case r of + Yield x s' -> Yield x (Right s') + Skip s' -> Skip (Right s') + Done -> Done + ) (step s) + +-- | The first @n@ elements +take :: Monad m => Int -> Stream m a -> Stream m a +{-# INLINE_FUSED take #-} +take n (Stream step t) = n `seq` Stream step' (t, 0) + where + {-# INLINE_INNER step' #-} + step' (s, i) | i < n = liftM (\r -> + case r of + Yield x s' -> Yield x (s', i+1) + Skip s' -> Skip (s', i) + Done -> Done + ) (step s) + step' (_, _) = return Done + +-- | All but the first @n@ elements +drop :: Monad m => Int -> Stream m a -> Stream m a +{-# INLINE_FUSED drop #-} +drop n (Stream step t) = Stream step' (t, Just n) + where + {-# INLINE_INNER step' #-} + step' (s, Just i) | i > 0 = liftM (\r -> + case r of + Yield _ s' -> Skip (s', Just (i-1)) + Skip s' -> Skip (s', Just i) + Done -> Done + ) (step s) + | otherwise = return $ Skip (s, Nothing) + + step' (s, Nothing) = liftM (\r -> + case r of + Yield x s' -> Yield x (s', Nothing) + Skip s' -> Skip (s', Nothing) + Done -> Done + ) (step s) + +-- Mapping +-- ------- + +instance Monad m => Functor (Stream m) where + {-# INLINE fmap #-} + fmap = map + +-- | Map a function over a 'Stream' +map :: Monad m => (a -> b) -> Stream m a -> Stream m b +{-# INLINE map #-} +map f = mapM (return . f) + + +-- | Map a monadic function over a 'Stream' +mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b +{-# INLINE_FUSED mapM #-} +mapM f (Stream step t) = Stream step' t + where + {-# INLINE_INNER step' #-} + step' s = do + r <- step s + case r of + Yield x s' -> liftM (`Yield` s') (f x) + Skip s' -> return (Skip s') + Done -> return Done + +consume :: Monad m => Stream m a -> m () +{-# INLINE_FUSED consume #-} +consume (Stream step t) = consume_loop SPEC t + where + consume_loop !_ s + = do + r <- step s + case r of + Yield _ s' -> consume_loop SPEC s' + Skip s' -> consume_loop SPEC s' + Done -> return () + +-- | Execute a monadic action for each element of the 'Stream' +mapM_ :: Monad m => (a -> m b) -> Stream m a -> m () +{-# INLINE_FUSED mapM_ #-} +mapM_ m = consume . mapM m + +-- | Transform a 'Stream' to use a different monad +trans :: (Monad m, Monad m') + => (forall z. m z -> m' z) -> Stream m a -> Stream m' a +{-# INLINE_FUSED trans #-} +trans f (Stream step s) = Stream (f . step) s + +unbox :: Monad m => Stream m (Box a) -> Stream m a +{-# INLINE_FUSED unbox #-} +unbox (Stream step t) = Stream step' t + where + {-# INLINE_INNER step' #-} + step' s = do + r <- step s + case r of + Yield (Box x) s' -> return $ Yield x s' + Skip s' -> return $ Skip s' + Done -> return $ Done + +-- Zipping +-- ------- + +-- | Pair each element in a 'Stream' with its index +indexed :: Monad m => Stream m a -> Stream m (Int,a) +{-# INLINE_FUSED indexed #-} +indexed (Stream step t) = Stream step' (t,0) + where + {-# INLINE_INNER step' #-} + step' (s,i) = i `seq` + do + r <- step s + case r of + Yield x s' -> return $ Yield (i,x) (s', i+1) + Skip s' -> return $ Skip (s', i) + Done -> return Done + +-- | Pair each element in a 'Stream' with its index, starting from the right +-- and counting down +indexedR :: Monad m => Int -> Stream m a -> Stream m (Int,a) +{-# INLINE_FUSED indexedR #-} +indexedR m (Stream step t) = Stream step' (t,m) + where + {-# INLINE_INNER step' #-} + step' (s,i) = i `seq` + do + r <- step s + case r of + Yield x s' -> let i' = i-1 + in + return $ Yield (i',x) (s', i') + Skip s' -> return $ Skip (s', i) + Done -> return Done + +-- | Zip two 'Stream's with the given monadic function +zipWithM :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c +{-# INLINE_FUSED zipWithM #-} +zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing) + where + {-# INLINE_INNER step #-} + step (sa, sb, Nothing) = liftM (\r -> + case r of + Yield x sa' -> Skip (sa', sb, Just x) + Skip sa' -> Skip (sa', sb, Nothing) + Done -> Done + ) (stepa sa) + + step (sa, sb, Just x) = do + r <- stepb sb + case r of + Yield y sb' -> + do + z <- f x y + return $ Yield z (sa, sb', Nothing) + Skip sb' -> return $ Skip (sa, sb', Just x) + Done -> return $ Done + +-- FIXME: This might expose an opportunity for inplace execution. +{-# RULES + +"zipWithM xs xs [Vector.Stream]" forall f xs. + zipWithM f xs xs = mapM (\x -> f x x) xs #-} + + +zipWithM_ :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> m () +{-# INLINE zipWithM_ #-} +zipWithM_ f sa sb = consume (zipWithM f sa sb) + +zipWith3M :: Monad m => (a -> b -> c -> m d) -> Stream m a -> Stream m b -> Stream m c -> Stream m d +{-# INLINE_FUSED zipWith3M #-} +zipWith3M f (Stream stepa ta) + (Stream stepb tb) + (Stream stepc tc) = Stream step (ta, tb, tc, Nothing) + where + {-# INLINE_INNER step #-} + step (sa, sb, sc, Nothing) = do + r <- stepa sa + return $ case r of + Yield x sa' -> Skip (sa', sb, sc, Just (x, Nothing)) + Skip sa' -> Skip (sa', sb, sc, Nothing) + Done -> Done + + step (sa, sb, sc, Just (x, Nothing)) = do + r <- stepb sb + return $ case r of + Yield y sb' -> Skip (sa, sb', sc, Just (x, Just y)) + Skip sb' -> Skip (sa, sb', sc, Just (x, Nothing)) + Done -> Done + + step (sa, sb, sc, Just (x, Just y)) = do + r <- stepc sc + case r of + Yield z sc' -> f x y z >>= (\res -> return $ Yield res (sa, sb, sc', Nothing)) + Skip sc' -> return $ Skip (sa, sb, sc', Just (x, Just y)) + Done -> return $ Done + +zipWith4M :: Monad m => (a -> b -> c -> d -> m e) + -> Stream m a -> Stream m b -> Stream m c -> Stream m d + -> Stream m e +{-# INLINE zipWith4M #-} +zipWith4M f sa sb sc sd + = zipWithM (\(a,b) (c,d) -> f a b c d) (zip sa sb) (zip sc sd) + +zipWith5M :: Monad m => (a -> b -> c -> d -> e -> m f) + -> Stream m a -> Stream m b -> Stream m c -> Stream m d + -> Stream m e -> Stream m f +{-# INLINE zipWith5M #-} +zipWith5M f sa sb sc sd se + = zipWithM (\(a,b,c) (d,e) -> f a b c d e) (zip3 sa sb sc) (zip sd se) + +zipWith6M :: Monad m => (a -> b -> c -> d -> e -> f -> m g) + -> Stream m a -> Stream m b -> Stream m c -> Stream m d + -> Stream m e -> Stream m f -> Stream m g +{-# INLINE zipWith6M #-} +zipWith6M fn sa sb sc sd se sf + = zipWithM (\(a,b,c) (d,e,f) -> fn a b c d e f) (zip3 sa sb sc) + (zip3 sd se sf) + +zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c +{-# INLINE zipWith #-} +zipWith f = zipWithM (\a b -> return (f a b)) + +zipWith3 :: Monad m => (a -> b -> c -> d) + -> Stream m a -> Stream m b -> Stream m c -> Stream m d +{-# INLINE zipWith3 #-} +zipWith3 f = zipWith3M (\a b c -> return (f a b c)) + +zipWith4 :: Monad m => (a -> b -> c -> d -> e) + -> Stream m a -> Stream m b -> Stream m c -> Stream m d + -> Stream m e +{-# INLINE zipWith4 #-} +zipWith4 f = zipWith4M (\a b c d -> return (f a b c d)) + +zipWith5 :: Monad m => (a -> b -> c -> d -> e -> f) + -> Stream m a -> Stream m b -> Stream m c -> Stream m d + -> Stream m e -> Stream m f +{-# INLINE zipWith5 #-} +zipWith5 f = zipWith5M (\a b c d e -> return (f a b c d e)) + +zipWith6 :: Monad m => (a -> b -> c -> d -> e -> f -> g) + -> Stream m a -> Stream m b -> Stream m c -> Stream m d + -> Stream m e -> Stream m f -> Stream m g +{-# INLINE zipWith6 #-} +zipWith6 fn = zipWith6M (\a b c d e f -> return (fn a b c d e f)) + +zip :: Monad m => Stream m a -> Stream m b -> Stream m (a,b) +{-# INLINE zip #-} +zip = zipWith (,) + +zip3 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m (a,b,c) +{-# INLINE zip3 #-} +zip3 = zipWith3 (,,) + +zip4 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m d + -> Stream m (a,b,c,d) +{-# INLINE zip4 #-} +zip4 = zipWith4 (,,,) + +zip5 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m d + -> Stream m e -> Stream m (a,b,c,d,e) +{-# INLINE zip5 #-} +zip5 = zipWith5 (,,,,) + +zip6 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m d + -> Stream m e -> Stream m f -> Stream m (a,b,c,d,e,f) +{-# INLINE zip6 #-} +zip6 = zipWith6 (,,,,,) + +-- Comparisons +-- ----------- + +-- | Check if two 'Stream's are equal +eqBy :: (Monad m) => (a -> b -> Bool) -> Stream m a -> Stream m b -> m Bool +{-# INLINE_FUSED eqBy #-} +eqBy eq (Stream step1 t1) (Stream step2 t2) = eq_loop0 SPEC t1 t2 + where + eq_loop0 !_ s1 s2 = do + r <- step1 s1 + case r of + Yield x s1' -> eq_loop1 SPEC x s1' s2 + Skip s1' -> eq_loop0 SPEC s1' s2 + Done -> eq_null s2 + + eq_loop1 !_ x s1 s2 = do + r <- step2 s2 + case r of + Yield y s2' + | eq x y -> eq_loop0 SPEC s1 s2' + | otherwise -> return False + Skip s2' -> eq_loop1 SPEC x s1 s2' + Done -> return False + + eq_null s2 = do + r <- step2 s2 + case r of + Yield _ _ -> return False + Skip s2' -> eq_null s2' + Done -> return True + +-- | Lexicographically compare two 'Stream's +cmpBy :: (Monad m) => (a -> b -> Ordering) -> Stream m a -> Stream m b -> m Ordering +{-# INLINE_FUSED cmpBy #-} +cmpBy cmp (Stream step1 t1) (Stream step2 t2) = cmp_loop0 SPEC t1 t2 + where + cmp_loop0 !_ s1 s2 = do + r <- step1 s1 + case r of + Yield x s1' -> cmp_loop1 SPEC x s1' s2 + Skip s1' -> cmp_loop0 SPEC s1' s2 + Done -> cmp_null s2 + + cmp_loop1 !_ x s1 s2 = do + r <- step2 s2 + case r of + Yield y s2' -> case x `cmp` y of + EQ -> cmp_loop0 SPEC s1 s2' + c -> return c + Skip s2' -> cmp_loop1 SPEC x s1 s2' + Done -> return GT + + cmp_null s2 = do + r <- step2 s2 + case r of + Yield _ _ -> return LT + Skip s2' -> cmp_null s2' + Done -> return EQ + +-- Filtering +-- --------- + +-- | Drop elements which do not satisfy the predicate +filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a +{-# INLINE filter #-} +filter f = filterM (return . f) + +mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b +{-# INLINE_FUSED mapMaybe #-} +mapMaybe f (Stream step t) = Stream step' t + where + {-# INLINE_INNER step' #-} + step' s = do + r <- step s + case r of + Yield x s' -> do + return $ case f x of + Nothing -> Skip s' + Just b' -> Yield b' s' + Skip s' -> return $ Skip s' + Done -> return $ Done + +-- | Drop elements which do not satisfy the monadic predicate +filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a +{-# INLINE_FUSED filterM #-} +filterM f (Stream step t) = Stream step' t + where + {-# INLINE_INNER step' #-} + step' s = do + r <- step s + case r of + Yield x s' -> do + b <- f x + return $ if b then Yield x s' + else Skip s' + Skip s' -> return $ Skip s' + Done -> return $ Done + +-- | Drop repeated adjacent elements. +uniq :: (Eq a, Monad m) => Stream m a -> Stream m a +{-# INLINE_FUSED uniq #-} +uniq (Stream step st) = Stream step' (Nothing,st) + where + {-# INLINE_INNER step' #-} + step' (Nothing, s) = do r <- step s + case r of + Yield x s' -> return $ Yield x (Just x , s') + Skip s' -> return $ Skip (Nothing, s') + Done -> return Done + step' (Just x0, s) = do r <- step s + case r of + Yield x s' | x == x0 -> return $ Skip (Just x0, s') + | otherwise -> return $ Yield x (Just x , s') + Skip s' -> return $ Skip (Just x0, s') + Done -> return Done + +-- | Longest prefix of elements that satisfy the predicate +takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a +{-# INLINE takeWhile #-} +takeWhile f = takeWhileM (return . f) + +-- | Longest prefix of elements that satisfy the monadic predicate +takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a +{-# INLINE_FUSED takeWhileM #-} +takeWhileM f (Stream step t) = Stream step' t + where + {-# INLINE_INNER step' #-} + step' s = do + r <- step s + case r of + Yield x s' -> do + b <- f x + return $ if b then Yield x s' else Done + Skip s' -> return $ Skip s' + Done -> return $ Done + +-- | Drop the longest prefix of elements that satisfy the predicate +dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a +{-# INLINE dropWhile #-} +dropWhile f = dropWhileM (return . f) + +data DropWhile s a = DropWhile_Drop s | DropWhile_Yield a s | DropWhile_Next s + +-- | Drop the longest prefix of elements that satisfy the monadic predicate +dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a +{-# INLINE_FUSED dropWhileM #-} +dropWhileM f (Stream step t) = Stream step' (DropWhile_Drop t) + where + -- NOTE: we jump through hoops here to have only one Yield; local data + -- declarations would be nice! + + {-# INLINE_INNER step' #-} + step' (DropWhile_Drop s) + = do + r <- step s + case r of + Yield x s' -> do + b <- f x + return $ if b then Skip (DropWhile_Drop s') + else Skip (DropWhile_Yield x s') + Skip s' -> return $ Skip (DropWhile_Drop s') + Done -> return $ Done + + step' (DropWhile_Yield x s) = return $ Yield x (DropWhile_Next s) + + step' (DropWhile_Next s) + = liftM (\r -> + case r of + Yield x s' -> Skip (DropWhile_Yield x s') + Skip s' -> Skip (DropWhile_Next s') + Done -> Done + ) (step s) + +-- Searching +-- --------- + +infix 4 `elem` +-- | Check whether the 'Stream' contains an element +elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool +{-# INLINE_FUSED elem #-} +elem x (Stream step t) = elem_loop SPEC t + where + elem_loop !_ s + = do + r <- step s + case r of + Yield y s' | x == y -> return True + | otherwise -> elem_loop SPEC s' + Skip s' -> elem_loop SPEC s' + Done -> return False + +infix 4 `notElem` +-- | Inverse of `elem` +notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool +{-# INLINE notElem #-} +notElem x s = liftM not (elem x s) + +-- | Yield 'Just' the first element that satisfies the predicate or 'Nothing' +-- if no such element exists. +find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a) +{-# INLINE find #-} +find f = findM (return . f) + +-- | Yield 'Just' the first element that satisfies the monadic predicate or +-- 'Nothing' if no such element exists. +findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a) +{-# INLINE_FUSED findM #-} +findM f (Stream step t) = find_loop SPEC t + where + find_loop !_ s + = do + r <- step s + case r of + Yield x s' -> do + b <- f x + if b then return $ Just x + else find_loop SPEC s' + Skip s' -> find_loop SPEC s' + Done -> return Nothing + +-- | Yield 'Just' the index of the first element that satisfies the predicate +-- or 'Nothing' if no such element exists. +findIndex :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe Int) +{-# INLINE_FUSED findIndex #-} +findIndex f = findIndexM (return . f) + +-- | Yield 'Just' the index of the first element that satisfies the monadic +-- predicate or 'Nothing' if no such element exists. +findIndexM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe Int) +{-# INLINE_FUSED findIndexM #-} +findIndexM f (Stream step t) = findIndex_loop SPEC t 0 + where + findIndex_loop !_ s i + = do + r <- step s + case r of + Yield x s' -> do + b <- f x + if b then return $ Just i + else findIndex_loop SPEC s' (i+1) + Skip s' -> findIndex_loop SPEC s' i + Done -> return Nothing + +-- Folding +-- ------- + +-- | Left fold +foldl :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a +{-# INLINE foldl #-} +foldl f = foldlM (\a b -> return (f a b)) + +-- | Left fold with a monadic operator +foldlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a +{-# INLINE_FUSED foldlM #-} +foldlM m w (Stream step t) = foldlM_loop SPEC w t + where + foldlM_loop !_ z s + = do + r <- step s + case r of + Yield x s' -> do { z' <- m z x; foldlM_loop SPEC z' s' } + Skip s' -> foldlM_loop SPEC z s' + Done -> return z + +-- | Same as 'foldlM' +foldM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a +{-# INLINE foldM #-} +foldM = foldlM + +-- | Left fold over a non-empty 'Stream' +foldl1 :: Monad m => (a -> a -> a) -> Stream m a -> m a +{-# INLINE foldl1 #-} +foldl1 f = foldl1M (\a b -> return (f a b)) + +-- | Left fold over a non-empty 'Stream' with a monadic operator +foldl1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a +{-# INLINE_FUSED foldl1M #-} +foldl1M f (Stream step t) = foldl1M_loop SPEC t + where + foldl1M_loop !_ s + = do + r <- step s + case r of + Yield x s' -> foldlM f x (Stream step s') + Skip s' -> foldl1M_loop SPEC s' + Done -> EMPTY_STREAM "foldl1M" + +-- | Same as 'foldl1M' +fold1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a +{-# INLINE fold1M #-} +fold1M = foldl1M + +-- | Left fold with a strict accumulator +foldl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a +{-# INLINE foldl' #-} +foldl' f = foldlM' (\a b -> return (f a b)) + +-- | Left fold with a strict accumulator and a monadic operator +foldlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a +{-# INLINE_FUSED foldlM' #-} +foldlM' m w (Stream step t) = foldlM'_loop SPEC w t + where + foldlM'_loop !_ z s + = z `seq` + do + r <- step s + case r of + Yield x s' -> do { z' <- m z x; foldlM'_loop SPEC z' s' } + Skip s' -> foldlM'_loop SPEC z s' + Done -> return z + +-- | Same as 'foldlM'' +foldM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a +{-# INLINE foldM' #-} +foldM' = foldlM' + +-- | Left fold over a non-empty 'Stream' with a strict accumulator +foldl1' :: Monad m => (a -> a -> a) -> Stream m a -> m a +{-# INLINE foldl1' #-} +foldl1' f = foldl1M' (\a b -> return (f a b)) + +-- | Left fold over a non-empty 'Stream' with a strict accumulator and a +-- monadic operator +foldl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> m a +{-# INLINE_FUSED foldl1M' #-} +foldl1M' f (Stream step t) = foldl1M'_loop SPEC t + where + foldl1M'_loop !_ s + = do + r <- step s + case r of + Yield x s' -> foldlM' f x (Stream step s') + Skip s' -> foldl1M'_loop SPEC s' + Done -> EMPTY_STREAM "foldl1M'" + +-- | Same as 'foldl1M'' +fold1M' :: Monad m => (a -> a -> m a) -> Stream m a -> m a +{-# INLINE fold1M' #-} +fold1M' = foldl1M' + +-- | Right fold +foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b +{-# INLINE foldr #-} +foldr f = foldrM (\a b -> return (f a b)) + +-- | Right fold with a monadic operator +foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m a -> m b +{-# INLINE_FUSED foldrM #-} +foldrM f z (Stream step t) = foldrM_loop SPEC t + where + foldrM_loop !_ s + = do + r <- step s + case r of + Yield x s' -> f x =<< foldrM_loop SPEC s' + Skip s' -> foldrM_loop SPEC s' + Done -> return z + +-- | Right fold over a non-empty stream +foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m a +{-# INLINE foldr1 #-} +foldr1 f = foldr1M (\a b -> return (f a b)) + +-- | Right fold over a non-empty stream with a monadic operator +foldr1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a +{-# INLINE_FUSED foldr1M #-} +foldr1M f (Stream step t) = foldr1M_loop0 SPEC t + where + foldr1M_loop0 !_ s + = do + r <- step s + case r of + Yield x s' -> foldr1M_loop1 SPEC x s' + Skip s' -> foldr1M_loop0 SPEC s' + Done -> EMPTY_STREAM "foldr1M" + + foldr1M_loop1 !_ x s + = do + r <- step s + case r of + Yield y s' -> f x =<< foldr1M_loop1 SPEC y s' + Skip s' -> foldr1M_loop1 SPEC x s' + Done -> return x + +-- Specialised folds +-- ----------------- + +and :: Monad m => Stream m Bool -> m Bool +{-# INLINE_FUSED and #-} +and (Stream step t) = and_loop SPEC t + where + and_loop !_ s + = do + r <- step s + case r of + Yield False _ -> return False + Yield True s' -> and_loop SPEC s' + Skip s' -> and_loop SPEC s' + Done -> return True + +or :: Monad m => Stream m Bool -> m Bool +{-# INLINE_FUSED or #-} +or (Stream step t) = or_loop SPEC t + where + or_loop !_ s + = do + r <- step s + case r of + Yield False s' -> or_loop SPEC s' + Yield True _ -> return True + Skip s' -> or_loop SPEC s' + Done -> return False + +concatMap :: Monad m => (a -> Stream m b) -> Stream m a -> Stream m b +{-# INLINE concatMap #-} +concatMap f = concatMapM (return . f) + +concatMapM :: Monad m => (a -> m (Stream m b)) -> Stream m a -> Stream m b +{-# INLINE_FUSED concatMapM #-} +concatMapM f (Stream step t) = Stream concatMap_go (Left t) + where + concatMap_go (Left s) = do + r <- step s + case r of + Yield a s' -> do + b_stream <- f a + return $ Skip (Right (b_stream, s')) + Skip s' -> return $ Skip (Left s') + Done -> return Done + concatMap_go (Right (Stream inner_step inner_s, s)) = do + r <- inner_step inner_s + case r of + Yield b inner_s' -> return $ Yield b (Right (Stream inner_step inner_s', s)) + Skip inner_s' -> return $ Skip (Right (Stream inner_step inner_s', s)) + Done -> return $ Skip (Left s) + +-- | Create a 'Stream' of values from a 'Stream' of streamable things +flatten :: Monad m => (a -> m s) -> (s -> m (Step s b)) -> Stream m a -> Stream m b +{-# INLINE_FUSED flatten #-} +flatten mk istep (Stream ostep u) = Stream step (Left u) + where + {-# INLINE_INNER step #-} + step (Left t) = do + r <- ostep t + case r of + Yield a t' -> do + s <- mk a + s `seq` return (Skip (Right (s,t'))) + Skip t' -> return $ Skip (Left t') + Done -> return $ Done + + + step (Right (s,t)) = do + r <- istep s + case r of + Yield x s' -> return $ Yield x (Right (s',t)) + Skip s' -> return $ Skip (Right (s',t)) + Done -> return $ Skip (Left t) + +-- Unfolding +-- --------- + +-- | Unfold +unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a +{-# INLINE_FUSED unfoldr #-} +unfoldr f = unfoldrM (return . f) + +-- | Unfold with a monadic function +unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a +{-# INLINE_FUSED unfoldrM #-} +unfoldrM f t = Stream step t + where + {-# INLINE_INNER step #-} + step s = liftM (\r -> + case r of + Just (x, s') -> Yield x s' + Nothing -> Done + ) (f s) + +unfoldrN :: Monad m => Int -> (s -> Maybe (a, s)) -> s -> Stream m a +{-# INLINE_FUSED unfoldrN #-} +unfoldrN n f = unfoldrNM n (return . f) + +-- | Unfold at most @n@ elements with a monadic functions +unfoldrNM :: Monad m => Int -> (s -> m (Maybe (a, s))) -> s -> Stream m a +{-# INLINE_FUSED unfoldrNM #-} +unfoldrNM m f t = Stream step (t,m) + where + {-# INLINE_INNER step #-} + step (s,n) | n <= 0 = return Done + | otherwise = liftM (\r -> + case r of + Just (x,s') -> Yield x (s',n-1) + Nothing -> Done + ) (f s) + +-- | Apply monadic function n times to value. Zeroth element is original value. +iterateNM :: Monad m => Int -> (a -> m a) -> a -> Stream m a +{-# INLINE_FUSED iterateNM #-} +iterateNM n f x0 = Stream step (x0,n) + where + {-# INLINE_INNER step #-} + step (x,i) | i <= 0 = return Done + | i == n = return $ Yield x (x,i-1) + | otherwise = do a <- f x + return $ Yield a (a,i-1) + +-- | Apply function n times to value. Zeroth element is original value. +iterateN :: Monad m => Int -> (a -> a) -> a -> Stream m a +{-# INLINE_FUSED iterateN #-} +iterateN n f x0 = iterateNM n (return . f) x0 + +-- Scans +-- ----- + +-- | Prefix scan +prescanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a +{-# INLINE prescanl #-} +prescanl f = prescanlM (\a b -> return (f a b)) + +-- | Prefix scan with a monadic operator +prescanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a +{-# INLINE_FUSED prescanlM #-} +prescanlM f w (Stream step t) = Stream step' (t,w) + where + {-# INLINE_INNER step' #-} + step' (s,x) = do + r <- step s + case r of + Yield y s' -> do + z <- f x y + return $ Yield x (s', z) + Skip s' -> return $ Skip (s', x) + Done -> return Done + +-- | Prefix scan with strict accumulator +prescanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a +{-# INLINE prescanl' #-} +prescanl' f = prescanlM' (\a b -> return (f a b)) + +-- | Prefix scan with strict accumulator and a monadic operator +prescanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a +{-# INLINE_FUSED prescanlM' #-} +prescanlM' f w (Stream step t) = Stream step' (t,w) + where + {-# INLINE_INNER step' #-} + step' (s,x) = x `seq` + do + r <- step s + case r of + Yield y s' -> do + z <- f x y + return $ Yield x (s', z) + Skip s' -> return $ Skip (s', x) + Done -> return Done + +-- | Suffix scan +postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a +{-# INLINE postscanl #-} +postscanl f = postscanlM (\a b -> return (f a b)) + +-- | Suffix scan with a monadic operator +postscanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a +{-# INLINE_FUSED postscanlM #-} +postscanlM f w (Stream step t) = Stream step' (t,w) + where + {-# INLINE_INNER step' #-} + step' (s,x) = do + r <- step s + case r of + Yield y s' -> do + z <- f x y + return $ Yield z (s',z) + Skip s' -> return $ Skip (s',x) + Done -> return Done + +-- | Suffix scan with strict accumulator +postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a +{-# INLINE postscanl' #-} +postscanl' f = postscanlM' (\a b -> return (f a b)) + +-- | Suffix scan with strict acccumulator and a monadic operator +postscanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a +{-# INLINE_FUSED postscanlM' #-} +postscanlM' f w (Stream step t) = w `seq` Stream step' (t,w) + where + {-# INLINE_INNER step' #-} + step' (s,x) = x `seq` + do + r <- step s + case r of + Yield y s' -> do + z <- f x y + z `seq` return (Yield z (s',z)) + Skip s' -> return $ Skip (s',x) + Done -> return Done + +-- | Haskell-style scan +scanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a +{-# INLINE scanl #-} +scanl f = scanlM (\a b -> return (f a b)) + +-- | Haskell-style scan with a monadic operator +scanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a +{-# INLINE scanlM #-} +scanlM f z s = z `cons` postscanlM f z s + +-- | Haskell-style scan with strict accumulator +scanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a +{-# INLINE scanl' #-} +scanl' f = scanlM' (\a b -> return (f a b)) + +-- | Haskell-style scan with strict accumulator and a monadic operator +scanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a +{-# INLINE scanlM' #-} +scanlM' f z s = z `seq` (z `cons` postscanlM f z s) + +-- | Scan over a non-empty 'Stream' +scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a +{-# INLINE scanl1 #-} +scanl1 f = scanl1M (\x y -> return (f x y)) + +-- | Scan over a non-empty 'Stream' with a monadic operator +scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a +{-# INLINE_FUSED scanl1M #-} +scanl1M f (Stream step t) = Stream step' (t, Nothing) + where + {-# INLINE_INNER step' #-} + step' (s, Nothing) = do + r <- step s + case r of + Yield x s' -> return $ Yield x (s', Just x) + Skip s' -> return $ Skip (s', Nothing) + Done -> EMPTY_STREAM "scanl1M" + + step' (s, Just x) = do + r <- step s + case r of + Yield y s' -> do + z <- f x y + return $ Yield z (s', Just z) + Skip s' -> return $ Skip (s', Just x) + Done -> return Done + +-- | Scan over a non-empty 'Stream' with a strict accumulator +scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a +{-# INLINE scanl1' #-} +scanl1' f = scanl1M' (\x y -> return (f x y)) + +-- | Scan over a non-empty 'Stream' with a strict accumulator and a monadic +-- operator +scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a +{-# INLINE_FUSED scanl1M' #-} +scanl1M' f (Stream step t) = Stream step' (t, Nothing) + where + {-# INLINE_INNER step' #-} + step' (s, Nothing) = do + r <- step s + case r of + Yield x s' -> x `seq` return (Yield x (s', Just x)) + Skip s' -> return $ Skip (s', Nothing) + Done -> EMPTY_STREAM "scanl1M" + + step' (s, Just x) = x `seq` + do + r <- step s + case r of + Yield y s' -> do + z <- f x y + z `seq` return (Yield z (s', Just z)) + Skip s' -> return $ Skip (s', Just x) + Done -> return Done + +-- Enumerations +-- ------------ + +-- The Enum class is broken for this, there just doesn't seem to be a +-- way to implement this generically. We have to specialise for as many types +-- as we can but this doesn't help in polymorphic loops. + +-- | Yield a 'Stream' of the given length containing the values @x@, @x+y@, +-- @x+y+y@ etc. +enumFromStepN :: (Num a, Monad m) => a -> a -> Int -> Stream m a +{-# INLINE_FUSED enumFromStepN #-} +enumFromStepN x y n = x `seq` y `seq` n `seq` Stream step (x,n) + where + {-# INLINE_INNER step #-} + step (w,m) | m > 0 = return $ Yield w (w+y,m-1) + | otherwise = return $ Done + +-- | Enumerate values +-- +-- /WARNING:/ This operation can be very inefficient. If at all possible, use +-- 'enumFromStepN' instead. +enumFromTo :: (Enum a, Monad m) => a -> a -> Stream m a +{-# INLINE_FUSED enumFromTo #-} +enumFromTo x y = fromList [x .. y] + +-- NOTE: We use (x+1) instead of (succ x) below because the latter checks for +-- overflow which can't happen here. + +-- FIXME: add "too large" test for Int +enumFromTo_small :: (Integral a, Monad m) => a -> a -> Stream m a +{-# INLINE_FUSED enumFromTo_small #-} +enumFromTo_small x y = x `seq` y `seq` Stream step x + where + {-# INLINE_INNER step #-} + step w | w <= y = return $ Yield w (w+1) + | otherwise = return $ Done + +{-# RULES + +"enumFromTo<Int8> [Stream]" + enumFromTo = enumFromTo_small :: Monad m => Int8 -> Int8 -> Stream m Int8 + +"enumFromTo<Int16> [Stream]" + enumFromTo = enumFromTo_small :: Monad m => Int16 -> Int16 -> Stream m Int16 + +"enumFromTo<Word8> [Stream]" + enumFromTo = enumFromTo_small :: Monad m => Word8 -> Word8 -> Stream m Word8 + +"enumFromTo<Word16> [Stream]" + enumFromTo = enumFromTo_small :: Monad m => Word16 -> Word16 -> Stream m Word16 #-} + + +#if WORD_SIZE_IN_BITS > 32 + +{-# RULES + +"enumFromTo<Int32> [Stream]" + enumFromTo = enumFromTo_small :: Monad m => Int32 -> Int32 -> Stream m Int32 + +"enumFromTo<Word32> [Stream]" + enumFromTo = enumFromTo_small :: Monad m => Word32 -> Word32 -> Stream m Word32 #-} + + +#endif + +-- NOTE: We could implement a generic "too large" test: +-- +-- len x y | x > y = 0 +-- | n > 0 && n <= fromIntegral (maxBound :: Int) = fromIntegral n +-- | otherwise = error +-- where +-- n = y-x+1 +-- +-- Alas, GHC won't eliminate unnecessary comparisons (such as n >= 0 for +-- unsigned types). See http://hackage.haskell.org/trac/ghc/ticket/3744 +-- + +enumFromTo_int :: forall m. Monad m => Int -> Int -> Stream m Int +{-# INLINE_FUSED enumFromTo_int #-} +enumFromTo_int x y = x `seq` y `seq` Stream step x + where + -- {-# INLINE [0] len #-} + -- len :: Int -> Int -> Int + -- len u v | u > v = 0 + -- | otherwise = BOUNDS_CHECK(check) "enumFromTo" "vector too large" + -- (n > 0) + -- $ n + -- where + -- n = v-u+1 + + {-# INLINE_INNER step #-} + step z | z <= y = return $ Yield z (z+1) + | otherwise = return $ Done + +enumFromTo_intlike :: (Integral a, Monad m) => a -> a -> Stream m a +{-# INLINE_FUSED enumFromTo_intlike #-} +enumFromTo_intlike x y = x `seq` y `seq` Stream step x + where + {-# INLINE_INNER step #-} + step z | z <= y = return $ Yield z (z+1) + | otherwise = return $ Done + +{-# RULES + +"enumFromTo<Int> [Stream]" + enumFromTo = enumFromTo_int :: Monad m => Int -> Int -> Stream m Int + +#if WORD_SIZE_IN_BITS > 32 + +"enumFromTo<Int64> [Stream]" + enumFromTo = enumFromTo_intlike :: Monad m => Int64 -> Int64 -> Stream m Int64 #-} + +#else + +"enumFromTo<Int32> [Stream]" + enumFromTo = enumFromTo_intlike :: Monad m => Int32 -> Int32 -> Stream m Int32 #-} + +#endif + +enumFromTo_big_word :: (Integral a, Monad m) => a -> a -> Stream m a +{-# INLINE_FUSED enumFromTo_big_word #-} +enumFromTo_big_word x y = x `seq` y `seq` Stream step x + where + {-# INLINE_INNER step #-} + step z | z <= y = return $ Yield z (z+1) + | otherwise = return $ Done + +{-# RULES + +"enumFromTo<Word> [Stream]" + enumFromTo = enumFromTo_big_word :: Monad m => Word -> Word -> Stream m Word + +"enumFromTo<Word64> [Stream]" + enumFromTo = enumFromTo_big_word + :: Monad m => Word64 -> Word64 -> Stream m Word64 + +#if WORD_SIZE_IN_BITS == 32 + +"enumFromTo<Word32> [Stream]" + enumFromTo = enumFromTo_big_word + :: Monad m => Word32 -> Word32 -> Stream m Word32 + +#endif + +"enumFromTo<Integer> [Stream]" + enumFromTo = enumFromTo_big_word + :: Monad m => Integer -> Integer -> Stream m Integer #-} + + + +#if WORD_SIZE_IN_BITS > 32 + +-- FIXME: the "too large" test is totally wrong +enumFromTo_big_int :: (Integral a, Monad m) => a -> a -> Stream m a +{-# INLINE_FUSED enumFromTo_big_int #-} +enumFromTo_big_int x y = x `seq` y `seq` Stream step x + where + {-# INLINE_INNER step #-} + step z | z <= y = return $ Yield z (z+1) + | otherwise = return $ Done + +{-# RULES + +"enumFromTo<Int64> [Stream]" + enumFromTo = enumFromTo_big_int :: Monad m => Int64 -> Int64 -> Stream m Int64 #-} + + + +#endif + +enumFromTo_char :: Monad m => Char -> Char -> Stream m Char +{-# INLINE_FUSED enumFromTo_char #-} +enumFromTo_char x y = x `seq` y `seq` Stream step xn + where + xn = ord x + yn = ord y + + {-# INLINE_INNER step #-} + step zn | zn <= yn = return $ Yield (unsafeChr zn) (zn+1) + | otherwise = return $ Done + +{-# RULES + +"enumFromTo<Char> [Stream]" + enumFromTo = enumFromTo_char #-} + + + +------------------------------------------------------------------------ + +-- Specialise enumFromTo for Float and Double. +-- Also, try to do something about pairs? + +enumFromTo_double :: (Monad m, Ord a, RealFrac a) => a -> a -> Stream m a +{-# INLINE_FUSED enumFromTo_double #-} +enumFromTo_double n m = n `seq` m `seq` Stream step n + where + lim = m + 1/2 -- important to float out + + {-# INLINE_INNER step #-} + step x | x <= lim = return $ Yield x (x+1) + | otherwise = return $ Done + +{-# RULES + +"enumFromTo<Double> [Stream]" + enumFromTo = enumFromTo_double :: Monad m => Double -> Double -> Stream m Double + +"enumFromTo<Float> [Stream]" + enumFromTo = enumFromTo_double :: Monad m => Float -> Float -> Stream m Float #-} + + + +------------------------------------------------------------------------ + +-- | Enumerate values with a given step. +-- +-- /WARNING:/ This operation is very inefficient. If at all possible, use +-- 'enumFromStepN' instead. +enumFromThenTo :: (Enum a, Monad m) => a -> a -> a -> Stream m a +{-# INLINE_FUSED enumFromThenTo #-} +enumFromThenTo x y z = fromList [x, y .. z] + +-- FIXME: Specialise enumFromThenTo. + +-- Conversions +-- ----------- + +-- | Convert a 'Stream' to a list +toList :: Monad m => Stream m a -> m [a] +{-# INLINE toList #-} +toList = foldr (:) [] + +-- | Convert a list to a 'Stream' +fromList :: Monad m => [a] -> Stream m a +{-# INLINE fromList #-} +fromList zs = Stream step zs + where + step (x:xs) = return (Yield x xs) + step [] = return Done + +-- | Convert the first @n@ elements of a list to a 'Bundle' +fromListN :: Monad m => Int -> [a] -> Stream m a +{-# INLINE_FUSED fromListN #-} +fromListN m zs = Stream step (zs,m) + where + {-# INLINE_INNER step #-} + step (_, n) | n <= 0 = return Done + step (x:xs,n) = return (Yield x (xs,n-1)) + step ([],_) = return Done + +{- +fromVector :: (Monad m, Vector v a) => v a -> Stream m a +{-# INLINE_FUSED fromVector #-} +fromVector v = v `seq` n `seq` Stream (Unf step 0) + (Unf vstep True) + (Just v) + (Exact n) + where + n = basicLength v + + {-# INLINE step #-} + step i | i >= n = return Done + | otherwise = case basicUnsafeIndexM v i of + Box x -> return $ Yield x (i+1) + + + {-# INLINE vstep #-} + vstep True = return (Yield (Chunk (basicLength v) (\mv -> basicUnsafeCopy mv v)) False) + vstep False = return Done + +fromVectors :: forall m a. (Monad m, Vector v a) => [v a] -> Stream m a +{-# INLINE_FUSED fromVectors #-} +fromVectors vs = Stream (Unf pstep (Left vs)) + (Unf vstep vs) + Nothing + (Exact n) + where + n = List.foldl' (\k v -> k + basicLength v) 0 vs + + pstep (Left []) = return Done + pstep (Left (v:vs)) = basicLength v `seq` return (Skip (Right (v,0,vs))) + + pstep (Right (v,i,vs)) + | i >= basicLength v = return $ Skip (Left vs) + | otherwise = case basicUnsafeIndexM v i of + Box x -> return $ Yield x (Right (v,i+1,vs)) + + -- FIXME: work around bug in GHC 7.6.1 + vstep :: [v a] -> m (Step [v a] (Chunk v a)) + vstep [] = return Done + vstep (v:vs) = return $ Yield (Chunk (basicLength v) + (\mv -> INTERNAL_CHECK(check) "concatVectors" "length mismatch" + (M.basicLength mv == basicLength v) + $ basicUnsafeCopy mv v)) vs + + +concatVectors :: (Monad m, Vector v a) => Stream m (v a) -> Stream m a +{-# INLINE_FUSED concatVectors #-} +concatVectors (Stream step s} + = Stream (Unf pstep (Left s)) + (Unf vstep s) + Nothing + Unknown + where + pstep (Left s) = do + r <- step s + case r of + Yield v s' -> basicLength v `seq` return (Skip (Right (v,0,s'))) + Skip s' -> return (Skip (Left s')) + Done -> return Done + + pstep (Right (v,i,s)) + | i >= basicLength v = return (Skip (Left s)) + | otherwise = case basicUnsafeIndexM v i of + Box x -> return (Yield x (Right (v,i+1,s))) + + + vstep s = do + r <- step s + case r of + Yield v s' -> return (Yield (Chunk (basicLength v) + (\mv -> INTERNAL_CHECK(check) "concatVectors" "length mismatch" + (M.basicLength mv == basicLength v) + $ basicUnsafeCopy mv v)) s') + Skip s' -> return (Skip s') + Done -> return Done + +reVector :: Monad m => Stream m a -> Stream m a +{-# INLINE_FUSED reVector #-} +reVector (Stream step s, sSize = n} = Stream step s n + +{-# RULES + +"reVector [Vector]" + reVector = id + +"reVector/reVector [Vector]" forall s. + reVector (reVector s) = s #-} + + +-} + diff --git a/third_party/bazel/rules_haskell/examples/vector/Data/Vector/Fusion/Util.hs b/third_party/bazel/rules_haskell/examples/vector/Data/Vector/Fusion/Util.hs new file mode 100644 index 0000000000..855bf5ddd4 --- /dev/null +++ b/third_party/bazel/rules_haskell/examples/vector/Data/Vector/Fusion/Util.hs @@ -0,0 +1,60 @@ +{-# LANGUAGE CPP #-} +-- | +-- Module : Data.Vector.Fusion.Util +-- Copyright : (c) Roman Leshchinskiy 2009 +-- License : BSD-style +-- +-- Maintainer : Roman Leshchinskiy <rl@cse.unsw.edu.au> +-- Stability : experimental +-- Portability : portable +-- +-- Fusion-related utility types +-- + +module Data.Vector.Fusion.Util ( + Id(..), Box(..), + + delay_inline, delayed_min +) where + +#if !MIN_VERSION_base(4,8,0) +import Control.Applicative (Applicative(..)) +#endif + +-- | Identity monad +newtype Id a = Id { unId :: a } + +instance Functor Id where + fmap f (Id x) = Id (f x) + +instance Applicative Id where + pure = Id + Id f <*> Id x = Id (f x) + +instance Monad Id where + return = pure + Id x >>= f = f x + +-- | Box monad +data Box a = Box { unBox :: a } + +instance Functor Box where + fmap f (Box x) = Box (f x) + +instance Applicative Box where + pure = Box + Box f <*> Box x = Box (f x) + +instance Monad Box where + return = pure + Box x >>= f = f x + +-- | Delay inlining a function until late in the game (simplifier phase 0). +delay_inline :: (a -> b) -> a -> b +{-# INLINE [0] delay_inline #-} +delay_inline f = f + +-- | `min` inlined in phase 0 +delayed_min :: Int -> Int -> Int +{-# INLINE [0] delayed_min #-} +delayed_min m n = min m n |