810ba5b · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Refactor teeWith

810ba5bHarendra Kumar5 years ago

src/Streamly/Internal/Data/Fold.hs+9−115

compare with master

unifiedsplit

@@ -239,6 +239,7 @@ where

239

240import Control.Monad (void)

241import Data.Bifunctor (first)

242+import Data.Either (isLeft, isRight, fromLeft, fromRight)

242import Data.Functor.Identity (Identity(..))

243import Data.Functor.Identity (Identity(..))

243import Data.Int (Int64)

244import Data.Int (Int64)

244import Data.Map.Strict (Map)

245import Data.Map.Strict (Map)

@@ -1162,67 +1163,10 @@ distribute = Prelude.foldr (teeWith (:)) (fromPure [])

1162{-# INLINE partitionByM #-}

1163{-# INLINE partitionByM #-}

1163partitionByM :: Monad m

1164partitionByM :: Monad m

1164 => (a -> m (Either b c)) -> Fold m b x -> Fold m c y -> Fold m a (x, y)

1165 => (a -> m (Either b c)) -> Fold m b x -> Fold m c y -> Fold m a (x, y)

1165−partitionByM f (Fold stepL beginL doneL) (Fold stepR beginR doneR) =

1166+partitionByM f fld1 fld2 =

1166− Fold step begin done

1167+ let l = lmap (fromLeft undefined) fld1 -- :: Fold m (Either b c) x

1167−

1168+ r = lmap (fromRight undefined) fld2 -- :: Fold m (Either b c) y

1168− where

1169+ in lmapM f (teeWith (,) (filter isLeft l) (filter isRight r))

1169−

1170− begin = do

1171− resL <- beginL

1172− resR <- beginR

1173− return

1174− $ case resL of

1175− Partial sL ->

1176− Partial

1177− $ case resR of

1178− Partial sR -> RunBoth sL sR

1179− Done bR -> RunLeft sL bR

1180− Done bL ->

1181− case resR of

1182− Partial sR -> Partial $ RunRight bL sR

1183− Done bR -> Done (bL, bR)

1184−

1185− step (RunBoth sL sR) a = do

1186− r <- f a

1187− case r of

1188− Left b -> do

1189− res <- stepL sL b

1190− return

1191− $ Partial

1192− $ case res of

1193− Partial sres -> RunBoth sres sR

1194− Done bres -> RunRight bres sR

1195− Right c -> do

1196− res <- stepR sR c

1197− return

1198− $ Partial

1199− $ case res of

1200− Partial sres -> RunBoth sL sres

1201− Done bres -> RunLeft sL bres

1202− step (RunLeft sL bR) a = do

1203− r <- f a

1204− case r of

1205− Left b -> do

1206− res <- stepL sL b

1207− return

1208− $ case res of

1209− Partial sres -> Partial $ RunLeft sres bR

1210− Done bres -> Done (bres, bR)

1211− Right _ -> return $ Partial $ RunLeft sL bR

1212− step (RunRight bL sR) a = do

1213− r <- f a

1214− case r of

1215− Left _ -> return $ Partial $ RunRight bL sR

1216− Right c -> do

1217− res <- stepR sR c

1218− return

1219− $ case res of

1220− Partial sres -> Partial $ RunRight bL sres

1221− Done bres -> Done (bL, bres)

1222−

1223− done (RunBoth sL sR) = (,) <$> doneL sL <*> doneR sR

1224− done (RunLeft sL bR) = (,bR) <$> doneL sL

1225− done (RunRight bL sR) = (bL,) <$> doneR sR

1226

1170

1227-- | Similar to 'partitionByM' but terminates when the first fold terminates.

1171-- | Similar to 'partitionByM' but terminates when the first fold terminates.

1228--

1172--

@@ -1659,60 +1603,10 @@ classify fld = classifyWith fst (lmap snd fld)

1659{-# INLINE unzipWithM #-}

1603{-# INLINE unzipWithM #-}

1660unzipWithM :: Monad m

1604unzipWithM :: Monad m

1661 => (a -> m (b,c)) -> Fold m b x -> Fold m c y -> Fold m a (x,y)

1605 => (a -> m (b,c)) -> Fold m b x -> Fold m c y -> Fold m a (x,y)

1662−unzipWithM f (Fold stepL beginL doneL) (Fold stepR beginR doneR) =

1606+unzipWithM f fld1 fld2 =

1663− Fold step begin done

1607+ let f1 = lmap fst fld1 -- :: Fold m (b, c) b

1664−

1608+ f2 = lmap snd fld2 -- :: Fold m (b, c) c

1665− where

1609+ in lmapM f (teeWith (,) f1 f2)

1666−

1667− begin = do

1668− resL <- beginL

1669− resR <- beginR

1670− return

1671− $ case resL of

1672− Partial sL ->

1673− Partial

1674− $ case resR of

1675− Partial sR -> RunBoth sL sR

1676− Done bR -> RunLeft sL bR

1677− Done bL ->

1678− case resR of

1679− Partial sR -> Partial $ RunRight bL sR

1680− Done bR -> Done (bL, bR)

1681−

1682− step (RunBoth sL sR) a = do

1683− (b, c) <- f a

1684− resL <- stepL sL b

1685− resR <- stepR sR c

1686− case resL of

1687− Partial sresL ->

1688− return

1689− $ Partial

1690− $ case resR of

1691− Partial sresR -> RunBoth sresL sresR

1692− Done bresR -> RunLeft sresL bresR

1693− Done bresL ->

1694− return

1695− $ case resR of

1696− Partial sresR -> Partial $ RunRight bresL sresR

1697− Done bresR -> Done (bresL, bresR)

1698− step (RunLeft sL bR) a = do

1699− (b, _) <- f a

1700− resL <- stepL sL b

1701− return

1702− $ case resL of

1703− Partial sresL -> Partial $ RunLeft sresL bR

1704− Done bresL -> Done (bresL, bR)

1705− step (RunRight bL sR) a = do

1706− (_, c) <- f a

1707− resR <- stepR sR c

1708− return

1709− $ case resR of

1710− Partial sresR -> Partial $ RunRight bL sresR

1711− Done bresR -> Done (bL, bresR)

1712−

1713− done (RunBoth sL sR) = (,) <$> doneL sL <*> doneR sR

1714− done (RunLeft sL bR) = (,bR) <$> doneL sL

1715− done (RunRight bL sR) = (bL,) <$> doneR sR

1716

1610

1717-- | Similar to 'unzipWithM' but terminates when the first fold terminates.

1611-- | Similar to 'unzipWithM' but terminates when the first fold terminates.

1718--

1612--

src/Streamly/Internal/Data/Fold/Type.hs+21−51

@@ -235,7 +235,6 @@ module Streamly.Internal.Data.Fold.Type

235 , serial_

236

237 -- ** Parallel Distribution

238− , GenericRunner(..)

239 , FstRunner(..)

238 , FstRunner(..)

240 , teeWith

239 , teeWith

241 , teeWithFst

240 , teeWithFst

@@ -747,11 +746,11 @@ serial_ (Fold stepL initialL _) (Fold stepR initialR extractR) =

747 Partial sR -> extractR sR

746 Partial sR -> extractR sR

748 Done rR -> return rR

747 Done rR -> return rR

749

748

750−{-# ANN type GenericRunner Fuse #-}

749+{-# ANN type TeeState Fuse #-}

751−data GenericRunner sL sR bL bR

750+data TeeState sL sR bL bR

752− = RunBoth !sL !sR

751+ = TeeBoth !sL !sR

753− | RunLeft !sL !bR

752+ | TeeLeft !bR !sL

754− | RunRight !bL !sR

753+ | TeeRight !bL !sR

755

754

756{-# ANN type FstRunner Fuse #-}

755{-# ANN type FstRunner Fuse #-}

757data FstRunner sL sR b

756data FstRunner sL sR b

@@ -776,61 +775,32 @@ data FstRunner sL sR b

776--

775--

777{-# INLINE teeWith #-}

776{-# INLINE teeWith #-}

778teeWith :: Monad m => (a -> b -> c) -> Fold m x a -> Fold m x b -> Fold m x c

777teeWith :: Monad m => (a -> b -> c) -> Fold m x a -> Fold m x b -> Fold m x c

779−teeWith f (Fold stepL beginL doneL) (Fold stepR beginR doneR) =

778+teeWith f (Fold stepL initialL extractL) (Fold stepR initialR extractR) =

780− Fold step begin done

779+ Fold step initial extract

781

780

782 where

781 where

783

782

784− begin = do

783+ runBoth actionL actionR = do

785− resL <- beginL

784+ resL <- actionL

786− resR <- beginR

785+ resR <- actionR

787 return

786 return

788 $ case resL of

787 $ case resL of

789 Partial sl ->

788 Partial sl ->

790 Partial

789 Partial

791 $ case resR of

790 $ case resR of

792− Partial sr -> RunBoth sl sr

791+ Partial sr -> TeeBoth sl sr

793− Done br -> RunLeft sl br

792+ Done br -> TeeLeft br sl

794− Done bl -> bimap (RunRight bl) (f bl) resR

793+ Done bl -> bimap (TeeRight bl) (f bl) resR

795

794

796− step (RunBoth sL sR) a = do

795+ initial = runBoth initialL initialR

797− resL <- stepL sL a

798− resR <- stepR sR a

799− case resL of

800− Partial sL1 ->

801− return

802− $ Partial

803− $ case resR of

804− Partial sR1 -> RunBoth sL1 sR1

805− Done bR -> RunLeft sL1 bR

806− Done bL ->

807− return

808− $ case resR of

809− Partial sR1 -> Partial $ RunRight bL sR1

810− Done bR -> Done $ f bL bR

811− step (RunLeft sL bR) a = do

812− resL <- stepL sL a

813− return

814− $ case resL of

815− Partial sL1 -> Partial $ RunLeft sL1 bR

816− Done bL -> Done $ f bL bR

817− step (RunRight bL sR) a = do

818− resR <- stepR sR a

819− return

820− $ case resR of

821− Partial sR1 -> Partial $ RunRight bL sR1

822− Done bR -> Done $ f bL bR

823

796

824− done (RunBoth sL sR) = do

797+ step (TeeBoth sL sR) a = runBoth (stepL sL a) (stepR sR a)

825− bL <- doneL sL

798+ step (TeeLeft bR sL) a = bimap (TeeLeft bR) (`f` bR) <$> stepL sL a

826− bR <- doneR sR

799+ step (TeeRight bL sR) a = bimap (TeeRight bL) (f bL) <$> stepR sR a

827− return $ f bL bR

800+

828− done (RunLeft sL bR) = do

801+ extract (TeeBoth sL sR) = f <$> extractL sL <*> extractR sR

829− bL <- doneL sL

802+ extract (TeeLeft bR sL) = (`f` bR) <$> extractL sL

830− return $ f bL bR

803+ extract (TeeRight bL sR) = f bL <$> extractR sR

831− done (RunRight bL sR) = do

832− bR <- doneR sR

833− return $ f bL bR

834

804

835-- | Like 'teeWith' but terminates as soon as the first fold terminates.

805-- | Like 'teeWith' but terminates as soon as the first fold terminates.

836--

806--