src/Streamly/Fold/Types.hs · streamly

git.iko.soy/streamly

on b9e99db help sign in

streamly/src/Streamly/Fold/Types.hs

437 lines 13 kB history raw

1{-# LANGUAGE CPP #-}

2{-# LANGUAGE ExistentialQuantification #-}

3{-# LANGUAGE FlexibleContexts #-}

5-- |

6-- Module : Streamly.Fold.Types

9-- License : BSD3

10-- Maintainer : harendra.kumar@gmail.com

11-- Stability : experimental

12-- Portability : GHC

14module Streamly.Fold.Types

15 ( Fold (..)

16 , toListRevF -- experimental

17 , lmap

18 , lmapM

19 , lfilter

20 , lfilterM

21 , lcatMaybes

22 , ltake

23 , ltakeWhile

24 , lsessionsOf

25 , lchunksOf

27 , duplicate

28 , initialize

29 , runStep

30 )

31where

33import Control.Applicative (liftA2)

34import Control.Concurrent (threadDelay, forkIO, killThread)

35import Control.Concurrent.MVar (MVar, newMVar, takeMVar, putMVar)

36import Control.Exception (SomeException(..), catch, mask)

37import Control.Monad (void)

38import Control.Monad.Catch (throwM)

39import Control.Monad.IO.Class (MonadIO(..))

40import Control.Monad.Trans.Control (control)

41import Data.Maybe (isJust, fromJust)

42#if __GLASGOW_HASKELL__ < 808

43import Data.Semigroup (Semigroup(..))

44#endif

45import Streamly.Strict (Tuple'(..), Tuple3'(..), Either'(..))

46import Streamly.SVar (MonadAsync)

48------------------------------------------------------------------------------

49-- Monadic left folds

50------------------------------------------------------------------------------

52-- | Represents a left fold over an input stream of values of type @a@ to a

53-- single value of type @b@ in 'Monad' @m@.

54--

55-- @since 0.7.0

57-- The fold uses an intermediate type @x@ as accumulator. The fold accumulator

58-- is initialized by calling the @init@ function and is then driven by calling

59-- the step function repeatedly. When the fold is done the @extract@ function

60-- is used to map the intermediate type @x@ to the final type @b@. This allows

61-- the state of the fold to be embedded in an arbitrary type @x@.

62data Fold m a b =

63 -- | @Fold @ @ step @ @ initial @ @ extract@

64 forall x. Fold (x -> a -> m x) (m x) (x -> m b)

66-- | Maps a function on the output of the fold (the type @b@).

67instance Applicative m => Functor (Fold m a) where

68 {-# INLINE fmap #-}

69 fmap f (Fold step start done) = Fold step start done'

70 where

71 done' x = fmap f $! done x

73 {-# INLINE (<$) #-}

74 (<$) b = \_ -> pure b

76-- | The fold resulting from '<*>' distributes its input to both the argument

77-- folds and combines their output using the supplied function.

78instance Applicative m => Applicative (Fold m a) where

79 {-# INLINE pure #-}

80 pure b = Fold (\() _ -> pure ()) (pure ()) (\() -> pure b)

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

83 (Fold stepL beginL doneL) <*> (Fold stepR beginR doneR) =

84 let step (Tuple' xL xR) a = Tuple' <$> stepL xL a <*> stepR xR a

85 begin = Tuple' <$> beginL <*> beginR

86 done (Tuple' xL xR) = doneL xL <*> doneR xR

87 in Fold step begin done

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

90 (<*) m = \_ -> m

92 {-# INLINE (*>) #-}

93 _ *> m = m

95-- | Combines the outputs of the folds (the type @b@) using their 'Semigroup'

96-- instances.

97instance (Semigroup b, Monad m) => Semigroup (Fold m a b) where

98 {-# INLINE (<>) #-}

99 (<>) = liftA2 (<>)

100

101-- | Combines the outputs of the folds (the type @b@) using their 'Monoid'

102-- instances.

103instance (Semigroup b, Monoid b, Monad m) => Monoid (Fold m a b) where

104 {-# INLINE mempty #-}

105 mempty = pure mempty

106

107 {-# INLINE mappend #-}

108 mappend = (<>)

109

110-- | Combines the fold outputs (type @b@) using their 'Num' instances.

111instance (Monad m, Num b) => Num (Fold m a b) where

112 {-# INLINE fromInteger #-}

113 fromInteger = pure . fromInteger

114

115 {-# INLINE negate #-}

116 negate = fmap negate

117

118 {-# INLINE abs #-}

119 abs = fmap abs

120

121 {-# INLINE signum #-}

122 signum = fmap signum

123

124 {-# INLINE (+) #-}

125 (+) = liftA2 (+)

126

127 {-# INLINE (*) #-}

128 (*) = liftA2 (*)

129

130 {-# INLINE (-) #-}

131 (-) = liftA2 (-)

132

133-- | Combines the fold outputs (type @b@) using their 'Fractional' instances.

134instance (Monad m, Fractional b) => Fractional (Fold m a b) where

135 {-# INLINE fromRational #-}

136 fromRational = pure . fromRational

137

138 {-# INLINE recip #-}

139 recip = fmap recip

140

141 {-# INLINE (/) #-}

142 (/) = liftA2 (/)

143

144-- | Combines the fold outputs using their 'Floating' instances.

145instance (Monad m, Floating b) => Floating (Fold m a b) where

146 {-# INLINE pi #-}

147 pi = pure pi

148

149 {-# INLINE exp #-}

150 exp = fmap exp

151

152 {-# INLINE sqrt #-}

153 sqrt = fmap sqrt

154

155 {-# INLINE log #-}

156 log = fmap log

157

158 {-# INLINE sin #-}

159 sin = fmap sin

160

161 {-# INLINE tan #-}

162 tan = fmap tan

163

164 {-# INLINE cos #-}

165 cos = fmap cos

166

167 {-# INLINE asin #-}

168 asin = fmap asin

169

170 {-# INLINE atan #-}

171 atan = fmap atan

172

173 {-# INLINE acos #-}

174 acos = fmap acos

175

176 {-# INLINE sinh #-}

177 sinh = fmap sinh

178

179 {-# INLINE tanh #-}

180 tanh = fmap tanh

181

182 {-# INLINE cosh #-}

183 cosh = fmap cosh

184

185 {-# INLINE asinh #-}

186 asinh = fmap asinh

187

188 {-# INLINE atanh #-}

189 atanh = fmap atanh

190

191 {-# INLINE acosh #-}

192 acosh = fmap acosh

193

194 {-# INLINE (**) #-}

195 (**) = liftA2 (**)

196

197 {-# INLINE logBase #-}

198 logBase = liftA2 logBase

199

200------------------------------------------------------------------------------

201-- Internal APIs

202------------------------------------------------------------------------------

203

204-- This is more efficient than 'toList'. toList is exactly the same as

205-- reversing the list after toListRev.

206--

207-- | Buffers the input stream to a list in the reverse order of the input.

208--

209-- /Warning!/ working on large lists accumulated as buffers in memory could be

210-- very inefficient, consider using "Streamly.Array" instead.

211--

212-- @since 0.7.0

213

214-- xn : ... : x2 : x1 : []

215{-# INLINABLE toListRevF #-}

216toListRevF :: Monad m => Fold m a [a]

217toListRevF = Fold (\xs x -> return $ x:xs) (return []) return

218

219-- | @(lmap f fold)@ maps the function @f@ on the input of the fold.

220--

221-- >>> S.runFold (FL.lmap (\x -> x * x) FL.sum) (S.enumerateFromTo 1 100)

222-- 338350

223--

224-- @since 0.7.0

225{-# INLINABLE lmap #-}

226lmap :: (a -> b) -> Fold m b r -> Fold m a r

227lmap f (Fold step begin done) = Fold step' begin done

228 where

229 step' x a = step x (f a)

230

231-- | @(lmapM f fold)@ maps the monadic function @f@ on the input of the fold.

232--

233-- @since 0.7.0

234{-# INLINABLE lmapM #-}

235lmapM :: Monad m => (a -> m b) -> Fold m b r -> Fold m a r

236lmapM f (Fold step begin done) = Fold step' begin done

237 where

238 step' x a = f a >>= step x

239

240------------------------------------------------------------------------------

241-- Filtering

242------------------------------------------------------------------------------

243

244-- | Include only those elements that pass a predicate.

245--

246-- >>> S.runFold (lfilter (> 5) FL.sum) [1..10]

247-- 40

248--

249-- @since 0.7.0

250{-# INLINABLE lfilter #-}

251lfilter :: Monad m => (a -> Bool) -> Fold m a r -> Fold m a r

252lfilter f (Fold step begin done) = Fold step' begin done

253 where

254 step' x a = if f a then step x a else return x

255

256-- | Like 'lfilter' but with a monadic predicate.

257--

258-- @since 0.7.0

259{-# INLINABLE lfilterM #-}

260lfilterM :: Monad m => (a -> m Bool) -> Fold m a r -> Fold m a r

261lfilterM f (Fold step begin done) = Fold step' begin done

262 where

263 step' x a = do

264 use <- f a

265 if use then step x a else return x

266

267-- | Transform a fold from a pure input to a 'Maybe' input, consuming only

268-- 'Just' values.

269{-# INLINE lcatMaybes #-}

270lcatMaybes :: Monad m => Fold m a b -> Fold m (Maybe a) b

271lcatMaybes = lfilter isJust . lmap fromJust

272

273-- | Take first 'n' elements from the stream and discard the rest.

274--

275-- @since 0.7.0

276{-# INLINABLE ltake #-}

277ltake :: Monad m => Int -> Fold m a b -> Fold m a b

278ltake n (Fold step initial done) = Fold step' initial' done'

279 where

280 initial' = fmap (Tuple' 0) initial

281 step' (Tuple' i r) a = do

282 if i < n

283 then do

284 res <- step r a

285 return $ Tuple' (i + 1) res

286 else return $ Tuple' i r

287 done' (Tuple' _ r) = done r

288

289-- | Takes elements from the input as long as the predicate succeeds.

290--

291-- @since 0.7.0

292{-# INLINABLE ltakeWhile #-}

293ltakeWhile :: Monad m => (a -> Bool) -> Fold m a b -> Fold m a b

294ltakeWhile predicate (Fold step initial done) = Fold step' initial' done'

295 where

296 initial' = fmap Left' initial

297 step' (Left' r) a = do

298 if predicate a

299 then fmap Left' $ step r a

300 else return (Right' r)

301 step' r _ = return r

302 done' (Left' r) = done r

303 done' (Right' r) = done r

304

305------------------------------------------------------------------------------

306-- Nesting

307------------------------------------------------------------------------------

308--

309-- | Modify the fold such that when the fold is done, instead of returning the

310-- accumulator, it returns a fold. The returned fold starts from where we left

311-- i.e. it uses the last accumulator value as the initial value of the

312-- accumulator. Thus we can resume the fold later and feed it more input.

313--

314-- >> do

315-- > more <- S.runFold (FL.duplicate FL.sum) (S.enumerateFromTo 1 10)

316-- > evenMore <- S.runFold (FL.duplicate more) (S.enumerateFromTo 11 20)

317-- > S.runFold evenMore (S.enumerateFromTo 21 30)

318-- > 465

319--

320-- @since 0.7.0

321{-# INLINABLE duplicate #-}

322duplicate :: Applicative m => Fold m a b -> Fold m a (Fold m a b)

323duplicate (Fold step begin done) =

324 Fold step begin (\x -> pure (Fold step (pure x) done))

325

326-- | Run the initialization effect of a fold. The returned fold would use the

327-- value returned by this effect as its initial value.

328--

329{-# INLINABLE initialize #-}

330initialize :: Monad m => Fold m a b -> m (Fold m a b)

331initialize (Fold step initial extract) = do

332 i <- initial

333 return $ Fold step (return i) extract

334

335-- | Run one step of a fold and store the accumulator as an initial value in

336-- the returned fold.

337{-# INLINABLE runStep #-}

338runStep :: Monad m => Fold m a b -> a -> m (Fold m a b)

339runStep (Fold step initial extract) a = do

340 i <- initial

341 r <- step i a

342 return $ (Fold step (return r) extract)

343

344-- | For every n input items, apply the first fold and supply the result to the

345-- next fold.

346--

347{-# INLINE lchunksOf #-}

348lchunksOf :: Monad m => Int -> Fold m a b -> Fold m b c -> Fold m a c

349lchunksOf n (Fold step1 initial1 extract1) (Fold step2 initial2 extract2) =

350 Fold step' initial' extract'

351

352 where

353

354 initial' = (Tuple3' 0) <$> initial1 <*> initial2

355 step' (Tuple3' i r1 r2) a = do

356 if i < n

357 then do

358 res <- step1 r1 a

359 return $ Tuple3' (i + 1) res r2

360 else do

361 res <- extract1 r1

362 acc2 <- step2 r2 res

363

364 i1 <- initial1

365 acc1 <- step1 i1 a

366 return $ Tuple3' 1 acc1 acc2

367 extract' (Tuple3' _ r1 r2) = do

368 res <- extract1 r1

369 acc2 <- step2 r2 res

370 extract2 acc2

371

372-- | Group the input stream into windows of n second each and then fold each

373-- group using the provided fold function.

374--

375-- For example, we can copy and distribute a stream to multiple folds where

376-- each fold can group the input differently e.g. by one second, one minute and

377-- one hour windows respectively and fold each resulting stream of folds.

378--

379-- @

380--

381-- -----Fold m a b----|-Fold n a c-|-Fold n a c-|-...-|----Fold m a c

382--

383-- @

384{-# INLINE lsessionsOf #-}

385lsessionsOf :: MonadAsync m => Double -> Fold m a b -> Fold m b c -> Fold m a c

386lsessionsOf n (Fold step1 initial1 extract1) (Fold step2 initial2 extract2) =

387 Fold step' initial' extract'

388

389 where

390

391 -- XXX MVar may be expensive we need a cheaper synch mechanism here

392 initial' = do

393 i1 <- initial1

394 i2 <- initial2

395 mv1 <- liftIO $ newMVar i1

396 mv2 <- liftIO $ newMVar (Right i2)

397 t <- control $ \run ->

398 mask $ \restore -> do

399 tid <- forkIO $ catch (restore $ void $ run (timerThread mv1 mv2))

400 (handleChildException mv2)

401 run (return tid)

402 return $ Tuple3' t mv1 mv2

403 step' acc@(Tuple3' _ mv1 _) a = do

404 r1 <- liftIO $ takeMVar mv1

405 res <- step1 r1 a

406 liftIO $ putMVar mv1 res

407 return acc

408 extract' (Tuple3' tid _ mv2) = do

409 r2 <- liftIO $ takeMVar mv2

410 liftIO $ killThread tid

411 case r2 of

412 Left e -> throwM e

413 Right x -> extract2 x

414

415 timerThread mv1 mv2 = do

416 liftIO $ threadDelay (round $ n * 1000000)

417

418 r1 <- liftIO $ takeMVar mv1

419 i1 <- initial1

420 liftIO $ putMVar mv1 i1

421

422 res1 <- extract1 r1

423 r2 <- liftIO $ takeMVar mv2

424 res <- case r2 of

425 Left _ -> return r2

426 Right x -> fmap Right $ step2 x res1

427 liftIO $ putMVar mv2 res

428 timerThread mv1 mv2

429

430 handleChildException ::

431 MVar (Either SomeException a) -> SomeException -> IO ()

432 handleChildException mv2 e = do

433 r2 <- takeMVar mv2

434 let r = case r2 of

435 Left _ -> r2

436 Right _ -> Left e

437 putMVar mv2 r