b47918e · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add Category and Arrow instances for unfolds

b47918eHarendra Kumar6 years ago

src/Streamly/Internal/Data/Unfold.hs+29−111

compare with master

unifiedsplit

@@ -8,6 +8,34 @@

8-- Stability : experimental

9-- Portability : GHC

10--

11+-- = Unfolds and Streams

12+--

13+-- An 'Unfold' type is the same as the direct style 'Stream' type except that

14+-- it uses an inject function to determine the initial state of the stream

15+-- based on an input. A stream is a special case of Unfold when the static

16+-- input is unit or Void.

17+--

18+-- This allows an important optimization to occur in several cases, making the

19+-- 'Unfold' a more efficient abstraction. Consider the 'concatMap' and

20+-- 'concatUnfold' operations, the latter is more efficient. 'concatMap'

21+-- generates a new stream object from each element in the stream by applying

22+-- the supplied function to the element, the stream object includes the "step"

23+-- function as well as the initial "state" of the stream. Since the stream is

24+-- generated dynamically the compiler does not know the step function or the

25+-- state type statically at compile time, therefore, it cannot inline it. On

26+-- the other hand in case of 'concatUnfold' the compiler has visibility into

27+-- the unfold's state generation function, therefore, the compiler knows all

28+-- the types statically and it can inline the inject as well as the step

29+-- functions, generating efficient code. Essentially, the stream is not opaque

30+-- to the consumer in case of unfolds, the consumer knows how to generate the

31+-- stream from a seed using a known "inject" and "step" functions.

32+--

33+-- A Stream is like a data object whereas unfold is like a function. Being

34+-- function like, an Unfold is an instance of 'Category' and 'Arrow' type

35+-- classes.

36+--

37+-- = Unfolds and Folds

38+--

11-- Streams forcing a closed control flow loop can be categorized under

39-- Streams forcing a closed control flow loop can be categorized under

12-- two types, unfolds and folds, both of these are duals of each other.

40-- two types, unfolds and folds, both of these are duals of each other.

13--

41--

@@ -149,15 +177,12 @@ import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_)

149import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)

177import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)

150import Data.Maybe (isNothing)

178import Data.Maybe (isNothing)

151import Data.Void (Void)

179import Data.Void (Void)

152−import Fusion.Plugin.Types (Fuse(..))

153import GHC.Types (SPEC(..))

180import GHC.Types (SPEC(..))

154import Streamly.Internal.Data.Fold.Types (Fold(..))

181import Streamly.Internal.Data.Fold.Types (Fold(..))

155import Streamly.Internal.Data.IOFinalizer

182import Streamly.Internal.Data.IOFinalizer

156 (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)

183 (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)

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

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

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

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

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

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

161import System.Mem (performMajorGC)

186import System.Mem (performMajorGC)

162

187

163import qualified Prelude

188import qualified Prelude

@@ -169,6 +194,7 @@ import qualified Streamly.Internal.Data.Stream.StreamK as K (uncons)

169import qualified Streamly.Internal.Data.Stream.StreamK.Type as K

194import qualified Streamly.Internal.Data.Stream.StreamK.Type as K

170

195

171import Streamly.Internal.Data.SVar

196import Streamly.Internal.Data.SVar

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

172import Prelude

198import Prelude

173 hiding (concat, map, mapM, takeWhile, take, filter, const, zipWith

199 hiding (concat, map, mapM, takeWhile, take, filter, const, zipWith

174 , drop, dropWhile)

200 , drop, dropWhile)

@@ -399,34 +425,6 @@ effect eff = Unfold step inject

399 step True = eff >>= \r -> return $ Yield r False

425 step True = eff >>= \r -> return $ Yield r False

400 step False = return Stop

426 step False = return Stop

401

427

402−-- XXX change it to yieldM or change yieldM in Prelude to singletonM

403−--

404−-- | Lift a monadic function into an unfold generating a singleton stream.

405−--

406−{-# INLINE singletonM #-}

407−singletonM :: Monad m => (a -> m b) -> Unfold m a b

408−singletonM f = Unfold step inject

409− where

410− inject x = return $ Just x

411− {-# INLINE_LATE step #-}

412− step (Just x) = f x >>= \r -> return $ Yield r Nothing

413− step Nothing = return Stop

414−

415−-- | Lift a pure function into an unfold generating a singleton stream.

416−--

417−{-# INLINE singleton #-}

418−singleton :: Monad m => (a -> b) -> Unfold m a b

419−singleton f = singletonM $ return . f

420−

421−-- | Identity unfold. Generates a singleton stream with the seed as the only

422−-- element in the stream.

423−--

424−-- > identity = singletonM return

425−--

426−{-# INLINE identity #-}

427−identity :: Monad m => Unfold m a a

428−identity = singletonM return

429−

430-- | Convert a list of pure values to a 'Stream'

428-- | Convert a list of pure values to a 'Stream'

431{-# INLINE_LATE fromList #-}

429{-# INLINE_LATE fromList #-}

432fromList :: Monad m => Unfold m [a] a

430fromList :: Monad m => Unfold m [a] a

@@ -842,53 +840,6 @@ fromProducer = Unfold step (return . FromSVarRead)

842-- Zipping

840-- Zipping

843-------------------------------------------------------------------------------

841-------------------------------------------------------------------------------

844

842

845−{-# INLINE_NORMAL zipWithM #-}

846−zipWithM :: Monad m

847− => (a -> b -> m c) -> Unfold m x a -> Unfold m y b -> Unfold m (x, y) c

848−zipWithM f (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject

849−

850− where

851−

852− inject (x, y) = do

853− s1 <- inject1 x

854− s2 <- inject2 y

855− return (s1, s2, Nothing)

856−

857− {-# INLINE_LATE step #-}

858− step (s1, s2, Nothing) = do

859− r <- step1 s1

860− return $

861− case r of

862− Yield x s -> Skip (s, s2, Just x)

863− Skip s -> Skip (s, s2, Nothing)

864− Stop -> Stop

865−

866− step (s1, s2, Just x) = do

867− r <- step2 s2

868− case r of

869− Yield y s -> do

870− z <- f x y

871− return $ Yield z (s1, s, Nothing)

872− Skip s -> return $ Skip (s1, s, Just x)

873− Stop -> return Stop

874−

875−-- | Divide the input into two unfolds and then zip the outputs to a single

876−-- stream.

877−--

878−-- @

879−-- S.mapM_ print

880−-- $ S.concatUnfold (UF.zipWith (,) UF.identity (UF.singleton sqrt))

881−-- $ S.map (\x -> (x,x))

882−-- $ S.fromList [1..10]

883−-- @

884−--

885−-- /Internal/

886−--

887−{-# INLINE zipWith #-}

888−zipWith :: Monad m

889− => (a -> b -> c) -> Unfold m x a -> Unfold m y b -> Unfold m (x, y) c

890−zipWith f = zipWithM (\a b -> return (f a b))

891−

892-- | Distribute the input to two unfolds and then zip the outputs to a single

843-- | Distribute the input to two unfolds and then zip the outputs to a single

893-- stream.

844-- stream.

894--

845--

@@ -907,39 +858,6 @@ teeZipWith f unf1 unf2 = lmap (\x -> (x,x)) $ zipWith f unf1 unf2

907-- Nested

858-- Nested

908-------------------------------------------------------------------------------

859-------------------------------------------------------------------------------

909

860

910−{-# ANN type ConcatState Fuse #-}

911−data ConcatState s1 s2 = ConcatOuter s1 | ConcatInner s1 s2

912−

913−-- | Apply the second unfold to each output element of the first unfold and

914−-- flatten the output in a single stream.

915−--

916−-- /Internal/

917−--

918−{-# INLINE_NORMAL concat #-}

919−concat :: Monad m => Unfold m a b -> Unfold m b c -> Unfold m a c

920−concat (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject

921− where

922− inject x = do

923− s <- inject1 x

924− return $ ConcatOuter s

925−

926− {-# INLINE_LATE step #-}

927− step (ConcatOuter st) = do

928− r <- step1 st

929− case r of

930− Yield x s -> do

931− innerSt <- inject2 x

932− return $ Skip (ConcatInner s innerSt)

933− Skip s -> return $ Skip (ConcatOuter s)

934− Stop -> return Stop

935−

936− step (ConcatInner ost ist) = do

937− r <- step2 ist

938− return $ case r of

939− Yield x s -> Yield x (ConcatInner ost s)

940− Skip s -> Skip (ConcatInner ost s)

941− Stop -> Skip (ConcatOuter ost)

942−

943data OuterProductState s1 s2 sy x y =

861data OuterProductState s1 s2 sy x y =

944 OuterProductOuter s1 y | OuterProductInner s1 sy s2 x

862 OuterProductOuter s1 y | OuterProductInner s1 sy s2 x

945

863

src/Streamly/Internal/Data/Unfold/Types.hs+153−4

@@ -8,6 +8,11 @@

9module Streamly.Internal.Data.Unfold.Types

10 ( Unfold (..)

11+ , singletonM

12+ , singleton

13+ , identity

14+ , ConcatState (..)

15+ , concat

11 , lmap

16 , lmap

12 , map

17 , map

13 , const

18 , const

@@ -16,19 +21,27 @@ module Streamly.Internal.Data.Unfold.Types

16 , cross

21 , cross

17 , concatMapM

22 , concatMapM

18 , concatMap

23 , concatMap

24+ , zipWithM

25+ , zipWith

19 )

26 )

20where

27where

22#include "inline.hs"

29#include "inline.hs"

31+import Control.Arrow (Arrow(..))

32+import Control.Category (Category(..))

33+import Fusion.Plugin.Types (Fuse(..))

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

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

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

36+import Prelude hiding (const, map, concat, concatMap, zipWith)

28------------------------------------------------------------------------------

38------------------------------------------------------------------------------

29-- Monadic Unfolds

39-- Monadic Unfolds

30------------------------------------------------------------------------------

40------------------------------------------------------------------------------

42+-- The order of arguments allows 'Category' and 'Arrow' instances but precludes

43+-- contravariant and contra-applicative.

44+--

32-- | An @Unfold m a b@ is a generator of a stream of values of type @b@ from a

45-- | An @Unfold m a b@ is a generator of a stream of values of type @b@ from a

33-- seed of type 'a' in 'Monad' @m@.

46-- seed of type 'a' in 'Monad' @m@.

34--

47--

@@ -47,7 +60,7 @@ data Unfold m a b =

47-- /Internal/

60-- /Internal/

48{-# INLINE_NORMAL lmap #-}

61{-# INLINE_NORMAL lmap #-}

49lmap :: (a -> c) -> Unfold m c b -> Unfold m a b

62lmap :: (a -> c) -> Unfold m c b -> Unfold m a b

50−lmap f (Unfold ustep uinject) = Unfold ustep (uinject . f)

63+lmap f (Unfold ustep uinject) = Unfold ustep (uinject Prelude.. f)

52------------------------------------------------------------------------------

65------------------------------------------------------------------------------

53-- Functor

66-- Functor

@@ -142,7 +155,7 @@ cross (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject

142--

155--

143instance Monad m => Applicative (Unfold m a) where

156instance Monad m => Applicative (Unfold m a) where

144 {-# INLINE pure #-}

157 {-# INLINE pure #-}

145− pure = const . return

158+ pure = const Prelude.. return

146

159

147 {-# INLINE (<*>) #-}

160 {-# INLINE (<*>) #-}

148 u1 <*> u2 = fmap (\(a, b) -> a b) (cross u1 u2)

161 u1 <*> u2 = fmap (\(a, b) -> a b) (cross u1 u2)

@@ -195,7 +208,7 @@ concatMapM f (Unfold step1 inject1) = Unfold step inject

195

208

196{-# INLINE concatMap #-}

209{-# INLINE concatMap #-}

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

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

198−concatMap f = concatMapM (return . f)

211+concatMap f = concatMapM (return Prelude.. f)

199

212

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

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

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

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

@@ -216,3 +229,139 @@ instance Monad m => Monad (Unfold m a) where

216

229

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

230 -- {-# INLINE (>>) #-}

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

231 -- (>>) = (*>)

232+

233+-------------------------------------------------------------------------------

234+-- Category

235+-------------------------------------------------------------------------------

236+

237+-- XXX change it to yieldM or change yieldM in Prelude to singletonM

238+--

239+-- | Lift a monadic function into an unfold generating a singleton stream.

240+--

241+{-# INLINE singletonM #-}

242+singletonM :: Monad m => (a -> m b) -> Unfold m a b

243+singletonM f = Unfold step inject

244+

245+ where

246+

247+ inject x = return $ Just x

248+

249+ {-# INLINE_LATE step #-}

250+ step (Just x) = f x >>= \r -> return $ Yield r Nothing

251+ step Nothing = return Stop

252+

253+-- | Lift a pure function into an unfold generating a singleton stream.

254+--

255+{-# INLINE singleton #-}

256+singleton :: Monad m => (a -> b) -> Unfold m a b

257+singleton f = singletonM $ return Prelude.. f

258+

259+-- | Identity unfold. Generates a singleton stream with the seed as the only

260+-- element in the stream.

261+--

262+-- > identity = singletonM return

263+--

264+{-# INLINE identity #-}

265+identity :: Monad m => Unfold m a a

266+identity = singletonM return

267+

268+{-# ANN type ConcatState Fuse #-}

269+data ConcatState s1 s2 = ConcatOuter s1 | ConcatInner s1 s2

270+

271+-- | Apply the second unfold to each output element of the first unfold and

272+-- flatten the output in a single stream.

273+--

274+-- /Internal/

275+--

276+{-# INLINE_NORMAL concat #-}

277+concat :: Monad m => Unfold m a b -> Unfold m b c -> Unfold m a c

278+concat (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject

279+

280+ where

281+

282+ inject x = do

283+ s <- inject1 x

284+ return $ ConcatOuter s

285+

286+ {-# INLINE_LATE step #-}

287+ step (ConcatOuter st) = do

288+ r <- step1 st

289+ case r of

290+ Yield x s -> do

291+ innerSt <- inject2 x

292+ return $ Skip (ConcatInner s innerSt)

293+ Skip s -> return $ Skip (ConcatOuter s)

294+ Stop -> return Stop

295+

296+ step (ConcatInner ost ist) = do

297+ r <- step2 ist

298+ return $ case r of

299+ Yield x s -> Yield x (ConcatInner ost s)

300+ Skip s -> Skip (ConcatInner ost s)

301+ Stop -> Skip (ConcatOuter ost)

302+

303+instance Monad m => Category (Unfold m) where

304+ {-# INLINE id #-}

305+ id = identity

306+

307+ {-# INLINE (.) #-}

308+ (.) = flip concat

309+

310+-------------------------------------------------------------------------------

311+-- Arrow

312+-------------------------------------------------------------------------------

313+

314+{-# INLINE_NORMAL zipWithM #-}

315+zipWithM :: Monad m

316+ => (a -> b -> m c) -> Unfold m x a -> Unfold m y b -> Unfold m (x, y) c

317+zipWithM f (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject

318+

319+ where

320+

321+ inject (x, y) = do

322+ s1 <- inject1 x

323+ s2 <- inject2 y

324+ return (s1, s2, Nothing)

325+

326+ {-# INLINE_LATE step #-}

327+ step (s1, s2, Nothing) = do

328+ r <- step1 s1

329+ return $

330+ case r of

331+ Yield x s -> Skip (s, s2, Just x)

332+ Skip s -> Skip (s, s2, Nothing)

333+ Stop -> Stop

334+

335+ step (s1, s2, Just x) = do

336+ r <- step2 s2

337+ case r of

338+ Yield y s -> do

339+ z <- f x y

340+ return $ Yield z (s1, s, Nothing)

341+ Skip s -> return $ Skip (s1, s, Just x)

342+ Stop -> return Stop

343+

344+-- | Divide the input into two unfolds and then zip the outputs to a single

345+-- stream.

346+--

347+-- @

348+-- S.mapM_ print

349+-- $ S.concatUnfold (UF.zipWith (,) UF.identity (UF.singleton sqrt))

350+-- $ S.map (\x -> (x,x))

351+-- $ S.fromList [1..10]

352+-- @

353+--

354+-- /Internal/

355+--

356+{-# INLINE zipWith #-}

357+zipWith :: Monad m

358+ => (a -> b -> c) -> Unfold m x a -> Unfold m y b -> Unfold m (x, y) c

359+zipWith f = zipWithM (\a b -> return (f a b))

360+

361+{-# ANN module "HLint: ignore Use zip" #-}

362+instance Monad m => Arrow (Unfold m) where

363+ {-# INLINE arr #-}

364+ arr = singleton

365+

366+ {-# INLINE (***) #-}

367+ (***) = zipWith (,)