content: {-# LANGUAGE OverloadedStrings  #-}
  {-# LANGUAGE RecordWildCards    #-}
  
  module Main where
  
  import Control.Applicative
  import Data.Monoid ((<>))
  import Data.Text (Text)
  import Turtle (ExitCode(..), Parser, ShellFailed(..), s, (%))
  
  import qualified Control.Exception
  import qualified Criterion
  import qualified Criterion.Main         as Criterion
  import qualified Criterion.Main.Options as Criterion
  import qualified Data.Text              as Text
  import qualified Data.Version
  import qualified Options.Applicative
  import qualified Paths_bench
  import qualified System.IO              as IO
  import qualified System.IO.Silently     as Silently
  import qualified System.Process
  import qualified Turtle
  
  version :: Text
  version = Text.pack (Data.Version.showVersion Paths_bench.version)
  
  data Options
      = Options (Maybe Text) (Maybe Text) [Text] Criterion.Mode
      | Version deriving (Show)
  
  before :: Parser Text
  before =
      Options.Applicative.strOption
          (   Options.Applicative.long "before"
          <>  Options.Applicative.help "Specify a command to run before each run of the benchmark"
          <>  Options.Applicative.metavar "command"
          )
  
  after :: Parser Text
  after =
      Options.Applicative.strOption
          (   Options.Applicative.long "after"
          <>  Options.Applicative.help "Specify a command to run after each run of the benchmark"
          <>  Options.Applicative.metavar "command"
          )
  
  parser :: Parser Options
  parser =
              Version
          <$  Options.Applicative.flag'
                  ()
                  (   Options.Applicative.long "version"
                  <>  Options.Applicative.help "Version number"
                  )
      <|>     Options
          <$> optional before
          <*> optional after
          <*> some
                  (Turtle.argText "command(s)" "The command line(s) to benchmark")
          <*> Criterion.parseWith Criterion.defaultConfig
  
  main :: IO ()
  main = do
      x <- Turtle.options "Command-line tool to benchmark other programs" parser
      case x of
          Options maybeBefore maybeAfter commands mode -> do
              let benches = map (buildBench maybeBefore maybeAfter) commands
  
              let benches' = case commands of
                      [_] -> benches
                      _   -> [ Criterion.bgroup "bench" benches ]
  
              Criterion.runMode mode benches'
          Version -> do
              Turtle.printf ("bench version "%s%"\n") version
  
  buildBench :: Maybe Text -> Maybe Text -> Text -> Criterion.Benchmark
  buildBench maybeBefore maybeAfter command = do
      let io cmd = Silently.hSilence [IO.stdout, IO.stderr] $ do
              let createProcess =
                      (System.Process.shell (Text.unpack cmd))
                          { System.Process.std_in  = System.Process.NoStream
                          }
  
              exitCode <- Turtle.system createProcess empty
              case exitCode of
                  ExitFailure _ -> do
                      Control.Exception.throwIO (ShellFailed cmd exitCode)
                  _  -> do
                      return ()
  
      let benchmarkable = case (maybeBefore, maybeAfter) of
              (Just before_, Just after_) ->
                  Criterion.perRunEnvWithCleanup
                      (io before_)
                      (\_ -> io after_)
                      (\_ -> io command)
              (Just before_, Nothing) ->
                  Criterion.perRunEnvWithCleanup
                      (io before_)
                      (\_ -> return ())
                      (\_ -> io command)
              (Nothing, Just after_) ->
                  Criterion.perRunEnvWithCleanup
                      (return ())
                      (\_ -> io after_)
                      (\_ -> io command)
              (Nothing, Nothing) ->
                  Criterion.nfIO (io command)
  
      Criterion.bench (Text.unpack command) benchmarkable
  