leveldb-haskell-0.5.1: src/Data/Stream/Monadic.hs
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ExistentialQuantification #-}
-- |
-- Module : Data.Stream.Monadic
-- Copyright : (c) 2014 Kim Altintop
-- License : BSD3
-- Maintainer : kim.altintop@gmail.com
-- Stability : experimental
-- Portability : non-portable
--
-- (Mostly mechanical) adaptation of the
-- <http://hackage.haskell.org/package/stream-fusion/docs/Data-Stream.html Data.Stream>
-- module from the
-- <http://hackage.haskell.org/package/stream-fusion stream-fusion> package to a
-- monadic 'Stream' datatype similar to the one
-- <https://www.fpcomplete.com/blog/2014/08/conduit-stream-fusion proposed> by
-- Michael Snoyman for the <http://hackage.haskell.org/package/conduit conduit>
-- package.
--
-- The intention here is to provide a high-level, "Data.List"-like interface to
-- "Database.LevelDB.Iterator"s with predictable space and time complexity (see
-- "Database.LevelDB.Streaming"), and without introducing a dependency eg. on
-- one of the streaming libraries (all relevant datatypes are fully exported,
-- though, so it should be straightforward to write wrappers for your favourite
-- streaming library).
--
-- Fusion and inlining rules and strictness annotations have been put in place
-- faithfully, and may need further profiling. Also, some functions (from
-- "Data.List") have been omitted as either no obvious solution exists (notably
-- @mapM@), they didn't seem too useful in the given context (eg. @lookup@), or
-- I was just too lazy. Missing functions may be added upon
-- <https://github.com/kim/leveldb-haskell/pulls request>.
module Data.Stream.Monadic
( Step (..)
, Stream (..)
-- * Conversion with lists
, toList
, fromList
-- * Basic functions
, append
, cons
, snoc
, head
, last
, tail
, init
, null
, length
-- * Transformations
, map
-- , mapM
, intersperse
-- * Folds
, foldl
, foldl'
-- , foldl1
-- , foldl1'
, foldr
-- , foldr1
, foldMap
, foldM
, foldM_
-- * Special folds
-- , concat
, concatMap
-- , and
-- , or
-- , any
-- , all
-- , sum
-- , product
-- , maximum
-- , minimum
-- , scanl
-- , scanl1
-- * Infinite streams
, iterate
, repeat
, replicate
, cycle
-- * Unfolding
, unfoldr
, unfoldrM
-- , isPrefixOf
-- * Searching streams
-- , elem
-- , lookup
-- , find
, filter
-- , index
-- , findIndex
-- , elemIndex
-- , elemIndices
-- , findIndices
-- * Substreams
, take
, drop
-- , splitAt
, takeWhile
, dropWhile
-- * Zipping and unzipping
, zip
-- , zip3
-- , zip4
, zipWith
-- , zipWith3
-- , zipWith4
, unzip
-- , insertBy
-- , maximumBy
-- , minimumBy
-- , genericLength
-- , genericTake
-- , genericDrop
-- , genericIndex
-- , genericSplitAt
-- , enumFromToInt
-- , enumFromToChar
-- , enumDeltaInteger
)
where
import Control.Applicative
import Data.Monoid
import Prelude (Bool (..), Either (..), Eq (..), Functor (..), Int, Maybe (..),
Monad (..), Num (..), Ord (..), error, otherwise, ($), (&&),
(.), (=<<))
data Step a s
= Yield a !s
| Skip !s
| Done
data Stream m a = forall s. Stream (s -> m (Step a s)) (m s)
instance Monad m => Functor (Stream m) where
fmap = map
toList :: (Functor m, Monad m) => Stream m a -> m [a]
toList (Stream next s0) = unfold =<< s0
where
unfold !s = do
step <- next s
case step of
Done -> return []
Skip s' -> unfold s'
Yield x s' -> (x :) <$> unfold s'
fromList :: Monad m => [a] -> Stream m a
fromList xs = Stream next (return xs)
where
{-# INLINE next #-}
next [] = return Done
next (x:xs') = return $ Yield x xs'
{-# RULES
"Stream fromList/toList fusion" forall s.
fmap fromList (toList s) = return s
#-}
append :: (Functor m, Monad m) => Stream m a -> Stream m a -> Stream m a
append (Stream next0 s0) (Stream next1 s1) = Stream next (Left <$> s0)
where
{-# INLINE next #-}
next (Left s) = do
step <- next0 s
case step of
Done -> Skip . Right <$> s1
Skip s' -> return $ Skip (Left s')
Yield x s' -> return $ Yield x (Left s')
next (Right s) = do
step <- next1 s
return $ case step of
Done -> Done
Skip s' -> Skip (Right s')
Yield x s' -> Yield x (Right s')
{-# INLINE [0] append #-}
cons :: (Functor m, Monad m) => a -> Stream m a -> Stream m a
cons w (Stream next0 s0) = Stream next ((,) S2 <$> s0)
where
{-# INLINE next #-}
next (S2, s) = return $ Yield w (S1, s)
next (S1, s) = do
step <- next0 s
return $ case step of
Done -> Done
Skip s' -> Skip (S1, s')
Yield x s' -> Yield x (S1, s')
{-# INLINE [0] cons #-}
snoc :: (Functor m, Monad m) => Stream m a -> a -> Stream m a
snoc (Stream next0 s0) y = Stream next (Just <$> s0)
where
{-# INLINE next #-}
next Nothing = return Done
next (Just s) = do
step <- next0 s
return $ case step of
Done -> Yield y Nothing
Skip s' -> Skip (Just s')
Yield x s' -> Yield x (Just s')
{-# INLINE [0] snoc #-}
-- | Unlike 'Data.List.head', this function does not diverge if the 'Stream' is
-- empty. Instead, 'Nothing' is returned.
head :: Monad m => Stream m a -> m (Maybe a)
head (Stream next s0) = loop =<< s0
where
loop !s = do
step <- next s
case step of
Yield x _ -> return $ Just x
Skip s' -> loop s'
Done -> return Nothing
{-# INLINE [0] head #-}
-- | Unlike 'Data.List.last', this function does not diverge if the 'Stream' is
-- empty. Instead, 'Nothing' is returned.
last :: Monad m => Stream m a -> m (Maybe a)
last (Stream next s0) = loop =<< s0
where
loop !s = do
step <- next s
case step of
Done -> return Nothing
Skip s' -> loop s'
Yield x s' -> loop' x s'
loop' x !s = do
step <- next s
case step of
Done -> return $ Just x
Skip s' -> loop' x s'
Yield x' s' -> loop' x' s'
{-# INLINE [0] last #-}
data Switch = S1 | S2
-- | Unlike 'Data.List.tail', this function does not diverge if the 'Stream' is
-- empty. Instead, it is the identity in this case.
tail :: (Functor m, Monad m) => Stream m a -> Stream m a
tail (Stream next0 s0) = Stream next ((,) S1 <$> s0)
where
{-# INLINE next #-}
next (S1, s) = do
step <- next0 s
return $ case step of
Done -> Done
Skip s' -> Skip (S1, s')
Yield _ s' -> Skip (S2, s')
next (S2, s) = do
step <- next0 s
return $ case step of
Done -> Done
Skip s' -> Skip (S2, s')
Yield x s' -> Yield x (S2, s')
{-# INLINE [0] tail #-}
-- | Unlike 'Data.List.init', this function does not diverge if the 'Stream' is
-- empty. Instead, it is the identity in this case.
init :: (Functor m, Monad m) => Stream m a -> Stream m a
init (Stream next0 s0) = Stream next ((,) Nothing <$> s0)
where
{-# INLINE next #-}
next (Nothing, s) = do
step <- next0 s
return $ case step of
Done -> Done
Skip s' -> Skip (Nothing, s')
Yield x s' -> Skip (Just x , s')
next (Just x, s) = do
step <- next0 s
return $ case step of
Done -> Done
Skip s' -> Skip (Just x , s')
Yield x' s' -> Yield x (Just x', s')
{-# INLINE [0] init #-}
null :: Monad m => Stream m a -> m Bool
null (Stream next s0) = loop =<< s0
where
loop !s = do
step <- next s
case step of
Done -> return True
Yield _ _ -> return False
Skip s' -> loop s'
{-# INLINE [0] null #-}
length :: Monad m => Stream m a -> m Int
length (Stream next s0) = loop 0 =<< s0
where
loop !z !s = do
step <- next s
case step of
Done -> return z
Skip s' -> loop z s'
Yield _ s' -> loop (z+1) s'
{-# INLINE [0] length #-}
filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
filter p (Stream next0 s0) = Stream next s0
where
{-# INLINE next #-}
next !s = do
step <- next0 s
return $ case step of
Done -> Done
Skip s' -> Skip s'
Yield x s' | p x -> Yield x s'
| otherwise -> Skip s'
{-# INLINE [0] filter #-}
{-# RULES
"Stream filter/filter fusion" forall p q s.
filter p (filter q s) = filter (\x -> q x && p x) s
#-}
map :: Monad m => (a -> b) -> Stream m a -> Stream m b
map f (Stream next0 s0) = Stream next s0
where
{-# INLINE next #-}
next !s = do
step <- next0 s
return $ case step of
Done -> Done
Skip s' -> Skip s'
Yield x s' -> Yield (f x) s'
{-# INLINE [0] map #-}
{-# RULES
"Stream map/map fusion" forall f g s.
map f (map g s) = map (\x -> f (g x)) s
#-}
-- 'mapM' is tricky:
--
-- > mapM :: (Monad m, Monad n) => (a -> n b) -> Stream m a -> n (Stream n b)
--
-- we would need a constraint which specifies how to lift any monad /m/ into
-- some monad /n/ (or specialise /m/ to 'IO').
--
-- alternatively, we may define:
--
-- > mapM :: Monad m => (a -> m b) -> Stream m a -> m (Stream m b)
--
-- or rather:
--
-- > mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
--
-- not sure how useful this would be.
intersperse :: (Functor m, Monad m) => a -> Stream m a -> Stream m a
intersperse sep (Stream next0 s0) = Stream next ((,,) Nothing S1 <$> s0)
where
{-# INLINE next #-}
next (Nothing, S1, s) = do
step <- next0 s
return $ case step of
Done -> Done
Skip s' -> Skip (Nothing, S1, s')
Yield x s' -> Skip (Just x , S1, s')
next (Just x, S1, s) = return $ Yield x (Nothing, S2, s)
next (Nothing, S2, s) = do
step <- next0 s
return $ case step of
Done -> Done
Skip s' -> Skip (Nothing, S2, s')
Yield x s' -> Yield sep (Just x , S1, s')
next (Just _, S2, _) = error "Data.Stream.Monadic.intersperse: impossible"
{-# INLINE [0] intersperse #-}
foldMap :: (Monoid m, Functor n, Monad n) => (a -> m) -> Stream n a -> n m
foldMap f (Stream next s0) = loop mempty =<< s0
where
loop z !s = do
step <- next s
case step of
Done -> return z
Skip s' -> loop z s'
Yield x s' -> loop (z <> f x) s'
{-# INLINE [0] foldMap #-}
-- | Left-associative fold.
--
-- Note that the /direction/ of the traversal is not defined here.
foldl :: Monad m => (b -> a -> b) -> b -> Stream m a -> m b
foldl f z0 (Stream next s0) = loop z0 =<< s0
where
loop z !s = do
step <- next s
case step of
Done -> return z
Skip s' -> loop z s'
Yield x s' -> loop (f z x) s'
{-# INLINE [0] foldl #-}
-- | Left-associative fold with strict accumulator.
--
-- Note that the /direction/ of the traversal is not defined here.
foldl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> m b
foldl' f z0 (Stream next s0) = loop z0 =<< s0
where
loop !z !s = do
step <- next s
case step of
Done -> return z
Skip s' -> loop z s'
Yield x s' -> loop (f z x) s'
{-# INLINE [0] foldl' #-}
-- | Right-associative fold.
--
-- Note that the /direction/ of the traversal is not defined here.
foldr :: (Functor m, Monad m) => (a -> b -> b) -> b -> Stream m a -> m b
foldr f z (Stream next s0) = loop =<< s0
where
loop !s = do
step <- next s
case step of
Done -> return z
Skip s' -> loop s'
Yield x s' -> f x <$> loop s'
{-# INLINE [0] foldr #-}
foldM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> m b
foldM f z0 (Stream next s0) = loop z0 =<< s0
where
loop z !s = do
step <- next s
case step of
Done -> return z
Skip s' -> loop z s'
Yield x s' -> f z x >>= (`loop` s')
{-# INLINE [0] foldM #-}
foldM_ :: Monad m => (b -> a -> m b) -> b -> Stream m a -> m ()
foldM_ f z0 (Stream next s0) = loop z0 =<< s0
where
loop z !s = do
step <- next s
case step of
Done -> return ()
Skip s' -> loop z s'
Yield x s' -> f z x >>= (`loop` s')
{-# INLINE [0] foldM_ #-}
concatMap :: (Functor m, Monad m) => (a -> Stream m b) -> Stream m a -> Stream m b
concatMap f (Stream next0 s0) = Stream next ((,) Nothing <$> s0)
where
{-# INLINE next #-}
next (Nothing, s) = do
step <- next0 s
return $ case step of
Done -> Done
Skip s' -> Skip (Nothing , s')
Yield x s' -> Skip (Just (f x), s')
next (Just (Stream g t), s) = do
step <- g =<< t
return $ case step of
Done -> Skip (Nothing , s)
Skip t' -> Skip (Just (Stream g (return t')), s)
Yield x t' -> Yield x (Just (Stream g (return t')), s)
{-# INLINE [0] concatMap #-}
iterate :: Monad m => (a -> a) -> a -> Stream m a
iterate f x0 = Stream next (return x0)
where
{-# INLINE next #-}
next x = return $ Yield x (f x)
{-# INLINE [0] iterate #-}
repeat :: Monad m => a -> Stream m a
repeat x = Stream next (return ())
where
{-# INLINE next #-}
next _ = return $ Yield x ()
{-# INLINE [0] repeat #-}
{-# RULES
"map/repeat" forall f x. map f (repeat x) = repeat (f x)
#-}
replicate :: Monad m => Int -> a -> Stream m a
replicate n x = Stream next (return n)
where
{-# INLINE next #-}
next !i | i <= 0 = return Done
| otherwise = return $ Yield x (i-1)
{-# INLINE [0] replicate #-}
{-# RULES
"map/replicate" forall f n x. map f (replicate n x) = replicate n (f x)
#-}
-- | Unlike 'Data.List.cycle', this function does not diverge if the 'Stream' is
-- empty. Instead, it is the identity in this case.
cycle :: (Functor m, Monad m) => Stream m a -> Stream m a
cycle (Stream next0 s0) = Stream next ((,) S1 <$> s0)
where
{-# INLINE next #-}
next (S1, s) = do
step <- next0 s
return $ case step of
Done -> Done -- error?
Skip s' -> Skip (S1, s')
Yield x s' -> Yield x (S2, s')
next (S2, s) = do
step <- next0 s
case step of
Done -> Skip . ((,) S2) <$> s0
Skip s' -> return $ Skip (S2, s')
Yield x s' -> return $ Yield x (S2, s')
{-# INLINE [0] cycle #-}
unfoldr :: Monad m => (b -> Maybe (a, b)) -> b -> Stream m a
unfoldr f s0 = Stream next (return s0)
where
{-# INLINE next #-}
next s = return $ case f s of
Nothing -> Done
Just (w, s') -> Yield w s'
{-# INLINE [0] unfoldr #-}
-- | Build a stream from a monadic seed (or state function).
unfoldrM :: (Functor m, Monad m) => (b -> Maybe (a, m b)) -> m b -> Stream m a
unfoldrM f s0 = Stream next s0
where
{-# INLINE next #-}
next s = case f s of
Nothing -> return Done
Just (w, s') -> Yield w <$> s'
{-# INLINE [0] unfoldrM #-}
take :: (Functor m, Monad m) => Int -> Stream m a -> Stream m a
take n0 (Stream next0 s0) = Stream next ((,) n0 <$> s0)
where
{-# INLINE next #-}
next (!n, s)
| n <= 0 = return Done
| otherwise = do
step <- next0 s
return $ case step of
Done -> Done
Skip s' -> Skip (n , s')
Yield x s' -> Yield x (n-1, s')
{-# INLINE [0] take #-}
drop :: (Functor m, Monad m) => Int -> Stream m a -> Stream m a
drop n0 (Stream next0 s0) = Stream next ((,) (Just (max 0 n0)) <$> s0)
where
{-# INLINE next #-}
next (Just !n, s)
| n == 0 = return $ Skip (Nothing, s)
| otherwise = do
step <- next0 s
return $ case step of
Done -> Done
Skip s' -> Skip (Just n , s')
Yield _ s' -> Skip (Just (n-1), s')
next (Nothing, s) = do
step <- next0 s
return $ case step of
Done -> Done
Skip s' -> Skip (Nothing, s')
Yield x s' -> Yield x (Nothing, s')
{-# INLINE [0] drop #-}
takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
takeWhile p (Stream next0 s0) = Stream next s0
where
{-# INLINE next #-}
next !s = do
step <- next0 s
return $ case step of
Done -> Done
Skip s' -> Skip s'
Yield x s' | p x -> Yield x s'
| otherwise -> Done
{-# INLINE [0] takeWhile #-}
dropWhile :: (Functor m, Monad m) => (a -> Bool) -> Stream m a -> Stream m a
dropWhile p (Stream next0 s0) = Stream next ((,) S1 <$> s0)
where
{-# INLINE next #-}
next (S1, s) = do
step <- next0 s
return $ case step of
Done -> Done
Skip s' -> Skip (S1, s')
Yield x s' | p x -> Skip (S1, s')
| otherwise -> Yield x (S2, s')
next (S2, s) = do
step <- next0 s
return $ case step of
Done -> Done
Skip s' -> Skip (S2, s')
Yield x s' -> Yield x (S2, s')
{-# INLINE [0] dropWhile #-}
zip :: (Functor m, Applicative m, Monad m)
=> Stream m a
-> Stream m b
-> Stream m (a, b)
zip = zipWith (,)
{-# INLINE zip #-}
zipWith :: (Functor m, Applicative m, Monad m)
=> (a -> b -> c)
-> Stream m a
-> Stream m b
-> Stream m c
zipWith f (Stream nexta sa0) (Stream nextb sb0) =
Stream next ((,,) Nothing <$> sa0 <*> sb0)
where
{-# INLINE next #-}
next (Nothing, sa, sb) = do
step <- nexta sa
return $ case step of
Done -> Done
Skip sa' -> Skip (Nothing, sa', sb)
Yield a sa' -> Skip (Just a , sa', sb)
next (Just a, sa', sb) = do
step <- nextb sb
return $ case step of
Done -> Done
Skip sb' -> Skip (Just a, sa', sb')
Yield b sb' -> Yield (f a b) (Nothing, sa', sb')
{-# INLINE [0] zipWith #-}
unzip :: (Functor m, Monad m) => Stream m (a, b) -> m ([a], [b])
unzip = foldr (\(a,b) ~(as, bs) -> (a:as, b:bs)) ([], [])
{-# INLINE unzip #-}