@@ -239,7 +239,7 @@ where
239239
240240import Control.Monad (void)
241241import Data.Bifunctor (first)
242−import Data.Either (isLeft, isRight, fromLeft, fromRight)
242+import Data.Either (isLeft, isRight)
243243import Data.Functor.Identity (Identity(..))
244244import Data.Int (Int64)
245245import Data.Map.Strict (Map)
@@ -247,6 +247,7 @@ import Data.Maybe (isJust, fromJust)
247247#if __GLASGOW_HASKELL__ < 804
248248import Data.Semigroup (Semigroup((<>)))
249249#endif
250+import Streamly.Internal.BaseCompat (fromLeft, fromRight)
250251import Streamly.Internal.Data.Either.Strict
251252 (Either'(..), fromLeft', fromRight', isLeft', isRight')
252253import Streamly.Internal.Data.Pipe.Type (Pipe (..), PipeState(..))
@@ -1109,6 +1110,22 @@ distribute = Prelude.foldr (teeWith (:)) (fromPure [])
11091110-- Partitioning
11101111------------------------------------------------------------------------------
11111112
1113+{-# INLINE partitionByMUsing #-}
1114+partitionByMUsing :: Monad m =>
1115+ ( (x -> y -> (x, y))
1116+ -> Fold m (Either b c) x
1117+ -> Fold m (Either b c) y
1118+ -> Fold m (Either b c) (x, y)
1119+ )
1120+ -> (a -> m (Either b c))
1121+ -> Fold m b x
1122+ -> Fold m c y
1123+ -> Fold m a (x, y)
1124+partitionByMUsing t f fld1 fld2 =
1125+ let l = lmap (fromLeft undefined) fld1 -- :: Fold m (Either b c) x
1126+ r = lmap (fromRight undefined) fld2 -- :: Fold m (Either b c) y
1127+ in lmapM f (t (,) (filter isLeft l) (filter isRight r))
1128+
11121129-- | Partition the input over two folds using an 'Either' partitioning
11131130-- predicate.
11141131--
@@ -1163,10 +1180,7 @@ distribute = Prelude.foldr (teeWith (:)) (fromPure [])
11631180{-# INLINE partitionByM #-}
11641181partitionByM :: Monad m
11651182 => (a -> m (Either b c)) -> Fold m b x -> Fold m c y -> Fold m a (x, y)
1166−partitionByM f fld1 fld2 =
1167− let l = lmap (fromLeft undefined) fld1 -- :: Fold m (Either b c) x
1168− r = lmap (fromRight undefined) fld2 -- :: Fold m (Either b c) y
1169− in lmapM f (teeWith (,) (filter isLeft l) (filter isRight r))
1183+partitionByM = partitionByMUsing teeWith
11701184
11711185-- | Similar to 'partitionByM' but terminates when the first fold terminates.
11721186--
@@ -1174,58 +1188,7 @@ partitionByM f fld1 fld2 =
11741188{-# INLINE partitionByFstM #-}
11751189partitionByFstM :: Monad m
11761190 => (a -> m (Either b c)) -> Fold m b x -> Fold m c y -> Fold m a (x, y)
1177−partitionByFstM f (Fold stepL beginL doneL) (Fold stepR beginR doneR) =
1178− Fold step begin done
1179−
1180− where
1181−
1182− begin = do
1183− resL <- beginL
1184− resR <- beginR
1185− case resL of
1186− Partial sL -> return $
1187− case resR of
1188− Partial sR -> Partial $ RunFstBoth sL sR
1189− Done bR -> Partial $ RunFstLeft sL bR
1190− Done bL -> do
1191− case resR of
1192− Partial sR -> do
1193− bR <- doneR sR
1194− return $ Done (bL, bR)
1195− Done bR -> return $ Done (bL, bR)
1196−
1197− step (RunFstBoth sL sR) a = do
1198− r <- f a
1199− case r of
1200− Left b -> do
1201− res <- stepL sL b
1202− case res of
1203− Partial sres -> return $ Partial $ RunFstBoth sres sR
1204− Done bL -> do
1205− bR <- doneR sR
1206− return $ Done (bL, bR)
1207−
1208− Right c -> do
1209− res <- stepR sR c
1210− return
1211− $ Partial
1212− $ case res of
1213− Partial sres -> RunFstBoth sL sres
1214− Done bres -> RunFstLeft sL bres
1215− step (RunFstLeft sL bR) a = do
1216− r <- f a
1217− case r of
1218− Left b -> do
1219− res <- stepL sL b
1220− return
1221− $ case res of
1222− Partial sres -> Partial $ RunFstLeft sres bR
1223− Done bres -> Done (bres, bR)
1224− Right _ -> return $ Partial $ RunFstLeft sL bR
1225−
1226− done (RunFstBoth sL sR) = (,) <$> doneL sL <*> doneR sR
1227−
1228− done (RunFstLeft sL bR) = (,bR) <$> doneL sL
1191+partitionByFstM = partitionByMUsing teeWithFst
12291192
12301193-- | Similar to 'partitionByM' but terminates when any fold terminates.
12311194--
@@ -1233,48 +1196,7 @@ partitionByFstM f (Fold stepL beginL doneL) (Fold stepR beginR doneR) =
12331196{-# INLINE partitionByMinM #-}
12341197partitionByMinM :: Monad m =>
12351198 (a -> m (Either b c)) -> Fold m b x -> Fold m c y -> Fold m a (x, y)
1236−partitionByMinM f (Fold stepL beginL doneL) (Fold stepR beginR doneR) =
1237− Fold step begin done
1238−
1239− where
1240−
1241− begin = do
1242− resL <- beginL
1243− resR <- beginR
1244− case resL of
1245− Partial sL ->
1246− case resR of
1247− Partial sR -> return $ Partial $ Tuple' sL sR
1248− Done bR -> do
1249− bL <- doneL sL
1250− return $ Done (bL, bR)
1251− Done bL -> do
1252− case resR of
1253− Partial sR -> do
1254− bR <- doneR sR
1255− return $ Done (bL, bR)
1256− Done bR -> return $ Done (bL, bR)
1257−
1258− step (Tuple' sL sR) a = do
1259− r <- f a
1260− case r of
1261− Left b -> do
1262− res <- stepL sL b
1263− case res of
1264− Partial sres -> return $ Partial $ Tuple' sres sR
1265− Done bL -> do
1266− bR <- doneR sR
1267− return $ Done (bL, bR)
1268−
1269− Right c -> do
1270− res <- stepR sR c
1271− case res of
1272− Partial sres -> return $ Partial $ Tuple' sL sres
1273− Done bR -> do
1274− bL <- doneL sL
1275− return $ Done (bL, bR)
1276−
1277− done (Tuple' sL sR) = (,) <$> doneL sL <*> doneR sR
1199+partitionByMinM = partitionByMUsing teeWithMin
12781200
12791201-- Note: we could use (a -> Bool) instead of (a -> Either b c), but the latter
12801202-- makes the signature clearer as to which case belongs to which fold.
@@ -1595,6 +1517,22 @@ classify fld = classifyWith fst (lmap snd fld)
15951517-- Unzipping
15961518------------------------------------------------------------------------------
15971519
1520+{-# INLINE unzipWithMUsing #-}
1521+unzipWithMUsing :: Monad m =>
1522+ ( (x -> y -> (x, y))
1523+ -> Fold m (b, c) x
1524+ -> Fold m (b, c) y
1525+ -> Fold m (b, c) (x, y)
1526+ )
1527+ -> (a -> m (b, c))
1528+ -> Fold m b x
1529+ -> Fold m c y
1530+ -> Fold m a (x, y)
1531+unzipWithMUsing t f fld1 fld2 =
1532+ let f1 = lmap fst fld1 -- :: Fold m (b, c) b
1533+ f2 = lmap snd fld2 -- :: Fold m (b, c) c
1534+ in lmapM f (t (,) f1 f2)
1535+
15981536-- | Like 'unzipWith' but with a monadic splitter function.
15991537--
16001538-- @unzipWithM k f1 f2 = lmapM k (unzip f1 f2)@
@@ -1603,10 +1541,7 @@ classify fld = classifyWith fst (lmap snd fld)
16031541{-# INLINE unzipWithM #-}
16041542unzipWithM :: Monad m
16051543 => (a -> m (b,c)) -> Fold m b x -> Fold m c y -> Fold m a (x,y)
1606−unzipWithM f fld1 fld2 =
1607− let f1 = lmap fst fld1 -- :: Fold m (b, c) b
1608− f2 = lmap snd fld2 -- :: Fold m (b, c) c
1609− in lmapM f (teeWith (,) f1 f2)
1544+unzipWithM = unzipWithMUsing teeWith
16101545
16111546-- | Similar to 'unzipWithM' but terminates when the first fold terminates.
16121547--
@@ -1614,52 +1549,7 @@ unzipWithM f fld1 fld2 =
16141549{-# INLINE unzipWithFstM #-}
16151550unzipWithFstM :: Monad m =>
16161551 (a -> m (b, c)) -> Fold m b x -> Fold m c y -> Fold m a (x, y)
1617−unzipWithFstM f (Fold stepL beginL doneL) (Fold stepR beginR doneR) =
1618− Fold step begin done
1619−
1620− where
1621−
1622− begin = do
1623− resL <- beginL
1624− resR <- beginR
1625− case resL of
1626− Partial sL -> return $ Partial $
1627− case resR of
1628− Partial sR -> RunFstBoth sL sR
1629− Done bR -> RunFstLeft sL bR
1630− Done bL -> do
1631− case resR of
1632− Partial sR -> do
1633− bR <- doneR sR
1634− return $ Done (bL, bR)
1635− Done bR -> return $ Done (bL, bR)
1636−
1637− step (RunFstBoth sL sR) a = do
1638− (b, c) <- f a
1639− resL <- stepL sL b
1640− resR <- stepR sR c
1641− case resL of
1642− Partial sresL -> return $ Partial $
1643− case resR of
1644− Partial sresR -> RunFstBoth sresL sresR
1645− Done bresR -> RunFstLeft sresL bresR
1646− Done bL -> do
1647− case resR of
1648− Partial sresR -> do
1649− bR <- doneR sresR
1650− return $ Done (bL, bR)
1651− Done bR -> return $ Done (bL, bR)
1652−
1653− step (RunFstLeft sL bR) a = do
1654− (b, _) <- f a
1655− resL <- stepL sL b
1656− return
1657− $ case resL of
1658− Partial sresL -> Partial $ RunFstLeft sresL bR
1659− Done bresL -> Done (bresL, bR)
1660−
1661− done (RunFstBoth sL sR) = (,) <$> doneL sL <*> doneR sR
1662− done (RunFstLeft sL bR) = (,bR) <$> doneL sL
1552+unzipWithFstM = unzipWithMUsing teeWithFst
16631553
16641554-- | Similar to 'unzipWithM' but terminates when any fold terminates.
16651555--
@@ -1667,47 +1557,7 @@ unzipWithFstM f (Fold stepL beginL doneL) (Fold stepR beginR doneR) =
16671557{-# INLINE unzipWithMinM #-}
16681558unzipWithMinM :: Monad m =>
16691559 (a -> m (b,c)) -> Fold m b x -> Fold m c y -> Fold m a (x,y)
1670−unzipWithMinM f (Fold stepL beginL doneL) (Fold stepR beginR doneR) =
1671− Fold step begin done
1672−
1673− where
1674−
1675− begin = do
1676− resL <- beginL
1677− resR <- beginR
1678− case resL of
1679− Partial sL ->
1680− case resR of
1681− Partial sR -> return $ Partial $ Tuple' sL sR
1682− Done bR -> do
1683− bL <- doneL sL
1684− return $ Done (bL, bR)
1685− Done bL -> do
1686− case resR of
1687− Partial sR -> do
1688− bR <- doneR sR
1689− return $ Done (bL, bR)
1690− Done bR -> return $ Done (bL, bR)
1691−
1692− step (Tuple' sL sR) a = do
1693− (b, c) <- f a
1694− resL <- stepL sL b
1695− resR <- stepR sR c
1696− case resL of
1697− Partial sresL ->
1698− case resR of
1699− Partial sresR -> return $ Partial $ Tuple' sresL sresR
1700− Done bR -> do
1701− bL <- doneL sL
1702− return $ Done (bL, bR)
1703− Done bL -> do
1704− case resR of
1705− Partial sresR -> do
1706− bR <- doneR sresR
1707− return $ Done (bL, bR)
1708− Done bR -> return $ Done (bL, bR)
1709−
1710− done (Tuple' sL sR) = (,) <$> doneL sL <*> doneR sR
1560+unzipWithMinM = unzipWithMUsing teeWithMin
17111561
17121562-- | Split elements in the input stream into two parts using a pure splitter
17131563-- function, direct each part to a different fold and zip the results.