git.iko.soy/streamly
on f5c55df help sign in
streamly/benchmark/Common.hs
95 lines 3.1 kB history raw
1-- |
2-- Module : Main
3-- Copyright : (c) 2019 Composewell Technologies
4--
5-- License : BSD3
6-- Maintainer : streamly@composewell.com
8module Common (parseCLIOpts) where
10import Control.Exception (evaluate)
11import Control.Monad (when)
12import Data.List (scanl')
13import Data.Maybe (catMaybes)
14import System.Console.GetOpt
15 (OptDescr(..), ArgDescr(..), ArgOrder(..), getOpt')
16import System.Environment (getArgs, lookupEnv, setEnv)
17import Text.Read (readMaybe)
19import Gauge
21-------------------------------------------------------------------------------
22-- Parse custom CLI options
23-------------------------------------------------------------------------------
25data BenchOpts = StreamSize Int deriving Show
27getStreamSize :: String -> Int
28getStreamSize size =
29 case (readMaybe size :: Maybe Int) of
30 Just x -> x
31 Nothing -> error "Stream size must be numeric"
33options :: [OptDescr BenchOpts]
34options =
35 [
36 Option [] ["stream-size"] (ReqArg getSize "COUNT") "Stream element count"
37 ]
39 where
41 getSize = StreamSize . getStreamSize
43deleteOptArgs
44 :: (Maybe String, Maybe String) -- (prev, yielded)
45 -> String
46 -> (Maybe String, Maybe String)
47deleteOptArgs (Nothing, _) opt =
48 if opt == "--stream-size"
49 then (Just opt, Nothing)
50 else (Just opt, Just opt)
52deleteOptArgs (Just prev, _) opt =
53 if opt == "--stream-size" || prev == "--stream-size"
54 then (Just opt, Nothing)
55 else (Just opt, Just opt)
57parseCLIOpts :: Int -> IO (Int, Config, [String])
58parseCLIOpts defaultStreamSize = do
59 args <- getArgs
61 -- Parse custom options
62 let (opts, _, _, errs) = getOpt' Permute options args
63 when (not $ null errs) $ error $ concat errs
64 (streamSize, args') <-
65 case opts of
66 StreamSize x : _ -> do
67 -- When using the gauge "--measure-with" option we need to make
68 -- sure that we pass the stream size to child process forked by
69 -- gauge. So we use this env var for that purpose.
70 setEnv "STREAM_SIZE" (show x)
71 -- Hack! remove the option and its argument from args
72 -- getOpt should have a way to return the unconsumed args in
73 -- correct order.
74 newArgs <-
75 evaluate
76 $ catMaybes
77 $ map snd
78 $ scanl' deleteOptArgs (Nothing, Nothing) args
79 return (x, newArgs)
80 _ -> do
81 r <- lookupEnv "STREAM_SIZE"
82 case r of
83 Just x -> do
84 s <- evaluate $ getStreamSize x
85 return (s, args)
86 Nothing -> return (defaultStreamSize, args)
88 -- Parse gauge options
89 let config = defaultConfig
90 { timeLimit = Just 1
91 , minDuration = 0
92 , includeFirstIter = streamSize > defaultStreamSize
93 }
94 let (cfg, benches) = parseWith config args'
95 streamSize `seq` return (streamSize, cfg, benches)