diff options
author | Vincent Ambo <tazjin@google.com> | 2019-07-04T10·18+0100 |
---|---|---|
committer | Vincent Ambo <tazjin@google.com> | 2019-07-04T10·18+0100 |
commit | f723b8b878a3c4a4687b9e337a875500bebb39b1 (patch) | |
tree | e85204cf042c355e90cff61c111e7d8cd15df311 /third_party/bazel/rules_haskell/examples/vector/Data/Vector/Fusion/Bundle.hs | |
parent | 2eb1dc26e42ffbdc168f05ef744bd4b4f3e4c36f (diff) |
feat(third_party/bazel): Check in rules_haskell from Tweag r/17
Diffstat (limited to 'third_party/bazel/rules_haskell/examples/vector/Data/Vector/Fusion/Bundle.hs')
-rw-r--r-- | third_party/bazel/rules_haskell/examples/vector/Data/Vector/Fusion/Bundle.hs | 655 |
1 files changed, 655 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 000000000000..6b6b6236d7cb --- /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 + |