e669760 · streamly · git.iko.soy

git.iko.soy/streamly

Add missing init/tail functions in Stream/StreamK

commite669760f55d8757d1e65bca6c00e58687899ae22 authorHarendra Kumar <harendra@composewell.com> dateMon 22 Jan 2024, 18:10 (3 years ago) parent8ef0939

3 files changed+137−13

compare with master

unifiedsplit

core/src/Streamly/Internal/Data/Stream/Eliminate.hs+30−2 core/src/Streamly/Internal/Data/Stream/Transform.hs+51 core/src/Streamly/Internal/Data/StreamK/Type.hs+56−11

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

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

4343 , null

4444 , head

4545 , headElse

46+ , init

4647 , tail

4748 , last

4849 , elem

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

9091import GHC.Exts (SpecConstrAnnotation(..))

9192import GHC.Types (SPEC(..))

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

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

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

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

9596

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

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

105106

106107import Prelude hiding

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

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

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

109110import Data.Foldable (length)

110111import Streamly.Internal.Data.Stream.Type

111112

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

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

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

393394

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+--

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

419+--

395420{-# INLINE_NORMAL tail #-}

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

422+tail = fmap (fmap snd) . uncons

423+{-

397424tail (UnStream step state) = go SPEC state

398425 where

399426 go !_ st = do

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

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

403430 Skip s -> go SPEC s

404431 Stop -> return Nothing

432+-}

405433

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

407435{-# INLINE_NORMAL last #-}

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

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

8383

8484 -- * Trimming

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

86+ , initNonEmpty

87+ , tailNonEmpty

8688 , take

8789 , takeWhile

8890 , takeWhileM

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

10741076-- Trimming

10751077------------------------------------------------------------------------------

10761078

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+

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

10781129-- predicate is true.

10791130--

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

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

9797 -- ** Specific Folds

9898 , drain

9999 , null

100+ , headNonEmpty

100101 , tail

102+ , tailNonEmpty

101103 , init

104+ , initNonEmpty

102105

103106 -- * Mapping

104107 , map

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

18111814 !next <- s

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

18131816

1814−{-# INLINE headPartial #-}

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

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

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

1818+--

1819+{-# INLINE headNonEmpty #-}

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

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

18171843

1818−{-# INLINE tailPartial #-}

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

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

1821− let stop = error "tail of nil"

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

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

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

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

18221853 single _ = stp

18231854 yieldk _ r = foldStream st yld sng stp r

18241855 in foldStream st yieldk single stop m

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

18831914 where

18841915

18851916 -- fix the head element of the stream

1886− xs = fix (f . headPartial)

1917+ xs = fix (f . headNonEmpty)

18871918

18881919 -- now fix the tail recursively

1889− ys = mfix (tailPartial . f)

1920+ ys = mfix (tailNonEmpty . f)

18901921

18911922-------------------------------------------------------------------------------

18921923-- Conversions

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

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

19181949 in foldStream defState yieldk single stop m

19191950

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+--

19201960{-# INLINE tail #-}

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

19221962tail =

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

19251965 yieldk _ r = pure $ Just r

19261966 in foldStream defState yieldk single stop

19271967

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.

19291970--

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.

19311976--

19321977-- /Pre-release/

19331978{-# INLINE init #-}