mirror of
https://github.com/composewell/streamly.git
synced 2024-11-10 12:47:22 +03:00
358 lines
11 KiB
Haskell
358 lines
11 KiB
Haskell
-- |
|
|
-- Module : StreamDOps
|
|
-- Copyright : (c) 2018 Harendra Kumar
|
|
--
|
|
-- License : BSD3
|
|
-- Maintainer : streamly@composewell.com
|
|
|
|
{-# LANGUAGE FlexibleContexts #-}
|
|
{-# LANGUAGE ScopedTypeVariables #-}
|
|
|
|
module StreamDOps where
|
|
|
|
import Control.Monad (when)
|
|
import Data.Maybe (isJust)
|
|
import Prelude
|
|
(Monad, Int, (+), ($), (.), return, (>), even, (<=), div,
|
|
subtract, undefined, Maybe(..), not, (>>=),
|
|
maxBound, fmap, odd, (==), flip, (<$>), (<*>), round, (/), (**), (<))
|
|
import qualified Prelude as P
|
|
|
|
import qualified Streamly.Internal.Data.Stream.StreamD as S
|
|
import qualified Streamly.Internal.Data.Unfold as UF
|
|
|
|
-- We try to keep the total number of iterations same irrespective of nesting
|
|
-- of the loops so that the overhead is easy to compare.
|
|
value, value2, value3, value16, maxValue :: Int
|
|
value = 100000
|
|
value2 = round (P.fromIntegral value**(1/2::P.Double)) -- double nested loop
|
|
value3 = round (P.fromIntegral value**(1/3::P.Double)) -- triple nested loop
|
|
value16 = round (P.fromIntegral value**(1/16::P.Double)) -- triple nested loop
|
|
maxValue = value
|
|
|
|
-------------------------------------------------------------------------------
|
|
-- Stream generation and elimination
|
|
-------------------------------------------------------------------------------
|
|
|
|
type Stream m a = S.Stream m a
|
|
|
|
{-# INLINE sourceUnfoldr #-}
|
|
sourceUnfoldr :: Monad m => Int -> Stream m Int
|
|
sourceUnfoldr n = S.unfoldr step n
|
|
where
|
|
step cnt =
|
|
if cnt > n + value
|
|
then Nothing
|
|
else Just (cnt, cnt + 1)
|
|
|
|
{-# INLINE sourceUnfoldrN #-}
|
|
sourceUnfoldrN :: Monad m => Int -> Int -> Stream m Int
|
|
sourceUnfoldrN m n = S.unfoldr step n
|
|
where
|
|
step cnt =
|
|
if cnt > n + m
|
|
then Nothing
|
|
else Just (cnt, cnt + 1)
|
|
|
|
{-# INLINE sourceUnfoldrMN #-}
|
|
sourceUnfoldrMN :: Monad m => Int -> Int -> Stream m Int
|
|
sourceUnfoldrMN m n = S.unfoldrM step n
|
|
where
|
|
step cnt =
|
|
if cnt > n + m
|
|
then return Nothing
|
|
else return (Just (cnt, cnt + 1))
|
|
|
|
{-# INLINE sourceUnfoldrM #-}
|
|
sourceUnfoldrM :: Monad m => Int -> Stream m Int
|
|
sourceUnfoldrM n = S.unfoldrM step n
|
|
where
|
|
step cnt =
|
|
if cnt > n + value
|
|
then return Nothing
|
|
else return (Just (cnt, cnt + 1))
|
|
|
|
{-# INLINE sourceIntFromTo #-}
|
|
sourceIntFromTo :: Monad m => Int -> Stream m Int
|
|
sourceIntFromTo n = S.enumerateFromToIntegral n (n + value)
|
|
|
|
{-# INLINE sourceFromList #-}
|
|
sourceFromList :: Monad m => Int -> Stream m Int
|
|
sourceFromList n = S.fromList [n..n+value]
|
|
|
|
{-# INLINE source #-}
|
|
source :: Monad m => Int -> Stream m Int
|
|
source = sourceUnfoldrM
|
|
|
|
-------------------------------------------------------------------------------
|
|
-- Elimination
|
|
-------------------------------------------------------------------------------
|
|
|
|
{-# INLINE runStream #-}
|
|
runStream :: Monad m => Stream m a -> m ()
|
|
runStream = S.drain
|
|
|
|
{-# INLINE mapM_ #-}
|
|
mapM_ :: Monad m => Stream m a -> m ()
|
|
mapM_ = S.mapM_ (\_ -> return ())
|
|
|
|
{-# INLINE toNull #-}
|
|
toNull :: Monad m => Stream m Int -> m ()
|
|
toNull = runStream
|
|
|
|
{-# INLINE uncons #-}
|
|
{-# INLINE nullTail #-}
|
|
{-# INLINE headTail #-}
|
|
uncons, nullTail, headTail
|
|
:: Monad m
|
|
=> Stream m Int -> m ()
|
|
|
|
uncons s = do
|
|
r <- S.uncons s
|
|
case r of
|
|
Nothing -> return ()
|
|
Just (_, t) -> uncons t
|
|
|
|
{-# INLINE tail #-}
|
|
tail :: Monad m => Stream m a -> m ()
|
|
tail s = S.tail s >>= P.mapM_ tail
|
|
|
|
nullTail s = do
|
|
r <- S.null s
|
|
when (not r) $ S.tail s >>= P.mapM_ nullTail
|
|
|
|
headTail s = do
|
|
h <- S.head s
|
|
when (isJust h) $ S.tail s >>= P.mapM_ headTail
|
|
|
|
{-# INLINE toList #-}
|
|
toList :: Monad m => Stream m Int -> m [Int]
|
|
toList = S.toList
|
|
|
|
{-# INLINE foldl #-}
|
|
foldl :: Monad m => Stream m Int -> m Int
|
|
foldl = S.foldl' (+) 0
|
|
|
|
{-# INLINE last #-}
|
|
last :: Monad m => Stream m Int -> m (Maybe Int)
|
|
last = S.last
|
|
|
|
-------------------------------------------------------------------------------
|
|
-- Transformation
|
|
-------------------------------------------------------------------------------
|
|
|
|
{-# INLINE transform #-}
|
|
transform :: Monad m => Stream m a -> m ()
|
|
transform = runStream
|
|
|
|
{-# INLINE composeN #-}
|
|
composeN
|
|
:: Monad m
|
|
=> Int -> (Stream m Int -> Stream m Int) -> Stream m Int -> m ()
|
|
composeN n f =
|
|
case n of
|
|
1 -> transform . f
|
|
2 -> transform . f . f
|
|
3 -> transform . f . f . f
|
|
4 -> transform . f . f . f . f
|
|
_ -> undefined
|
|
|
|
{-# INLINE scan #-}
|
|
{-# INLINE map #-}
|
|
{-# INLINE fmap #-}
|
|
{-# INLINE mapM #-}
|
|
{-# INLINE mapMaybe #-}
|
|
{-# INLINE mapMaybeM #-}
|
|
{-# INLINE filterEven #-}
|
|
{-# INLINE filterAllOut #-}
|
|
{-# INLINE filterAllIn #-}
|
|
{-# INLINE takeOne #-}
|
|
{-# INLINE takeAll #-}
|
|
{-# INLINE takeWhileTrue #-}
|
|
{-# INLINE takeWhileMTrue #-}
|
|
{-# INLINE dropOne #-}
|
|
{-# INLINE dropAll #-}
|
|
{-# INLINE dropWhileTrue #-}
|
|
{-# INLINE dropWhileMTrue #-}
|
|
{-# INLINE dropWhileFalse #-}
|
|
{-# INLINE foldrS #-}
|
|
{-# INLINE foldlS #-}
|
|
{-# INLINE concatMap #-}
|
|
{-# INLINE intersperse #-}
|
|
scan, map, fmap, mapM, mapMaybe, mapMaybeM, filterEven, filterAllOut,
|
|
filterAllIn, takeOne, takeAll, takeWhileTrue, takeWhileMTrue, dropOne,
|
|
dropAll, dropWhileTrue, dropWhileMTrue, dropWhileFalse, foldrS, foldlS,
|
|
concatMap, intersperse
|
|
:: Monad m
|
|
=> Int -> Stream m Int -> m ()
|
|
|
|
scan n = composeN n $ S.scanl' (+) 0
|
|
fmap n = composeN n $ Prelude.fmap (+1)
|
|
map n = composeN n $ S.map (+1)
|
|
mapM n = composeN n $ S.mapM return
|
|
mapMaybe n = composeN n $ S.mapMaybe
|
|
(\x -> if Prelude.odd x then Nothing else Just x)
|
|
mapMaybeM n = composeN n $ S.mapMaybeM
|
|
(\x -> if Prelude.odd x then return Nothing else return $ Just x)
|
|
filterEven n = composeN n $ S.filter even
|
|
filterAllOut n = composeN n $ S.filter (> maxValue)
|
|
filterAllIn n = composeN n $ S.filter (<= maxValue)
|
|
takeOne n = composeN n $ S.take 1
|
|
takeAll n = composeN n $ S.take maxValue
|
|
takeWhileTrue n = composeN n $ S.takeWhile (<= maxValue)
|
|
takeWhileMTrue n = composeN n $ S.takeWhileM (return . (<= maxValue))
|
|
dropOne n = composeN n $ S.drop 1
|
|
dropAll n = composeN n $ S.drop maxValue
|
|
dropWhileTrue n = composeN n $ S.dropWhile (<= maxValue)
|
|
dropWhileMTrue n = composeN n $ S.dropWhileM (return . (<= maxValue))
|
|
dropWhileFalse n = composeN n $ S.dropWhile (> maxValue)
|
|
foldrS n = composeN n $ S.foldrS S.cons S.nil
|
|
foldlS n = composeN n $ S.foldlS (flip S.cons) S.nil
|
|
concatMap n = composeN n $ (\s -> S.concatMap (\_ -> s) s)
|
|
intersperse n = composeN n $ S.intersperse maxValue
|
|
|
|
-------------------------------------------------------------------------------
|
|
-- Iteration
|
|
-------------------------------------------------------------------------------
|
|
|
|
iterStreamLen, maxIters :: Int
|
|
iterStreamLen = 10
|
|
maxIters = 10000
|
|
|
|
{-# INLINE iterateSource #-}
|
|
iterateSource
|
|
:: Monad m
|
|
=> (Stream m Int -> Stream m Int) -> Int -> Int -> Stream m Int
|
|
iterateSource g i n = f i (sourceUnfoldrMN iterStreamLen n)
|
|
where
|
|
f (0 :: Int) m = g m
|
|
f x m = g (f (x P.- 1) m)
|
|
|
|
{-# INLINE iterateMapM #-}
|
|
{-# INLINE iterateScan #-}
|
|
{-# INLINE iterateFilterEven #-}
|
|
{-# INLINE iterateTakeAll #-}
|
|
{-# INLINE iterateDropOne #-}
|
|
{-# INLINE iterateDropWhileFalse #-}
|
|
{-# INLINE iterateDropWhileTrue #-}
|
|
iterateMapM, iterateScan, iterateFilterEven, iterateTakeAll, iterateDropOne,
|
|
iterateDropWhileFalse, iterateDropWhileTrue
|
|
:: Monad m
|
|
=> Int -> Stream m Int
|
|
|
|
-- this is quadratic
|
|
iterateScan = iterateSource (S.scanl' (+) 0) (maxIters `div` 10)
|
|
iterateDropWhileFalse = iterateSource (S.dropWhile (> maxValue))
|
|
(maxIters `div` 10)
|
|
|
|
iterateMapM = iterateSource (S.mapM return) maxIters
|
|
iterateFilterEven = iterateSource (S.filter even) maxIters
|
|
iterateTakeAll = iterateSource (S.take maxValue) maxIters
|
|
iterateDropOne = iterateSource (S.drop 1) maxIters
|
|
iterateDropWhileTrue = iterateSource (S.dropWhile (<= maxValue)) maxIters
|
|
|
|
{-# INLINE iterateM #-}
|
|
iterateM :: Monad m => Int -> Stream m Int
|
|
iterateM i = S.take maxIters (S.iterateM (\x -> return (x + 1)) (return i))
|
|
|
|
-------------------------------------------------------------------------------
|
|
-- Zipping and concat
|
|
-------------------------------------------------------------------------------
|
|
|
|
{-# INLINE eqBy #-}
|
|
eqBy :: (Monad m, P.Eq a) => S.Stream m a -> m P.Bool
|
|
eqBy src = S.eqBy (==) src src
|
|
|
|
{-# INLINE cmpBy #-}
|
|
cmpBy :: (Monad m, P.Ord a) => S.Stream m a -> m P.Ordering
|
|
cmpBy src = S.cmpBy P.compare src src
|
|
|
|
{-# INLINE zip #-}
|
|
zip :: Monad m => Stream m Int -> m ()
|
|
zip src = transform $ S.zipWith (,) src src
|
|
|
|
{-# INLINE concatMapRepl4xN #-}
|
|
concatMapRepl4xN :: Monad m => Stream m Int -> m ()
|
|
concatMapRepl4xN src = transform $ (S.concatMap (S.replicate 4) src)
|
|
|
|
{-# INLINE concatMapURepl4xN #-}
|
|
concatMapURepl4xN :: Monad m => Stream m Int -> m ()
|
|
concatMapURepl4xN src = transform $ S.concatMapU (UF.replicateM 4) src
|
|
|
|
-------------------------------------------------------------------------------
|
|
-- Mixed Composition
|
|
-------------------------------------------------------------------------------
|
|
|
|
{-# INLINE scanMap #-}
|
|
{-# INLINE dropMap #-}
|
|
{-# INLINE dropScan #-}
|
|
{-# INLINE takeDrop #-}
|
|
{-# INLINE takeScan #-}
|
|
{-# INLINE takeMap #-}
|
|
{-# INLINE filterDrop #-}
|
|
{-# INLINE filterTake #-}
|
|
{-# INLINE filterScan #-}
|
|
{-# INLINE filterMap #-}
|
|
scanMap, dropMap, dropScan, takeDrop, takeScan, takeMap, filterDrop,
|
|
filterTake, filterScan, filterMap
|
|
:: Monad m => Int -> Stream m Int -> m ()
|
|
|
|
scanMap n = composeN n $ S.map (subtract 1) . S.scanl' (+) 0
|
|
dropMap n = composeN n $ S.map (subtract 1) . S.drop 1
|
|
dropScan n = composeN n $ S.scanl' (+) 0 . S.drop 1
|
|
takeDrop n = composeN n $ S.drop 1 . S.take maxValue
|
|
takeScan n = composeN n $ S.scanl' (+) 0 . S.take maxValue
|
|
takeMap n = composeN n $ S.map (subtract 1) . S.take maxValue
|
|
filterDrop n = composeN n $ S.drop 1 . S.filter (<= maxValue)
|
|
filterTake n = composeN n $ S.take maxValue . S.filter (<= maxValue)
|
|
filterScan n = composeN n $ S.scanl' (+) 0 . S.filter (<= maxBound)
|
|
filterMap n = composeN n $ S.map (subtract 1) . S.filter (<= maxValue)
|
|
|
|
-------------------------------------------------------------------------------
|
|
-- Nested Composition
|
|
-------------------------------------------------------------------------------
|
|
|
|
{-# INLINE toNullApNested #-}
|
|
toNullApNested :: Monad m => Stream m Int -> m ()
|
|
toNullApNested s = runStream $ do
|
|
(+) <$> s <*> s
|
|
|
|
{-# INLINE toNullNested #-}
|
|
toNullNested :: Monad m => Stream m Int -> m ()
|
|
toNullNested s = runStream $ do
|
|
x <- s
|
|
y <- s
|
|
return $ x + y
|
|
|
|
{-# INLINE toNullNested3 #-}
|
|
toNullNested3 :: Monad m => Stream m Int -> m ()
|
|
toNullNested3 s = runStream $ do
|
|
x <- s
|
|
y <- s
|
|
z <- s
|
|
return $ x + y + z
|
|
|
|
{-# INLINE filterAllOutNested #-}
|
|
filterAllOutNested
|
|
:: Monad m
|
|
=> Stream m Int -> m ()
|
|
filterAllOutNested str = runStream $ do
|
|
x <- str
|
|
y <- str
|
|
let s = x + y
|
|
if s < 0
|
|
then return s
|
|
else S.nil
|
|
|
|
{-# INLINE filterAllInNested #-}
|
|
filterAllInNested
|
|
:: Monad m
|
|
=> Stream m Int -> m ()
|
|
filterAllInNested str = runStream $ do
|
|
x <- str
|
|
y <- str
|
|
let s = x + y
|
|
if s > 0
|
|
then return s
|
|
else S.nil
|