path: root/third_party/bazel/rules_haskell/examples/vector/Data/Vector/Fusion/Bundle.hs

{-# 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