{-# LANGUAGE CPP #-}
module StgStats ( showStgStats ) where
#include "HsVersions.h"
import GhcPrelude
import StgSyn
import Id (Id)
import Panic
import Data.Map (Map)
import qualified Data.Map as Map
data CounterType
  = Literals
  | Applications
  | ConstructorApps
  | PrimitiveApps
  | LetNoEscapes
  | StgCases
  | FreeVariables
  | ConstructorBinds Bool
  | ReEntrantBinds   Bool
  | SingleEntryBinds Bool
  | UpdatableBinds   Bool
  deriving (Eq, Ord)
type Count      = Int
type StatEnv    = Map CounterType Count
emptySE :: StatEnv
emptySE = Map.empty
combineSE :: StatEnv -> StatEnv -> StatEnv
combineSE = Map.unionWith (+)
combineSEs :: [StatEnv] -> StatEnv
combineSEs = foldr combineSE emptySE
countOne :: CounterType -> StatEnv
countOne c = Map.singleton c 1
showStgStats :: [StgTopBinding] -> String
showStgStats prog
  = "STG Statistics:\n\n"
    ++ concat (map showc (Map.toList (gatherStgStats prog)))
  where
    showc (x,n) = (showString (s x) . shows n) "\n"
    s Literals                = "Literals                   "
    s Applications            = "Applications               "
    s ConstructorApps         = "ConstructorApps            "
    s PrimitiveApps           = "PrimitiveApps              "
    s LetNoEscapes            = "LetNoEscapes               "
    s StgCases                = "StgCases                   "
    s FreeVariables           = "FreeVariables              "
    s (ConstructorBinds True) = "ConstructorBinds_Top       "
    s (ReEntrantBinds True)   = "ReEntrantBinds_Top         "
    s (SingleEntryBinds True) = "SingleEntryBinds_Top       "
    s (UpdatableBinds True)   = "UpdatableBinds_Top         "
    s (ConstructorBinds _)    = "ConstructorBinds_Nested    "
    s (ReEntrantBinds _)      = "ReEntrantBindsBinds_Nested "
    s (SingleEntryBinds _)    = "SingleEntryBinds_Nested    "
    s (UpdatableBinds _)      = "UpdatableBinds_Nested      "
gatherStgStats :: [StgTopBinding] -> StatEnv
gatherStgStats binds = combineSEs (map statTopBinding binds)
statTopBinding :: StgTopBinding -> StatEnv
statTopBinding (StgTopStringLit _ _) = countOne Literals
statTopBinding (StgTopLifted bind) = statBinding True bind
statBinding :: Bool 
            -> StgBinding
            -> StatEnv
statBinding top (StgNonRec b rhs)
  = statRhs top (b, rhs)
statBinding top (StgRec pairs)
  = combineSEs (map (statRhs top) pairs)
statRhs :: Bool -> (Id, StgRhs) -> StatEnv
statRhs top (_, StgRhsCon _ _ _)
  = countOne (ConstructorBinds top)
statRhs top (_, StgRhsClosure _ _ u _ body)
  = statExpr body `combineSE`
    countOne (
      case u of
        ReEntrant   -> ReEntrantBinds   top
        Updatable   -> UpdatableBinds   top
        SingleEntry -> SingleEntryBinds top
    )
statExpr :: StgExpr -> StatEnv
statExpr (StgApp _ _)     = countOne Applications
statExpr (StgLit _)       = countOne Literals
statExpr (StgConApp _ _ _)= countOne ConstructorApps
statExpr (StgOpApp _ _ _) = countOne PrimitiveApps
statExpr (StgTick _ e)    = statExpr e
statExpr (StgLetNoEscape _ binds body)
  = statBinding False binds    `combineSE`
    statExpr body                               `combineSE`
    countOne LetNoEscapes
statExpr (StgLet _ binds body)
  = statBinding False binds    `combineSE`
    statExpr body
statExpr (StgCase expr _ _ alts)
  = statExpr expr       `combineSE`
    stat_alts alts      `combineSE`
    countOne StgCases
  where
    stat_alts alts
        = combineSEs (map statExpr [ e | (_,_,e) <- alts ])
statExpr (StgLam {}) = panic "statExpr StgLam"