e669760 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add missing init/tail functions in Stream/StreamK

e669760Harendra Kumar3 years ago

core/src/Streamly/Internal/Data/Stream/Eliminate.hs+30−2

compare with master

unifiedsplit

@@ -43,6 +43,7 @@ module Streamly.Internal.Data.Stream.Eliminate

43 , null

44 , head

45 , headElse

46+ , init

46 , tail

47 , tail

47 , last

48 , last

48 , elem

49 , elem

@@ -90,7 +91,7 @@ import Foreign.Storable (Storable)

90import GHC.Exts (SpecConstrAnnotation(..))

91import GHC.Exts (SpecConstrAnnotation(..))

91import GHC.Types (SPEC(..))

92import GHC.Types (SPEC(..))

92import Streamly.Internal.Data.Parser (ParseError(..))

93import Streamly.Internal.Data.Parser (ParseError(..))

93−import Streamly.Internal.Data.SVar.Type (defState)

94+import Streamly.Internal.Data.SVar.Type (adaptState, defState)

94import Streamly.Internal.Data.Unbox (Unbox)

95import Streamly.Internal.Data.Unbox (Unbox)

96import Streamly.Internal.Data.Maybe.Strict (Maybe'(..))

97import Streamly.Internal.Data.Maybe.Strict (Maybe'(..))

@@ -105,7 +106,7 @@ import qualified Streamly.Internal.Data.Stream.Transform as StreamD

105

106

106import Prelude hiding

107import Prelude hiding

107 ( Foldable(..), all, any, head, last, lookup, mapM, mapM_

108 ( Foldable(..), all, any, head, last, lookup, mapM, mapM_

108− , notElem, splitAt, tail, (!!))

109+ , notElem, splitAt, init, tail, (!!))

109import Data.Foldable (length)

110import Data.Foldable (length)

110import Streamly.Internal.Data.Stream.Type

111import Streamly.Internal.Data.Stream.Type

111

112

@@ -391,9 +392,35 @@ head = foldrM (\x _ -> return (Just x)) (return Nothing)

391headElse :: Monad m => a -> Stream m a -> m a

392headElse :: Monad m => a -> Stream m a -> m a

392headElse a = foldrM (\x _ -> return x) (return a)

393headElse a = foldrM (\x _ -> return x) (return a)

393

394

395+{-# INLINE_NORMAL init #-}

396+init :: Monad m => Stream m a -> m (Maybe (Stream m a))

397+init stream = do

398+ r <- uncons stream

399+ case r of

400+ Nothing -> return Nothing

401+ Just (h, Stream step1 state1) ->

402+ return $ Just $ Stream step (h, state1)

403+

404+ where

405+

406+ step gst (a, s1) = do

407+ res <- step1 (adaptState gst) s1

408+ return $

409+ case res of

410+ Yield x s -> Yield a (x, s)

411+ Skip s -> Skip (a, s)

412+ Stop -> Stop

413+

414+-- | Same as:

415+--

416+-- >>> tail = fmap (fmap snd) . Stream.uncons

417+--

394-- Does not fuse, has the same performance as the StreamK version.

418-- Does not fuse, has the same performance as the StreamK version.

419+--

395{-# INLINE_NORMAL tail #-}

420{-# INLINE_NORMAL tail #-}

396tail :: Monad m => Stream m a -> m (Maybe (Stream m a))

421tail :: Monad m => Stream m a -> m (Maybe (Stream m a))

422+tail = fmap (fmap snd) . uncons

423+{-

397tail (UnStream step state) = go SPEC state

424tail (UnStream step state) = go SPEC state

398 where

425 where

399 go !_ st = do

426 go !_ st = do

@@ -402,6 +429,7 @@ tail (UnStream step state) = go SPEC state

402 Yield _ s -> return (Just $ Stream step s)

429 Yield _ s -> return (Just $ Stream step s)

403 Skip s -> go SPEC s

430 Skip s -> go SPEC s

404 Stop -> return Nothing

431 Stop -> return Nothing

432+-}

405

433

406-- XXX will it fuse? need custom impl?

434-- XXX will it fuse? need custom impl?

407{-# INLINE_NORMAL last #-}

435{-# INLINE_NORMAL last #-}

core/src/Streamly/Internal/Data/Stream/Transform.hs+51

@@ -83,6 +83,8 @@ module Streamly.Internal.Data.Stream.Transform

84 -- * Trimming

85 -- | Produce a subset of the stream trimmed at ends.

86+ , initNonEmpty

87+ , tailNonEmpty

86 , take

88 , take

87 , takeWhile

89 , takeWhile

88 , takeWhileM

90 , takeWhileM

@@ -1074,6 +1076,55 @@ repeated = undefined

1074-- Trimming

1076-- Trimming

1075------------------------------------------------------------------------------

1077------------------------------------------------------------------------------

1076

1078

1079+-- | init for non-empty streams, fails for empty stream case.

1080+--

1081+{-# INLINE initNonEmpty #-}

1082+initNonEmpty :: Monad m => Stream m a -> Stream m a

1083+initNonEmpty (Stream step1 state1) = Stream step (Nothing, state1)

1084+

1085+ where

1086+

1087+ step gst (Nothing, s1) = do

1088+ r <- step1 (adaptState gst) s1

1089+ return $

1090+ case r of

1091+ Yield x s -> Skip (Just x, s)

1092+ Skip s -> Skip (Nothing, s)

1093+ Stop -> error "initNonEmpty: empty Stream"

1094+

1095+ step gst (Just a, s1) = do

1096+ r <- step1 (adaptState gst) s1

1097+ return $

1098+ case r of

1099+ Yield x s -> Yield a (Just x, s)

1100+ Skip s -> Skip (Just a, s)

1101+ Stop -> Stop

1102+

1103+-- | tail for non-empty streams, fails for empty stream case.

1104+--

1105+-- See also 'tail' for a non-partial version of this function..

1106+{-# INLINE tailNonEmpty #-}

1107+tailNonEmpty :: Monad m => Stream m a -> Stream m a

1108+tailNonEmpty (Stream step1 state1) = Stream step (Nothing, state1)

1109+

1110+ where

1111+

1112+ step gst (Nothing, s1) = do

1113+ r <- step1 (adaptState gst) s1

1114+ return $

1115+ case r of

1116+ Yield x s -> Skip (Just x, s)

1117+ Skip s -> Skip (Nothing, s)

1118+ Stop -> error "tailNonEmpty: empty Stream"

1119+

1120+ step gst (Just a, s1) = do

1121+ r <- step1 (adaptState gst) s1

1122+ return $

1123+ case r of

1124+ Yield x s -> Yield x (Just x, s)

1125+ Skip s -> Skip (Just a, s)

1126+ Stop -> Stop

1127+

1077-- | Take all consecutive elements at the end of the stream for which the

1128-- | Take all consecutive elements at the end of the stream for which the

1078-- predicate is true.

1129-- predicate is true.

1079--

1130--

core/src/Streamly/Internal/Data/StreamK/Type.hs+56−11

@@ -97,8 +97,11 @@ module Streamly.Internal.Data.StreamK.Type

97 -- ** Specific Folds

98 , drain

99 , null

100+ , headNonEmpty

100 , tail

101 , tail

102+ , tailNonEmpty

101 , init

103 , init

104+ , initNonEmpty

102

105

103 -- * Mapping

106 -- * Mapping

104 , map

107 , map

@@ -1811,14 +1814,42 @@ iterateMWith cns step = go

1811 !next <- s

1814 !next <- s

1812 foldStreamShared st stp sng yld (return next `cns` go (step next))

1815 foldStreamShared st stp sng yld (return next `cns` go (step next))

1813

1816

1814−{-# INLINE headPartial #-}

1817+-- | head for non-empty streams, fails for empty stream case.

1818+--

1819+{-# INLINE headNonEmpty #-}

1815−headPartial :: Monad m => StreamK m a -> m a

1820+headNonEmpty :: Monad m => StreamK m a -> m a

1816−headPartial = foldrM (\x _ -> return x) (error "head of nil")

1821+headNonEmpty = foldrM (\x _ -> return x) (error "headNonEmpty: empty stream")

1822+

1823+-- | init for non-empty streams, fails for empty stream case.

1824+--

1825+-- See also 'init' for a non-partial version of this function..

1826+{-# INLINE initNonEmpty #-}

1827+initNonEmpty :: Stream m a -> Stream m a

1828+initNonEmpty = go0

1829+

1830+ where

1831+

1832+ go0 m = mkStream $ \st yld sng stp ->

1833+ let stop = error "initNonEmpty: Empty Stream."

1834+ single _ = stp

1835+ yieldk a r = foldStream st yld sng stp (go1 a r)

1836+ in foldStream st yieldk single stop m

1837+

1838+ go1 a r = mkStream $ \st yld sng stp ->

1839+ let stop = stp

1840+ single _ = sng a

1841+ yieldk a1 r1 = yld a (go1 a1 r1)

1842+ in foldStream st yieldk single stop r

1817

1843

1818−{-# INLINE tailPartial #-}

1844+-- | tail for non-empty streams, fails for empty stream case.

1845+--

1846+-- See also 'tail' for a non-partial version of this function..

1847+--

1848+-- Note: this is same as "drop 1" with error on empty stream.

1849+{-# INLINE tailNonEmpty #-}

1819−tailPartial :: StreamK m a -> StreamK m a

1850+tailNonEmpty :: StreamK m a -> StreamK m a

1820−tailPartial m = mkStream $ \st yld sng stp ->

1851+tailNonEmpty m = mkStream $ \st yld sng stp ->

1821− let stop = error "tail of nil"

1852+ let stop = error "tailNonEmpty: empty stream"

1822 single _ = stp

1853 single _ = stp

1823 yieldk _ r = foldStream st yld sng stp r

1854 yieldk _ r = foldStream st yld sng stp r

1824 in foldStream st yieldk single stop m

1855 in foldStream st yieldk single stop m

@@ -1883,10 +1914,10 @@ mfix f = mkStream $ \st yld sng stp ->

1883 where

1914 where

1884

1915

1885 -- fix the head element of the stream

1916 -- fix the head element of the stream

1886− xs = fix (f . headPartial)

1917+ xs = fix (f . headNonEmpty)

1887

1918

1888 -- now fix the tail recursively

1919 -- now fix the tail recursively

1889− ys = mfix (tailPartial . f)

1920+ ys = mfix (tailNonEmpty . f)

1890

1921

1891-------------------------------------------------------------------------------

1922-------------------------------------------------------------------------------

1892-- Conversions

1923-- Conversions

@@ -1917,6 +1948,15 @@ uncons m =

1917 yieldk a r = pure (Just (a, r))

1948 yieldk a r = pure (Just (a, r))

1918 in foldStream defState yieldk single stop m

1949 in foldStream defState yieldk single stop m

1919

1950

1951+-- Note that this is not a StreamK -> StreamK because then we cannot handle the

1952+-- empty stream case without making this a partial function.

1953+--

1954+-- See tailNonEmpty as well above.

1955+

1956+-- | Same as:

1957+--

1958+-- >>> tail = fmap (fmap snd) StreamK.uncons

1959+--

1920{-# INLINE tail #-}

1960{-# INLINE tail #-}

1921tail :: Applicative m => StreamK m a -> m (Maybe (StreamK m a))

1961tail :: Applicative m => StreamK m a -> m (Maybe (StreamK m a))

1922tail =

1962tail =

@@ -1925,9 +1965,14 @@ tail =

1925 yieldk _ r = pure $ Just r

1965 yieldk _ r = pure $ Just r

1926 in foldStream defState yieldk single stop

1966 in foldStream defState yieldk single stop

1927

1967

1928−-- | Extract all but the last element of the stream, if any.

1968+-- Note that this is not a StreamK -> StreamK because then we cannot handle the

1969+-- empty stream case without making this a partial function.

1929--

1970--

1930−-- Note: This will end up buffering the entire stream.

1971+-- XXX How do we implement unsnoc? Make StreamK a monad and return the

1972+-- remaining stream as a result value in the monad?

1973+

1974+-- | Extract all but the last element of the stream, if any. This will end up

1975+-- evaluating the last element as well to find out that it is last.

1931--

1976--

1932-- /Pre-release/

1977-- /Pre-release/

1933{-# INLINE init #-}

1978{-# INLINE init #-}