77d11d3 · streamly · git.iko.soy

git.iko.soy/streamly

Add Monad instance for Unfold

commit77d11d37bb8a402dd00fbcd69c5e61773aa2c396 authorHarendra Kumar <harendra@composewell.com> dateThu 18 Feb 2021, 03:27 (6 years ago) parent4f9fc7d

5 files changed+215−165

compare with master

unifiedsplit

benchmark/Streamly/Benchmark/Data/NestedUnfoldOps.hs−86 benchmark/Streamly/Benchmark/Data/Unfold.hs+141−34 src/Streamly/Internal/Data/Unfold.hs+2−36 src/Streamly/Internal/Data/Unfold/Types.hs+67−1 test/Streamly/Test/Data/Unfold.hs+5−8

benchmark/Streamly/Benchmark/Data/NestedUnfoldOps.hs−86

@@ -18,14 +18,6 @@ concatCount :: Int -> Int

1818concatCount linearCount =

1919 round (((1 + 8 * fromIntegral linearCount)**(1/2::Double) - 1) / 2)

2020

21−-- double nested loop

22−nestedCount2 :: Int -> Int

23−nestedCount2 linearCount = round (fromIntegral linearCount**(1/2::Double))

24−

25−-- triple nested loop

26−nestedCount3 :: Int -> Int

27−nestedCount3 linearCount = round (fromIntegral linearCount**(1/3::Double))

28−

2921-------------------------------------------------------------------------------

3022-- Stream generation and elimination

3123-------------------------------------------------------------------------------

@@ -39,26 +31,6 @@ source n = UF.enumerateFromToIntegral n

3931-- Benchmark ops

4032-------------------------------------------------------------------------------

4133

42−{-# INLINE toNull #-}

43−toNull :: MonadIO m => Int -> Int -> m ()

44−toNull linearCount start = do

45− let end = start + nestedCount2 linearCount

46− UF.fold

47− (UF.map (\(x, y) -> x + y)

48− $ UF.outerProduct (source end) (source end))

49− FL.drain (start, start)

50−

51−{-# INLINE toNull3 #-}

52−toNull3 :: MonadIO m => Int -> Int -> m ()

53−toNull3 linearCount start = do

54− let end = start + nestedCount3 linearCount

55− UF.fold

56− (UF.map (\(x, y) -> x + y)

57− $ UF.outerProduct (source end)

58− ((UF.map (\(x, y) -> x + y)

59− $ UF.outerProduct (source end) (source end))))

60− FL.drain (start, (start, start))

61−

6234{-# INLINE concat #-}

6335concat :: MonadIO m => Int -> Int -> m ()

6436concat linearCount start = do

@@ -66,61 +38,3 @@ concat linearCount start = do

6638 UF.fold

6739 (UF.concat (source end) (source end))

6840 FL.drain start

69−

70−{-# INLINE toList #-}

71−toList :: MonadIO m => Int -> Int -> m [Int]

72−toList linearCount start = do

73− let end = start + nestedCount2 linearCount

74− UF.fold

75− (UF.map (\(x, y) -> x + y)

76− $ UF.outerProduct (source end) (source end))

77− FL.toList (start, start)

78−

79−{-# INLINE toListSome #-}

80−toListSome :: MonadIO m => Int -> Int -> m [Int]

81−toListSome linearCount start = do

82− let end = start + nestedCount2 linearCount

83− UF.fold

84− (UF.take 1000 $ (UF.map (\(x, y) -> x + y)

85− $ UF.outerProduct (source end) (source end)))

86− FL.toList (start, start)

87−

88−{-# INLINE filterAllOut #-}

89−filterAllOut :: MonadIO m => Int -> Int -> m ()

90−filterAllOut linearCount start = do

91− let end = start + nestedCount2 linearCount

92− UF.fold

93− (UF.filter (< 0)

94− $ UF.map (\(x, y) -> x + y)

95− $ UF.outerProduct (source end) (source end))

96− FL.drain (start, start)

97−

98−{-# INLINE filterAllIn #-}

99−filterAllIn :: MonadIO m => Int -> Int -> m ()

100−filterAllIn linearCount start = do

101− let end = start + nestedCount2 linearCount

102− UF.fold

103− (UF.filter (> 0)

104− $ UF.map (\(x, y) -> x + y)

105− $ UF.outerProduct (source end) (source end))

106− FL.drain (start, start)

107−

108−{-# INLINE filterSome #-}

109−filterSome :: MonadIO m => Int -> Int -> m ()

110−filterSome linearCount start = do

111− let end = start + nestedCount2 linearCount

112− UF.fold

113− (UF.filter (> 1100000)

114− $ UF.map (\(x, y) -> x + y)

115− $ UF.outerProduct (source end) (source end))

116− FL.drain (start, start)

117−

118−{-# INLINE breakAfterSome #-}

119−breakAfterSome :: MonadIO m => Int -> Int -> m ()

120−breakAfterSome linearCount start = do

121− let end = start + nestedCount2 linearCount

122− UF.fold

123− (UF.takeWhile (<= 1100000)

124− $ UF.map (\(x, y) -> x + y)

125− $ UF.outerProduct (source end) (source end))

126− FL.drain (start, start)

benchmark/Streamly/Benchmark/Data/Unfold.hs+141−34

@@ -20,7 +20,7 @@

2020module Main (main) where

2121

2222import Control.DeepSeq (NFData(..))

23−import Control.Exception (SomeException)

23+import Control.Exception (SomeException, ErrorCall, try)

2424import Streamly.Internal.Data.Unfold (Unfold)

2525import System.IO (Handle, hClose)

2626import System.Random (randomRIO)

@@ -397,29 +397,13 @@ teeZipWith size start =

397397 drainProductDefault (size + start) (UF.teeZipWith (+)) start

398398

399399-------------------------------------------------------------------------------

400−-- Nested

400+-- Applicative

401401-------------------------------------------------------------------------------

402402

403−{-# INLINE concatMapM #-}

404−concatMapM :: Monad m => Int -> Int -> m ()

405−concatMapM size start =

406− drainGeneration (UF.concatMapM unfoldInGen unfoldOut) start

407−

408− where

409−

410− sizeOuter = 100

411− sizeInner = size `div` sizeOuter

412−

413− unfoldInGen i =

414− return

415− $ UF.supply (UF.enumerateFromToIntegral (i + sizeInner)) i

416−

417− unfoldOut = UF.enumerateFromToIntegral (start + sizeOuter)

418−

419403{-# INLINE toNullAp #-}

420404toNullAp :: Monad m => Int -> Int -> m ()

421−toNullAp linearCount start =

422− let end = start + Nested.nestedCount2 linearCount

405+toNullAp value start =

406+ let end = start + nthRoot 2 value

423407 s = Nested.source end

424408 in UF.fold ((+) <$> s <*> s) FL.drain start

425409

@@ -431,6 +415,128 @@ _apDiscardFst = undefined

431415_apDiscardSnd :: Int -> Int -> m ()

432416_apDiscardSnd = undefined

433417

418+-------------------------------------------------------------------------------

419+-- Monad

420+-------------------------------------------------------------------------------

421+

422+nthRoot :: Double -> Int -> Int

423+nthRoot n value = round (fromIntegral value**(1/n))

424+

425+{-# INLINE concatMapM #-}

426+concatMapM :: Monad m => Int -> Int -> m ()

427+concatMapM value start =

428+ val `seq` drainGeneration (UF.concatMapM unfoldInGen unfoldOut) start

429+

430+ where

431+

432+ val = nthRoot 2 value

433+ unfoldInGen i = return (UF.enumerateFromToIntegral (i + val))

434+ unfoldOut = UF.enumerateFromToIntegral (start + val)

435+

436+{-# INLINE toNull #-}

437+toNull :: Monad m => Int -> Int -> m ()

438+toNull value start =

439+ let end = start + nthRoot 2 value

440+ src = Nested.source end

441+ u = do

442+ x <- src

443+ y <- src

444+ return (x + y)

445+ in UF.fold u FL.drain start

446+

447+

448+{-# INLINE toNull3 #-}

449+toNull3 :: Monad m => Int -> Int -> m ()

450+toNull3 value start =

451+ let end = start + nthRoot 3 value

452+ src = Nested.source end

453+ u = do

454+ x <- src

455+ y <- src

456+ z <- src

457+ return (x + y + z)

458+ in UF.fold u FL.drain start

459+

460+{-# INLINE toList #-}

461+toList :: Monad m => Int -> Int -> m [Int]

462+toList value start = do

463+ let end = start + nthRoot 2 value

464+ src = Nested.source end

465+ u = do

466+ x <- src

467+ y <- src

468+ return (x + y)

469+ in UF.fold u FL.toList start

470+

471+{-# INLINE toListSome #-}

472+toListSome :: Monad m => Int -> Int -> m [Int]

473+toListSome value start = do

474+ let end = start + nthRoot 2 value

475+ src = Nested.source end

476+ u = do

477+ x <- src

478+ y <- src

479+ return (x + y)

480+ in UF.fold (UF.take 1000 u) FL.toList start

481+

482+{-# INLINE filterAllOut #-}

483+filterAllOut :: Monad m => Int -> Int -> m ()

484+filterAllOut value start = do

485+ let end = start + nthRoot 2 value

486+ src = Nested.source end

487+ u = do

488+ x <- src

489+ y <- src

490+ let s = x + y

491+ if s < 0

492+ then return s

493+ else UF.nilM (return . const ())

494+ in UF.fold u FL.drain start

495+

496+{-# INLINE filterAllIn #-}

497+filterAllIn :: Monad m => Int -> Int -> m ()

498+filterAllIn value start = do

499+ let end = start + nthRoot 2 value

500+ src = Nested.source end

501+ u = do

502+ x <- src

503+ y <- src

504+ let s = x + y

505+ if s > 0

506+ then return s

507+ else UF.nilM (return . const ())

508+ in UF.fold u FL.drain start

509+

510+{-# INLINE filterSome #-}

511+filterSome :: Monad m => Int -> Int -> m ()

512+filterSome value start = do

513+ let end = start + nthRoot 2 value

514+ src = Nested.source end

515+ u = do

516+ x <- src

517+ y <- src

518+ let s = x + y

519+ if s > 1100000

520+ then return s

521+ else UF.nilM (return . const ())

522+ in UF.fold u FL.drain start

523+

524+{-# INLINE breakAfterSome #-}

525+breakAfterSome :: Int -> Int -> IO ()

526+breakAfterSome value start =

527+ let end = start + nthRoot 2 value

528+ src = Nested.source end

529+ u = do

530+ x <- src

531+ y <- src

532+ let s = x + y

533+ if s > 1100000

534+ then error "break"

535+ else return s

536+ in do

537+ (_ :: Either ErrorCall ()) <- try $ UF.fold u FL.drain start

538+ return ()

539+

434540-------------------------------------------------------------------------------

435541-- Benchmarks

436542-------------------------------------------------------------------------------

@@ -529,29 +635,30 @@ o_1_space_zip size =

529635o_1_space_nested :: Int -> [Benchmark]

530636o_1_space_nested size =

531637 [ bgroup

532− "outer-product"

533− [ benchIO "toNullAp" $ toNullAp size

534− , benchIO "toNull" $ Nested.toNull size

535− , benchIO "toNull3" $ Nested.toNull3 size

536− , benchIO "concat" $ Nested.concat size

537− , benchIO "breakAfterSome" $ Nested.breakAfterSome size

538− , benchIO "filterAllOut" $ Nested.filterAllOut size

539− , benchIO "filterAllIn" $ Nested.filterAllIn size

540− , benchIO "filterSome" $ Nested.filterSome size

541− , benchIO "concatMapM (100 x n/100)" $ concatMapM size

638+ "nested"

639+ [ benchIO "(<*>) (sqrt n x sqrt n)" $ toNullAp size

542640 -- Unimplemented

543− -- , benchIO "ap" $ ap size

544641 -- , benchIO "apDiscardFst" $ apDiscardFst size

545642 -- , benchIO "apDiscardSnd" $ apDiscardSnd size

643+

644+ , benchIO "concatMapM (sqrt n x sqrt n)" $ concatMapM size

645+ , benchIO "(>>=) (sqrt n x sqrt n)" $ toNull size

646+ , benchIO "(>>=) (cubert n x cubert n x cubert n)" $ toNull3 size

647+ , benchIO "breakAfterSome" $ breakAfterSome size

648+ , benchIO "filterAllOut" $ filterAllOut size

649+ , benchIO "filterAllIn" $ filterAllIn size

650+ , benchIO "filterSome" $ filterSome size

651+

652+ , benchIO "concat" $ Nested.concat size

546653 ]

547654 ]

548655

549656o_n_space_nested :: Int -> [Benchmark]

550657o_n_space_nested size =

551658 [ bgroup

552− "outer-product"

553− [ benchIO "toList" $ Nested.toList size

554− , benchIO "toListSome" $ Nested.toListSome size

659+ "nested"

660+ [ benchIO "toList" $ toList size

661+ , benchIO "toListSome" $ toListSome size

555662 ]

556663 ]

557664

src/Streamly/Internal/Data/Unfold.hs+2−36

@@ -156,7 +156,8 @@ import Streamly.Internal.Data.IOFinalizer

156156 (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)

157157import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))

158158import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)

159−import Streamly.Internal.Data.Unfold.Types (Unfold(..), lmap, map, const)

159+import Streamly.Internal.Data.Unfold.Types

160+ (Unfold(..), lmap, map, const, concatMapM)

160161import System.Mem (performMajorGC)

161162

162163import qualified Prelude

@@ -973,41 +974,6 @@ outerProduct (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject

973974 Skip s -> Skip (OuterProductInner ost sy s x)

974975 Stop -> Skip (OuterProductOuter ost sy)

975976

976−-- XXX This can be used to implement a Monad instance for "Unfold m ()".

977−

978−data ConcatMapState s1 s2 = ConcatMapOuter s1 | ConcatMapInner s1 s2

979−

980−-- | Map an unfold generating action to each element of an unfold and

981−-- flatten the results into a single stream.

982−--

983−{-# INLINE_NORMAL concatMapM #-}

984−concatMapM :: Monad m

985− => (b -> m (Unfold m Void c)) -> Unfold m a b -> Unfold m a c

986−concatMapM f (Unfold step1 inject1) = Unfold step inject

987− where

988− inject x = do

989− s <- inject1 x

990− return $ ConcatMapOuter s

991−

992− {-# INLINE_LATE step #-}

993− step (ConcatMapOuter st) = do

994− r <- step1 st

995− case r of

996− Yield x s -> do

997− Unfold step2 inject2 <- f x

998− innerSt <- inject2 undefined

999− return $ Skip (ConcatMapInner s (Stream (\_ ss -> step2 ss)

1000− innerSt))

1001− Skip s -> return $ Skip (ConcatMapOuter s)

1002− Stop -> return Stop

1003−

1004− step (ConcatMapInner ost (UnStream istep ist)) = do

1005− r <- istep defState ist

1006− return $ case r of

1007− Yield x s -> Yield x (ConcatMapInner ost (Stream istep s))

1008− Skip s -> Skip (ConcatMapInner ost (Stream istep s))

1009− Stop -> Skip (ConcatMapOuter ost)

1010−

1011977------------------------------------------------------------------------------

1012978-- Exceptions

1013979------------------------------------------------------------------------------

src/Streamly/Internal/Data/Unfold/Types.hs+67−1

@@ -14,6 +14,8 @@ module Streamly.Internal.Data.Unfold.Types

1414 , apSequence

1515 , apDiscardSnd

1616 , cross

17+ , concatMapM

18+ , concatMap

1719 )

1820where

1921

@@ -21,7 +23,7 @@ where

2123

2224import Streamly.Internal.Data.Stream.StreamD.Step (Step(..))

2325

24−import Prelude hiding (const, map)

26+import Prelude hiding (const, map, concatMap)

2527

2628------------------------------------------------------------------------------

2729-- Monadic Unfolds

@@ -150,3 +152,67 @@ instance Monad m => Applicative (Unfold m a) where

150152

151153 -- {-# INLINE (<*) #-}

152154 -- (<*) = apDiscardSnd

155+

156+------------------------------------------------------------------------------

157+-- Monad

158+------------------------------------------------------------------------------

159+

160+data ConcatMapState m b s1 x =

161+ ConcatMapOuter x s1

162+ | forall s2. ConcatMapInner x s1 s2 (s2 -> m (Step s2 b))

163+

164+-- | Map an unfold generating action to each element of an unfold and

165+-- flatten the results into a single stream.

166+--

167+{-# INLINE_NORMAL concatMapM #-}

168+concatMapM :: Monad m

169+ => (b -> m (Unfold m a c)) -> Unfold m a b -> Unfold m a c

170+concatMapM f (Unfold step1 inject1) = Unfold step inject

171+

172+ where

173+

174+ inject x = do

175+ s <- inject1 x

176+ return $ ConcatMapOuter x s

177+

178+ {-# INLINE_LATE step #-}

179+ step (ConcatMapOuter seed st) = do

180+ r <- step1 st

181+ case r of

182+ Yield x s -> do

183+ Unfold step2 inject2 <- f x

184+ innerSt <- inject2 seed

185+ return $ Skip (ConcatMapInner seed s innerSt step2)

186+ Skip s -> return $ Skip (ConcatMapOuter seed s)

187+ Stop -> return Stop

188+

189+ step (ConcatMapInner seed ost ist istep) = do

190+ r <- istep ist

191+ return $ case r of

192+ Yield x s -> Yield x (ConcatMapInner seed ost s istep)

193+ Skip s -> Skip (ConcatMapInner seed ost s istep)

194+ Stop -> Skip (ConcatMapOuter seed ost)

195+

196+{-# INLINE concatMap #-}

197+concatMap :: Monad m => (b -> Unfold m a c) -> Unfold m a b -> Unfold m a c

198+concatMap f = concatMapM (return . f)

199+

200+-- Note: concatMap and Monad instance for unfolds have performance comparable

201+-- to Stream. In fact, concatMap is slower than Stream, that may be some

202+-- optimization issue though.

203+--

204+-- | Example:

205+--

206+-- >>> u = do { x <- Unfold.lmap fst Unfold.fromList; y <- Unfold.lmap snd Unfold.fromList; return (x,y); }

207+-- >>> Stream.toList $ Stream.unfold u ([1,2],[3,4])

208+-- [(1,3),(1,4),(2,3),(2,4)]

209+--

210+instance Monad m => Monad (Unfold m a) where

211+ {-# INLINE return #-}

212+ return = pure

213+

214+ {-# INLINE (>>=) #-}

215+ (>>=) = flip concatMap

216+

217+ -- {-# INLINE (>>) #-}

218+ -- (>>) = (*>)

test/Streamly/Test/Data/Unfold.hs+5−8

@@ -347,14 +347,11 @@ outerProduct =

347347

348348concatMapM :: Bool

349349concatMapM =

350− let unfInF b =

351− modify (+ 1)

352− >> return

353− (UF.supply (UF.replicateM 10) (modify (+ 1) >> return b))

354− listInF b = replicate 10 b

355− unfOut = UF.enumerateFromToIntegral 10

356− unf = UF.concatMapM unfInF unfOut

357− list = List.concatMap listInF [1 .. 10]

350+ let inner b =

351+ let u = UF.lmap (\_ -> modify (+ 1) >> return b) (UF.replicateM 10)

352+ in modify (+ 1) >> return u

353+ unf = UF.concatMapM inner (UF.enumerateFromToIntegral 10)

354+ list = List.concatMap (replicate 10) [1 .. 10]

358355 in testUnfoldMD unf 1 0 110 list

359356

360357-------------------------------------------------------------------------------