{-
(c) The GRASP/AQUA Project, Glasgow University, 1993-1998

\section[SimplStg]{Driver for simplifying @STG@ programs}
-}


{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}

module GHC.Stg.Pipeline
  ( StgPipelineOpts (..)
  , StgToDo (..)
  , stg2stg
  , StgCgInfos
  ) where

import GHC.Prelude

import GHC.Stg.Syntax

import GHC.Stg.Lint     ( lintStgTopBindings )
import GHC.Stg.Stats    ( showStgStats )
import GHC.Stg.FVs      ( depSortWithAnnotStgPgm )
import GHC.Stg.Unarise  ( unarise )
import GHC.Stg.BcPrep   ( bcPrep )
import GHC.Stg.CSE      ( stgCse )
import GHC.Stg.Lift     ( StgLiftConfig, stgLiftLams )
import GHC.Unit.Module ( Module )
import GHC.Runtime.Context ( InteractiveContext )

import GHC.Driver.Flags (DumpFlag(..))
import GHC.Utils.Error
import GHC.Types.Unique.Supply
import GHC.Utils.Outputable
import GHC.Utils.Logger
import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.Trans.Reader
import GHC.Settings (Platform)
import GHC.Stg.InferTags (inferTags)
import GHC.Types.Name.Env (NameEnv)
import GHC.Stg.InferTags.TagSig (TagSig)

data StgPipelineOpts = StgPipelineOpts
  { stgPipeline_phases      :: ![StgToDo]
  -- ^ Spec of what stg-to-stg passes to do
  , stgPipeline_lint        :: !(Maybe DiagOpts)
  -- ^ Should we lint the STG at various stages of the pipeline?
  , stgPipeline_pprOpts     :: !StgPprOpts
  , stgPlatform             :: !Platform
  }

newtype StgM a = StgM { _unStgM :: ReaderT Char IO a }
  deriving (Functor, Applicative, Monad, MonadIO)

-- | Information to be exposed in interface files which is produced
-- by the stg2stg passes.
type StgCgInfos = NameEnv TagSig

instance MonadUnique StgM where
  getUniqueSupplyM = StgM $ do { mask <- ask
                               ; liftIO $! mkSplitUniqSupply mask}
  getUniqueM = StgM $ do { mask <- ask
                         ; liftIO $! uniqFromMask mask}

runStgM :: Char -> StgM a -> IO a
runStgM mask (StgM m) = runReaderT m mask

stg2stg :: Logger
        -> InteractiveContext
        -> StgPipelineOpts
        -> Module                    -- module being compiled
        -> [StgTopBinding]           -- input program
        -> IO ([CgStgTopBinding], StgCgInfos) -- output program
stg2stg logger ictxt opts this_mod binds
  = do  { dump_when Opt_D_dump_stg_from_core "Initial STG:" binds
        ; showPass logger "Stg2Stg"
        -- Do the main business!
        ; binds' <- runStgM 'g' $
            foldM (do_stg_pass this_mod) binds (stgPipeline_phases opts)

          -- Dependency sort the program as last thing. The program needs to be
          -- in dependency order for the SRT algorithm to work (see
          -- CmmBuildInfoTables, which also includes a detailed description of
          -- the algorithm), and we don't guarantee that the program is already
          -- sorted at this point. #16192 is for simplifier not preserving
          -- dependency order. We also don't guarantee that StgLiftLams will
          -- preserve the order or only create minimal recursive groups, so a
          -- sorting pass is necessary.
          -- This pass will also augment each closure with non-global free variables
          -- annotations (which is used by code generator to compute offsets into closures)
        ; let binds_sorted_with_fvs = depSortWithAnnotStgPgm this_mod binds'
        -- See Note [Tag inference for interactive contexts]
        ; inferTags (stgPipeline_pprOpts opts) logger this_mod binds_sorted_with_fvs
   }

  where
    stg_linter unarised
      | Just diag_opts <- stgPipeline_lint opts
      = lintStgTopBindings
          (stgPlatform opts) logger
          diag_opts ppr_opts
          ictxt this_mod unarised
      | otherwise
      = \ _whodunnit _binds -> return ()

    -------------------------------------------
    do_stg_pass :: Module -> [StgTopBinding] -> StgToDo -> StgM [StgTopBinding]
    do_stg_pass this_mod binds to_do
      = case to_do of
          StgDoNothing ->
            return binds

          StgStats ->
            logTraceMsg logger "STG stats" (text (showStgStats binds)) (return binds)

          StgCSE -> do
            let binds' = {-# SCC "StgCse" #-} stgCse binds
            end_pass "StgCse" binds'

          StgLiftLams cfg -> do
            us <- getUniqueSupplyM
            --
            let binds' = {-# SCC "StgLiftLams" #-} stgLiftLams this_mod cfg us binds
            end_pass "StgLiftLams" binds'

          StgBcPrep -> do
            us <- getUniqueSupplyM
            let binds' = {-# SCC "StgBcPrep" #-} bcPrep us binds
            end_pass "StgBcPrep" binds'

          StgUnarise -> do
            us <- getUniqueSupplyM
            liftIO (stg_linter False "Pre-unarise" binds)
            let binds' = {-# SCC "StgUnarise" #-} unarise us binds
            liftIO (dump_when Opt_D_dump_stg_unarised "Unarised STG:" binds')
            liftIO (stg_linter True "Unarise" binds')
            return binds'

    ppr_opts = stgPipeline_pprOpts opts
    dump_when flag header binds
      = putDumpFileMaybe logger flag header FormatSTG (pprStgTopBindings ppr_opts binds)

    end_pass what binds2
      = liftIO $ do -- report verbosely, if required
          putDumpFileMaybe logger Opt_D_verbose_stg2stg what
            FormatSTG (vcat (map (pprStgTopBinding ppr_opts) binds2))
          stg_linter False what binds2
          return binds2

-- -----------------------------------------------------------------------------
-- StgToDo:  abstraction of stg-to-stg passes to run.

-- | Optional Stg-to-Stg passes.
data StgToDo
  = StgCSE
  -- ^ Common subexpression elimination
  | StgLiftLams StgLiftConfig
  -- ^ Lambda lifting closure variables, trading stack/register allocation for
  -- heap allocation
  | StgStats
  | StgUnarise
  -- ^ Mandatory unarise pass, desugaring unboxed tuple and sum binders
  | StgBcPrep
  -- ^ Mandatory when compiling to bytecode
  | StgDoNothing
  -- ^ Useful for building up 'getStgToDo'
  deriving (Show, Read, Eq, Ord)