content: {-# LANGUAGE DeriveGeneric     #-}
  {-# LANGUAGE OverloadedStrings #-}
  
  {-# OPTIONS_GHC -fno-warn-orphans #-}
  
  -- | This benchmark is derived from the large-output benchmark in prettyprinter, but contains additional annotations.
  module Main (main) where
  
  import Prelude        ()
  import Prelude.Compat
  
  import           Control.DeepSeq
  import           Control.Monad.Compat
  import           Data.Char
  import           Data.Map                              (Map)
  import qualified Data.Map                              as M
  import           Data.Text                             (Text)
  import qualified Data.Text                             as T
  import qualified Data.Text.IO                          as T
  import qualified Data.Text.Lazy                        as TL
  import           GHC.Generics
  import           Prettyprinter
  import           Prettyprinter.Render.Terminal         as Terminal
  import qualified Prettyprinter.Render.Text             as Text
  import           Test.QuickCheck
  import           Test.QuickCheck.Gen
  import           Test.QuickCheck.Random
  import           Test.Tasty.Bench
  
  
  
  newtype Program = Program Binds deriving (Show, Generic)
  newtype Binds = Binds (Map Text LambdaForm) deriving (Show, Generic)
  data LambdaForm = LambdaForm ![Text] ![Text] !Expr deriving (Show, Generic)
  data Expr
      = Let Binds Expr
      | Case Expr [Alt]
      | AppF Text [Text]
      | AppC Text [Text]
      | AppP Text Text Text
      | LitE Int
      deriving (Show, Generic)
  data Alt = Alt Text [Text] Expr deriving (Show, Generic)
  
  instance NFData Program
  instance NFData Binds
  instance NFData LambdaForm
  instance NFData Expr
  instance NFData Alt
  
  instance Arbitrary Program where arbitrary = fmap Program arbitrary
  instance Arbitrary Binds where
      arbitrary = do
          NonEmpty xs <- arbitrary
          pure (Binds (M.fromList xs))
  instance Arbitrary LambdaForm where
      arbitrary = LambdaForm <$> fromTo 0 2 arbitrary <*> fromTo 0 2 arbitrary <*> arbitrary
  
  instance Arbitrary Expr where
      arbitrary = (oneof . map scaled)
          [ Let <$> arbitrary <*> arbitrary
          , Case <$> arbitrary <*> (do NonEmpty xs <- arbitrary; pure xs)
          , AppF <$> arbitrary <*> fromTo 0 3 arbitrary
          , AppC <$> ucFirst arbitrary <*> fromTo 0 3 arbitrary
          , AppP <$> arbitrary <*> arbitrary <*> arbitrary
          , LitE <$> arbitrary ]
  instance Arbitrary Alt where arbitrary = Alt <$> ucFirst arbitrary <*> fromTo 0 3 arbitrary <*> arbitrary
  instance Arbitrary Text where
      arbitrary = do
          n <- choose (3,6)
          str <- replicateM n (elements ['a'..'z'])
          if str `elem` ["let", "in", "case", "of"]
              then arbitrary
              else pure (T.pack str)
  
  ucFirst :: Gen Text -> Gen Text
  ucFirst gen = do
      x <- gen
      case T.uncons x of
          Nothing -> pure x
          Just (t,ext) -> pure (T.cons (toUpper t) ext)
  
  anCol :: Color -> Doc AnsiStyle -> Doc AnsiStyle
  anCol = annotate . color
  
  prettyProgram :: Program -> Doc AnsiStyle
  prettyProgram (Program binds) = annotate italicized $ prettyBinds binds
  
  prettyBinds :: Binds -> Doc AnsiStyle
  prettyBinds (Binds bs) = align (vsep (map prettyBinding (M.assocs bs)))
    where
      prettyBinding (var, lambda) = pretty var <+> anCol Red "=" <+> prettyLambdaForm lambda
  
  prettyLambdaForm :: LambdaForm -> Doc AnsiStyle
  prettyLambdaForm (LambdaForm free bound body) = prettyExp . (<+> anCol Blue "->") . prettyBound . prettyFree $ anCol Blue "\\"
    where
      prettyFree | null free = id
                 | otherwise = (<> anCol Blue lparen <> hsep (map pretty free) <> anCol Blue rparen)
      prettyBound | null bound = id
                  | null free = (<> hsep (map pretty bound))
                  | otherwise = (<+> hsep (map pretty bound))
      prettyExp = (<+> prettyExpr body)
  
  prettyExpr :: Expr -> Doc AnsiStyle
  prettyExpr = \expr -> case expr of
      Let binds body ->
          align (vsep [ anCol Red "let" <+> align (prettyBinds binds)
                      , anCol Red "in" <+> prettyExpr body ])
  
      Case scrutinee alts -> vsep
          [ anCol Yellow "case" <+> prettyExpr scrutinee <+> anCol Yellow "of"
          , indent 4 (align (vsep (map prettyAlt alts))) ]
  
      AppF f [] -> annotate bold . anCol Green $ pretty f
      AppF f args -> annotate bold . anCol Green $ pretty f <+> hsep (map pretty args)
  
      AppC c [] -> annotate bold . anCol Green $ pretty c
      AppC c args -> annotate bold . anCol Green $ pretty c <+> hsep (map pretty args)
  
      AppP op x y -> annotate bold . anCol Green $ pretty op <+> pretty x <+> pretty y
  
      LitE lit -> annotate bold . anCol Green $ pretty lit
  
  prettyAlt :: Alt -> Doc AnsiStyle
  prettyAlt (Alt con [] body) = pretty con <+> anCol Yellow "->" <+> prettyExpr body
  prettyAlt (Alt con args body) = pretty con <+> hsep (map pretty args) <+> anCol Yellow "->" <+> prettyExpr body
  
  scaled :: Gen a -> Gen a
  scaled = scale (\n -> n * 2 `quot` 3)
  
  fromTo :: Int -> Int -> Gen b -> Gen b
  fromTo a b gen = do
      n <- choose (min a b, max a b)
      resize n gen
  
  randomProgram
      :: Int -- ^ Seed
      -> Int -- ^ Generator size
      -> Program
  randomProgram seed size = let MkGen gen = arbitrary in gen (mkQCGen seed) size
  
  main :: IO ()
  main = do
      let prog = randomProgram 1 60
          layoutOpts = defaultLayoutOptions { layoutPageWidth = Unbounded }
          renderedProg = (renderLazy . layoutPretty layoutOpts . prettyProgram) prog
          (progLines, progWidth) = let l = TL.lines renderedProg in (length l, maximum (map TL.length l))
      putDoc ("Program size:" <+> pretty progLines <+> "lines, maximum width:" <+> pretty progWidth)
  
      let render :: (SimpleDocStream AnsiStyle -> TL.Text) -> Program -> TL.Text
          render r = r . layoutPretty layoutOpts . prettyProgram
  
      rnf prog `seq` T.putStrLn "Starting benchmark…"
  
      defaultMain
          [ bench "prettyprinter-ansi-terminal" $ nf (render Terminal.renderLazy) prog
          , bench "prettyprinter" $ nf (render Text.renderLazy) prog
          ]
  