src/Streamly/Internal/Data/Stream/Prelude.hs · streamly

git.iko.soy/streamly

on c3dcabc help sign in

streamly/src/Streamly/Internal/Data/Stream/Prelude.hs

349 lines 10 kB history raw

1{-# OPTIONS_GHC -Wno-orphans #-}

3#include "inline.hs"

5-- |

6-- Module : Streamly.Internal.Data.Stream.Prelude

8--

9-- License : BSD3

10-- Maintainer : streamly@composewell.com

11-- Stability : experimental

12-- Portability : GHC

13--

14--

15module Streamly.Internal.Data.Stream.Prelude

16 (

17 -- * Stream Conversion

18 fromStreamS

19 , toStreamS

21 -- * Running Effects

22 , drain

24 -- * Conversion operations

25 , fromList

26 , toList

28 -- * Fold operations

29 , foldrM

30 , foldrMx

31 , foldr

33 , foldlx'

34 , foldlMx'

35 , foldl'

36 , runFold

37 , parselMx'

39 -- Lazy left folds are useful only for reversing the stream

40 , foldlS

41 , foldlT

43 , scanlx'

44 , scanlMx'

45 , postscanlx'

46 , postscanlMx'

48 -- * Zip style operations

49 , eqBy

50 , cmpBy

52 -- * Foldable instance

53 , minimum

54 , maximum

56 -- * Nesting

57 , K.concatMapBy

58 , K.concatMap

60 -- * Fold Utilities

61 , foldWith

62 , foldMapWith

63 , forEachWith

64 )

65where

67import Control.Monad.Catch (MonadThrow)

68import Control.Monad.Trans.Class (MonadTrans(..))

69import Prelude hiding (foldr, minimum, maximum)

70import qualified Prelude

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

73import Streamly.Internal.Data.Parser (Step)

75#ifdef USE_STREAMK_ONLY

76import qualified Streamly.Internal.Data.Stream.StreamK as S

77#else

78import qualified Streamly.Internal.Data.Stream.StreamD as S

79#endif

81import Streamly.Internal.Data.Stream.StreamK (IsStream(..))

82import qualified Streamly.Internal.Data.Stream.StreamK as K

83import qualified Streamly.Internal.Data.Stream.StreamD as D

85------------------------------------------------------------------------------

86-- Conversion to and from direct style stream

87------------------------------------------------------------------------------

89-- These definitions are dependent on what is imported as S

90{-# INLINE fromStreamS #-}

91fromStreamS :: (IsStream t, Monad m) => S.Stream m a -> t m a

92fromStreamS = fromStream . S.toStreamK

94{-# INLINE toStreamS #-}

95toStreamS :: (IsStream t, Monad m) => t m a -> S.Stream m a

96toStreamS = S.fromStreamK . toStream

98------------------------------------------------------------------------------

99-- Conversions

100------------------------------------------------------------------------------

101

102{-# INLINE_EARLY drain #-}

103drain :: (IsStream t, Monad m) => t m a -> m ()

104drain m = D.drain $ D.fromStreamK (toStream m)

105{-# RULES "drain fallback to CPS" [1]

106 forall a. D.drain (D.fromStreamK a) = K.drain a #-}

107

108------------------------------------------------------------------------------

109-- Conversions

110------------------------------------------------------------------------------

111

112-- |

113-- @

114-- fromList = 'Prelude.foldr' 'K.cons' 'K.nil'

115-- @

116--

117-- Construct a stream from a list of pure values. This is more efficient than

118-- 'K.fromFoldable' for serial streams.

119--

120-- @since 0.4.0

121{-# INLINE_EARLY fromList #-}

122fromList :: (Monad m, IsStream t) => [a] -> t m a

123fromList = fromStreamS . S.fromList

124{-# RULES "fromList fallback to StreamK" [1]

125 forall a. S.toStreamK (S.fromList a) = K.fromFoldable a #-}

126

127-- | Convert a stream into a list in the underlying monad.

128--

129-- @since 0.1.0

130{-# INLINE toList #-}

131toList :: (Monad m, IsStream t) => t m a -> m [a]

132toList m = S.toList $ toStreamS m

133

134------------------------------------------------------------------------------

135-- Folds

136------------------------------------------------------------------------------

137

138{-# INLINE foldrM #-}

139foldrM :: (Monad m, IsStream t) => (a -> m b -> m b) -> m b -> t m a -> m b

140foldrM step acc m = S.foldrM step acc $ toStreamS m

141

142{-# INLINE foldrMx #-}

143foldrMx :: (Monad m, IsStream t)

144 => (a -> m x -> m x) -> m x -> (m x -> m b) -> t m a -> m b

145foldrMx step final project m = D.foldrMx step final project $ D.toStreamD m

146

147{-# INLINE foldr #-}

148foldr :: (Monad m, IsStream t) => (a -> b -> b) -> b -> t m a -> m b

149foldr f z = foldrM (\a b -> f a <$> b) (return z)

150

151-- | Like 'foldlx'', but with a monadic step function.

152--

153-- @since 0.7.0

154{-# INLINE foldlMx' #-}

155foldlMx' :: (IsStream t, Monad m)

156 => (x -> a -> m x) -> m x -> (x -> m b) -> t m a -> m b

157foldlMx' step begin done m = S.foldlMx' step begin done $ toStreamS m

158

159{-# INLINE parselMx' #-}

160parselMx'

161 :: (IsStream t, MonadThrow m)

162 => (s -> a -> m (Step s b))

163 -> m s

164 -> (s -> m b)

165 -> t m a

166 -> m b

167parselMx' step initial extract m =

168 D.parselMx' step initial extract $ D.toStreamD m

169

170-- | Strict left fold with an extraction function. Like the standard strict

171-- left fold, but applies a user supplied extraction function (the third

172-- argument) to the folded value at the end. This is designed to work with the

173-- @foldl@ library. The suffix @x@ is a mnemonic for extraction.

174--

175-- @since 0.7.0

176{-# INLINE foldlx' #-}

177foldlx' :: (IsStream t, Monad m)

178 => (x -> a -> x) -> x -> (x -> b) -> t m a -> m b

179foldlx' step begin done m = S.foldlx' step begin done $ toStreamS m

180

181-- | Strict left associative fold.

182--

183-- @since 0.2.0

184{-# INLINE foldl' #-}

185foldl' :: (Monad m, IsStream t) => (b -> a -> b) -> b -> t m a -> m b

186foldl' step begin m = S.foldl' step begin $ toStreamS m

187

188{-# INLINE foldlS #-}

189foldlS :: IsStream t => (t m b -> a -> t m b) -> t m b -> t m a -> t m b

190foldlS = K.foldlS

191

192-- | Lazy left fold to a transformer monad.

193--

194-- For example, to reverse a stream:

195--

196-- > S.toList $ S.foldlT (flip S.cons) S.nil $ (S.fromList [1..5] :: SerialT IO Int)

197--

198{-# INLINE foldlT #-}

199foldlT :: (Monad m, IsStream t, Monad (s m), MonadTrans s)

200 => (s m b -> a -> s m b) -> s m b -> t m a -> s m b

201foldlT f z s = S.foldlT f z (toStreamS s)

202

203{-# INLINE runFold #-}

204runFold :: (Monad m, IsStream t) => Fold m a b -> t m a -> m b

205runFold (Fold step begin done) = foldlMx' step begin done

206

207------------------------------------------------------------------------------

208-- Scans

209------------------------------------------------------------------------------

210

211-- postscanlM' followed by mapM

212{-# INLINE postscanlMx' #-}

213postscanlMx' :: (IsStream t, Monad m)

214 => (x -> a -> m x) -> m x -> (x -> m b) -> t m a -> t m b

215postscanlMx' step begin done m =

216 D.fromStreamD $ D.postscanlMx' step begin done $ D.toStreamD m

217

218-- postscanl' followed by map

219{-# INLINE postscanlx' #-}

220postscanlx' :: (IsStream t, Monad m)

221 => (x -> a -> x) -> x -> (x -> b) -> t m a -> t m b

222postscanlx' step begin done m =

223 D.fromStreamD $ D.postscanlx' step begin done $ D.toStreamD m

224

225-- scanlM' followed by mapM

226--

227{-# INLINE scanlMx' #-}

228scanlMx' :: (IsStream t, Monad m)

229 => (x -> a -> m x) -> m x -> (x -> m b) -> t m a -> t m b

230scanlMx' step begin done m =

231 D.fromStreamD $ D.scanlMx' step begin done $ D.toStreamD m

232

233-- scanl followed by map

234--

235-- | Strict left scan with an extraction function. Like 'scanl'', but applies a

236-- user supplied extraction function (the third argument) at each step. This is

237-- designed to work with the @foldl@ library. The suffix @x@ is a mnemonic for

238-- extraction.

239--

240-- @since 0.7.0

241{-# INLINE scanlx' #-}

242scanlx' :: (IsStream t, Monad m)

243 => (x -> a -> x) -> x -> (x -> b) -> t m a -> t m b

244scanlx' step begin done m =

245 fromStreamS $ S.scanlx' step begin done $ toStreamS m

246

247------------------------------------------------------------------------------

248-- Comparison

249------------------------------------------------------------------------------

250

251-- | Compare two streams for equality

252--

253-- @since 0.5.3

254{-# INLINE eqBy #-}

255eqBy :: (IsStream t, Monad m) => (a -> b -> Bool) -> t m a -> t m b -> m Bool

256eqBy f m1 m2 = D.eqBy f (D.toStreamD m1) (D.toStreamD m2)

257

258-- | Compare two streams

259--

260-- @since 0.5.3

261{-# INLINE cmpBy #-}

262cmpBy

263 :: (IsStream t, Monad m)

264 => (a -> b -> Ordering) -> t m a -> t m b -> m Ordering

265cmpBy f m1 m2 = D.cmpBy f (D.toStreamD m1) (D.toStreamD m2)

266

267{-# INLINE minimum #-}

268minimum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)

269minimum m = S.minimum (toStreamS m)

270

271{-# INLINE maximum #-}

272maximum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)

273maximum m = S.maximum (toStreamS m)

274

275------------------------------------------------------------------------------

276-- Fold Utilities

277------------------------------------------------------------------------------

278

279{-

280-- XXX do we have facilities in Foldable to fold any Foldable in this manner?

281--

282-- | Perform a pair wise bottom up hierarchical fold of elements in the

283-- container using the given function as the merge function.

284--

285-- This will perform a balanced merge sort if the merge function is

286-- 'mergeBy compare'.

287--

288-- @since 0.7.0

289{-# INLINABLE foldbWith #-}

290foldbWith :: IsStream t

291 => (t m a -> t m a -> t m a) -> SerialT Identity (t m a) -> t m a

292foldbWith f = K.foldb f K.nil

293-}

294

295-- /Since: 0.7.0 ("Streamly.Prelude")/

296--

297-- | A variant of 'Data.Foldable.fold' that allows you to fold a 'Foldable'

298-- container of streams using the specified stream sum operation.

299--

300-- @foldWith 'async' $ map return [1..3]@

301--

302-- Equivalent to:

303--

304-- @

305-- foldWith f = S.foldMapWith f id

306-- @

307--

308-- /Since: 0.1.0 ("Streamly")/

309{-# INLINABLE foldWith #-}

310foldWith :: (IsStream t, Foldable f)

311 => (t m a -> t m a -> t m a) -> f (t m a) -> t m a

312foldWith f = Prelude.foldr f K.nil

313

314-- /Since: 0.7.0 ("Streamly.Prelude")/

315--

316-- | A variant of 'foldMap' that allows you to map a monadic streaming action

317-- on a 'Foldable' container and then fold it using the specified stream merge

318-- operation.

319--

320-- @foldMapWith 'async' return [1..3]@

321--

322-- Equivalent to:

323--

324-- @

325-- foldMapWith f g xs = S.concatMapWith f g (S.fromFoldable xs)

326-- @

327--

328-- /Since: 0.1.0 ("Streamly")/

329{-# INLINABLE foldMapWith #-}

330foldMapWith :: (IsStream t, Foldable f)

331 => (t m b -> t m b -> t m b) -> (a -> t m b) -> f a -> t m b

332foldMapWith f g = Prelude.foldr (f . g) K.nil

333

334-- /Since: 0.7.0 ("Streamly.Prelude")/

335--

336-- | Like 'foldMapWith' but with the last two arguments reversed i.e. the

337-- monadic streaming function is the last argument.

338--

339-- Equivalent to:

340--

341-- @

342-- forEachWith = flip S.foldMapWith

343-- @

344--

345-- /Since: 0.1.0 ("Streamly")/

346{-# INLINABLE forEachWith #-}

347forEachWith :: (IsStream t, Foldable f)

348 => (t m b -> t m b -> t m b) -> f a -> (a -> t m b) -> t m b

349forEachWith f xs g = Prelude.foldr (f . g) K.nil xs