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

git.iko.soy/streamly

on 3947fee help sign in

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

331 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.Fold.Types as FL

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

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

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

26import qualified Streamly.Prelude as S

28import Gauge

29import Streamly.Benchmark.Common

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

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

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

35{-# INLINE sourceUnfoldrM #-}

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

37sourceUnfoldrM value n = S.unfoldrM step n

38 where

39 step cnt =

40 if cnt > n + value

41 then return Nothing

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

44{-# INLINE source #-}

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

46source = sourceUnfoldrM

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

49{-# INLINE benchIOSink #-}

50benchIOSink

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

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

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

55-------------------------------------------------------------------------------

56-- Folds

57-------------------------------------------------------------------------------

59{-# INLINE any #-}

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

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

63{-# INLINE all #-}

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

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

67{-# INLINE takeLE #-}

68takeLE :: Monad m => Int -> SerialT m a -> m ()

69takeLE value = IP.fold (FL.takeLE value FL.drain)

71{-# INLINE sequence_ #-}

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

73sequence_ value = foldr f (pure ()) (Prelude.replicate value (FL.takeLE 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 sliceSepBy #-}

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

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

88{-# INLINE many #-}

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

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

92{-# INLINE splitAllAny #-}

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

94splitAllAny value =

95 IP.fold

96 (FL.splitWith (,)

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 teeApplicative #-}

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

123teeApplicative value =

124 IP.fold ((,) <$> FL.all (<= value) <*> FL.any (> value))

125

126{-# INLINE distribute #-}

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

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

129

130-------------------------------------------------------------------------------

131-- Partitioning

132-------------------------------------------------------------------------------

133

134{-# INLINE partition #-}

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

136partition =

137 let f a =

138 if odd a

139 then Left a

140 else Right a

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

142

143{-# INLINE demuxWith #-}

144demuxWith ::

145 (Monad m, Ord k)

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

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

148 -> SerialT m a

149 -> m (Map k b)

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

151

152{-# INLINE demuxDefaultWith #-}

153demuxDefaultWith ::

154 (Monad m, Ord k, Num b)

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

156 -> Map k (Fold m b b)

157 -> SerialT m a

158 -> m (Map k b, b)

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

160

161{-# INLINE classifyWith #-}

162classifyWith ::

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

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

165

166-------------------------------------------------------------------------------

167-- unzip

168-------------------------------------------------------------------------------

169

170{-# INLINE unzip #-}

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

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

173

174-------------------------------------------------------------------------------

175-- Benchmarks

176-------------------------------------------------------------------------------

177

178moduleName :: String

179moduleName = "Data.Fold"

180

181o_1_space_serial_elimination :: Int -> [Benchmark]

182o_1_space_serial_elimination value =

183 [ bgroup "elimination"

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

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

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

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

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

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

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

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

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

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

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

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

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

197 , benchIOSink

198 value

199 "mean"

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

201 , benchIOSink

202 value

203 "variance"

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

205 , benchIOSink

206 value

207 "stdDev"

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

209 , benchIOSink

210 value

211 "mconcat"

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

213 , benchIOSink

214 value

215 "foldMap"

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

217 , benchIOSink

218 value

219 "foldMapM"

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

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

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

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

224 , benchIOSink

225 value

226 "lookup"

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

228 , benchIOSink

229 value

230 "findIndex"

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

232 , benchIOSink

233 value

234 "elemIndex"

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

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

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

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

239 , benchIOSink value "all" $ all value

240 , benchIOSink value "any" $ any value

241 , benchIOSink value "takeLE" $ takeLE value

242 , benchIOSink value "sliceSepBy" $ sliceSepBy value

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

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

245 ]

246 ]

247

248o_1_space_serial_transformation :: Int -> [Benchmark]

249o_1_space_serial_transformation value =

250 [ bgroup "transformation"

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

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

253 fld = FL.mapMaybe f FL.drain

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

255 , benchIOSink

256 value

257 "rsequence"

258 (S.fold (FL.rsequence (return <$> FL.drain)))

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

260 , benchIOSink

261 value

262 "pipe-mapM"

263 (S.fold

264 (FL.transform

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

266 FL.drain))

267 ]

268 ]

269

270o_1_space_serial_composition :: Int -> [Benchmark]

271o_1_space_serial_composition value =

272 [ bgroup

273 "composition"

274 [ benchIOSink value "splitWith (all, any)" $ splitAllAny value

275 , benchIOSink value "teeApplicative (all, any)"

276 $ teeApplicative value

277 , benchIOSink value "many drain (takeLE 1)" many

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

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

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

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

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

283 $ demuxDefaultWith fn mp

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

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

286 ]

287 ]

288

289 where

290

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

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

293

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

295

296o_n_space_serial :: Int -> [Benchmark]

297o_n_space_serial value =

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

299 ]

300

301o_n_heap_serial :: Int -> [Benchmark]

302o_n_heap_serial value =

303 [ bgroup "elimination"

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

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

306 [

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

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

309 ]

310 ]

311

312-------------------------------------------------------------------------------

313-- Driver

314-------------------------------------------------------------------------------

315

316main :: IO ()

317main = do

318 (value, cfg, benches) <- parseCLIOpts defaultStreamSize

319 value `seq` runMode (mode cfg) cfg benches (allBenchmarks value)

320

321 where

322

323 allBenchmarks value =

324 [ bgroup (o_1_space_prefix moduleName) $ concat

325 [ o_1_space_serial_elimination value

326 , o_1_space_serial_transformation value

327 , o_1_space_serial_composition value

328 ]

329 , bgroup (o_n_space_prefix moduleName) (o_n_space_serial value)

330 , bgroup (o_n_heap_prefix moduleName) (o_n_heap_serial value)

331 ]