benchmark/Streamly/Benchmark/Data/Fold.hs · streamly

git.iko.soy/streamly

on 3feeec2 help sign in

streamly/benchmark/Streamly/Benchmark/Data/Fold.hs

334 lines 12 kB history raw

1-- |

2-- Module : Streamly.Benchmark.Data.Fold

4--

5-- License : MIT

6-- Maintainer : streamly@composewell.com

8{-# LANGUAGE FlexibleContexts #-}

10module Main (main) where

12import Control.DeepSeq (NFData(..))

13import Data.Map.Strict (Map)

14import Data.Monoid (Last(..), Sum(..))

15import System.Random (randomRIO)

17import Streamly.Prelude (SerialT)

18import Streamly.Internal.Data.Fold (Fold(..))

20import qualified Data.Map.Strict as Map

21import qualified Streamly.Internal.Data.Fold as FL

22import qualified Streamly.Internal.Data.Pipe as Pipe

23import qualified Streamly.Internal.Data.Sink as Sink

24import qualified Streamly.Internal.Data.Stream.IsStream as IP

25import qualified Streamly.Prelude as S

27import Gauge

28import Streamly.Benchmark.Common

29import Prelude hiding (all, any, take, unzip, sequence_)

31-- We need a monadic bind here to make sure that the function f does not get

32-- completely optimized out by the compiler in some cases.

34{-# INLINE sourceUnfoldrM #-}

35sourceUnfoldrM :: (S.IsStream t, S.MonadAsync m) => Int -> Int -> t m Int

36sourceUnfoldrM value n = S.unfoldrM step n

37 where

38 step cnt =

39 if cnt > n + value

40 then return Nothing

41 else return (Just (cnt, cnt + 1))

43{-# INLINE source #-}

44source :: (S.MonadAsync m, S.IsStream t) => Int -> Int -> t m Int

45source = sourceUnfoldrM

47-- | Takes a fold method, and uses it with a default source.

48{-# INLINE benchIOSink #-}

49benchIOSink

50 :: (S.IsStream t, NFData b)

51 => Int -> String -> (t IO Int -> IO b) -> Benchmark

52benchIOSink value name f = bench name $ nfIO $ randomRIO (1,1) >>= f . source value

54-------------------------------------------------------------------------------

55-- Folds

56-------------------------------------------------------------------------------

58{-# INLINE any #-}

59any :: (Monad m, Ord a) => a -> SerialT m a -> m Bool

60any value = IP.fold (FL.any (> value))

62{-# INLINE all #-}

63all :: (Monad m, Ord a) => a -> SerialT m a -> m Bool

64all value = IP.fold (FL.all (<= value))

66{-# INLINE take #-}

67take :: Monad m => Int -> SerialT m a -> m ()

68take value = IP.fold (FL.take value FL.drain)

70{-# INLINE sequence_ #-}

71sequence_ :: Monad m => Int -> Fold m a ()

72sequence_ value =

73 foldr f (FL.fromPure ()) (Prelude.replicate value (FL.take 1 FL.drain))

75 where

77 {-# INLINE f #-}

78 f m k = FL.concatMap (const k) m

80-------------------------------------------------------------------------------

81-- Splitting by serial application

82-------------------------------------------------------------------------------

84{-# INLINE takeEndBy_ #-}

85takeEndBy_ :: Monad m => Int -> SerialT m Int -> m ()

86takeEndBy_ value = IP.fold (FL.takeEndBy_ (>= value) FL.drain)

88{-# INLINE many #-}

89many :: Monad m => SerialT m Int -> m ()

90many = IP.fold (FL.many (FL.take 1 FL.drain) FL.drain)

92{-# INLINE splitAllAny #-}

93splitAllAny :: Monad m => Int -> SerialT m Int -> m (Bool, Bool)

94splitAllAny value =

95 IP.fold

96 (FL.serialWith (,)

97 (FL.all (<= (value `div` 2)))

98 (FL.any (> value))

99 )

100

101{-

102{-# INLINE split_ #-}

103split_ :: Monad m

104 => Int -> SerialT m Int -> m ()

105split_ value =

106 IP.fold

107 (FL.split_

108 (FL.all (<= (value `div` 2)))

109 (FL.any (> value))

110 )

111-}

112

113-------------------------------------------------------------------------------

114-- Distributing by parallel application

115-------------------------------------------------------------------------------

116

117{-# INLINE teeSumLength #-}

118teeSumLength :: Monad m => SerialT m Int -> m (Int, Int)

119teeSumLength = IP.fold (FL.teeWith (,) FL.sum FL.length)

120

121{-# INLINE teeAllAny #-}

122teeAllAny :: (Monad m, Ord a) => a -> SerialT m a -> m (Bool, Bool)

123teeAllAny value = IP.fold (FL.teeWith (,) all_ any_)

124

125 where

126

127 all_ = FL.all (<= value)

128 any_ = FL.any (> value)

129

130{-# INLINE distribute #-}

131distribute :: Monad m => SerialT m Int -> m [Int]

132distribute = IP.fold (FL.distribute [FL.sum, FL.length])

133

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

135-- Partitioning

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

137

138{-# INLINE partition #-}

139partition :: Monad m => SerialT m Int -> m (Int, Int)

140partition =

141 let f a =

142 if odd a

143 then Left a

144 else Right a

145 in IP.fold $ FL.map f (FL.partition FL.sum FL.length)

146

147{-# INLINE demuxWith #-}

148demuxWith ::

149 (Monad m, Ord k)

150 => (a -> (k, a'))

151 -> Map k (Fold m a' b)

152 -> SerialT m a

153 -> m (Map k b)

154demuxWith f mp = S.fold (FL.demuxWith f mp)

155

156{-# INLINE demuxDefaultWith #-}

157demuxDefaultWith ::

158 (Monad m, Ord k, Num b)

159 => (a -> (k, b))

160 -> Map k (Fold m b b)

161 -> SerialT m a

162 -> m (Map k b, b)

163demuxDefaultWith f mp = S.fold (FL.demuxDefaultWith f mp (FL.map snd FL.sum))

164

165{-# INLINE classifyWith #-}

166classifyWith ::

167 (Monad m, Ord k, Num a) => (a -> k) -> SerialT m a -> m (Map k a)

168classifyWith f = S.fold (FL.classifyWith f FL.sum)

169

170-------------------------------------------------------------------------------

171-- unzip

172-------------------------------------------------------------------------------

173

174{-# INLINE unzip #-}

175unzip :: Monad m => SerialT m Int -> m (Int, Int)

176unzip = IP.fold $ FL.map (\a -> (a, a)) (FL.unzip FL.sum FL.length)

177

178-------------------------------------------------------------------------------

179-- Benchmarks

180-------------------------------------------------------------------------------

181

182moduleName :: String

183moduleName = "Data.Fold"

184

185o_1_space_serial_elimination :: Int -> [Benchmark]

186o_1_space_serial_elimination value =

187 [ bgroup "elimination"

188 [ benchIOSink value "drain" (S.fold FL.drain)

189 , benchIOSink value "drainBy" (S.fold (FL.drainBy return))

190 , benchIOSink value "drainN" (S.fold (FL.drainN value))

191 , benchIOSink value "sink" (S.fold $ Sink.toFold Sink.drain)

192 , benchIOSink value "last" (S.fold FL.last)

193 , benchIOSink value "length" (S.fold FL.length)

194 , benchIOSink value "sum" (S.fold FL.sum)

195 , benchIOSink value "sum (foldMap)" (S.fold (FL.foldMap Sum))

196 , benchIOSink value "product" (S.fold FL.product)

197 , benchIOSink value "maximumBy" (S.fold (FL.maximumBy compare))

198 , benchIOSink value "maximum" (S.fold FL.maximum)

199 , benchIOSink value "minimumBy" (S.fold (FL.minimumBy compare))

200 , benchIOSink value "minimum" (S.fold FL.minimum)

201 , benchIOSink

202 value

203 "mean"

204 (S.fold FL.mean . S.map (fromIntegral :: Int -> Double))

205 , benchIOSink

206 value

207 "variance"

208 (S.fold FL.variance . S.map (fromIntegral :: Int -> Double))

209 , benchIOSink

210 value

211 "stdDev"

212 (S.fold FL.stdDev . S.map (fromIntegral :: Int -> Double))

213 , benchIOSink

214 value

215 "mconcat"

216 (S.fold FL.mconcat . S.map (Last . Just))

217 , benchIOSink

218 value

219 "foldMap"

220 (S.fold (FL.foldMap (Last . Just)))

221 , benchIOSink

222 value

223 "foldMapM"

224 (S.fold (FL.foldMapM (return . Last . Just)))

225 , benchIOSink value "index" (S.fold (FL.index (value + 1)))

226 , benchIOSink value "head" (S.fold FL.head)

227 , benchIOSink value "find" (S.fold (FL.find (== (value + 1))))

228 , benchIOSink

229 value

230 "lookup"

231 (S.fold (FL.map (\a -> (a, a)) (FL.lookup (value + 1))))

232 , benchIOSink

233 value

234 "findIndex"

235 (S.fold (FL.findIndex (== (value + 1))))

236 , benchIOSink

237 value

238 "elemIndex"

239 (S.fold (FL.elemIndex (value + 1)))

240 , benchIOSink value "null" (S.fold FL.null)

241 , benchIOSink value "elem" (S.fold (FL.elem (value + 1)))

242 , benchIOSink value "notElem" (S.fold (FL.notElem (value + 1)))

243 , benchIOSink value "all" $ all value

244 , benchIOSink value "any" $ any value

245 , benchIOSink value "take" $ take value

246 , benchIOSink value "takeEndBy_" $ takeEndBy_ value

247 , benchIOSink value "and" (S.fold FL.and . S.map (<= (value + 1)))

248 , benchIOSink value "or" (S.fold FL.or . S.map (> (value + 1)))

249 ]

250 ]

251

252o_1_space_serial_transformation :: Int -> [Benchmark]

253o_1_space_serial_transformation value =

254 [ bgroup "transformation"

255 [ benchIOSink value "map" (S.fold (FL.map (+ 1) FL.drain))

256 , let f x = if even x then Just x else Nothing

257 fld = FL.mapMaybe f FL.drain

258 in benchIOSink value "mapMaybe" (S.fold fld)

259 , benchIOSink

260 value

261 "rsequence"

262 (S.fold (FL.rmapM id (return <$> FL.drain)))

263 , benchIOSink value "rmapM" (S.fold (FL.rmapM return FL.drain))

264 , benchIOSink

265 value

266 "pipe-mapM"

267 (S.fold

268 (FL.transform

269 (Pipe.mapM (\x -> return $ x + 1))

270 FL.drain))

271 ]

272 ]

273

274o_1_space_serial_composition :: Int -> [Benchmark]

275o_1_space_serial_composition value =

276 [ bgroup

277 "composition"

278 [ benchIOSink value "serialWith (all, any)" $ splitAllAny value

279 , benchIOSink value "tee (all, any)" $ teeAllAny value

280 , benchIOSink value "many drain (take 1)" many

281 , benchIOSink value "tee (sum, length)" teeSumLength

282 , benchIOSink value "distribute [sum, length]" distribute

283 , benchIOSink value "partition (sum, length)" partition

284 , benchIOSink value "unzip (sum, length)" unzip

285 , benchIOSink value "demuxDefaultWith [sum, length] sum"

286 $ demuxDefaultWith fn mp

287 , benchIOSink value "demuxWith [sum, length]" $ demuxWith fn mp

288 , benchIOSink value "classifyWith sum" $ classifyWith (fst . fn)

289 ]

290 ]

291

292 where

293

294 -- We use three keys 0, 1, and 2. 0 and 1 are mapped and 3 is unmapped.

295 fn x = (x `mod` 3, x)

296

297 mp = Map.fromList [(0, FL.sum), (1, FL.length)]

298

299o_n_space_serial :: Int -> [Benchmark]

300o_n_space_serial value =

301 [ benchIOSink value "sequence_/100" $ S.fold (sequence_ (value `div` 100))

302 ]

303

304o_n_heap_serial :: Int -> [Benchmark]

305o_n_heap_serial value =

306 [ bgroup "elimination"

307 -- Left folds for building a structure are inherently non-streaming

308 -- as the structure cannot be lazily consumed until fully built.

309 [

310 benchIOSink value "toList" (S.fold FL.toList)

311 , benchIOSink value "toListRev" (S.fold FL.toListRev)

312 , benchIOSink value "toStream" (S.fold FL.toStream)

313 , benchIOSink value "toStreamRev" (S.fold FL.toStreamRev)

314 ]

315 ]

316

317-------------------------------------------------------------------------------

318-- Driver

319-------------------------------------------------------------------------------

320

321main :: IO ()

322main = runWithCLIOpts defaultStreamSize allBenchmarks

323

324 where

325

326 allBenchmarks value =

327 [ bgroup (o_1_space_prefix moduleName) $ concat

328 [ o_1_space_serial_elimination value

329 , o_1_space_serial_transformation value

330 , o_1_space_serial_composition value

331 ]

332 , bgroup (o_n_space_prefix moduleName) (o_n_space_serial value)

333 , bgroup (o_n_heap_prefix moduleName) (o_n_heap_serial value)

334 ]