-- %
-- % (c) The AQUA Project, Glasgow University, 1993-1998
-- %
-- \section[ScpMonad]{The Supercompiler Monad}
module ScpMonad (
        -- Some data types
	Savings(..),

	-- The monad
	ScpM,
	initScp,
	getDOptsSmpl, getSimplRules, 

        -- Unique supply
        MonadUnique(..), newId,

        -- Store
        addToStore, getStore, replaceStore,

        -- Binds
        getBinds, putBinds,

        -- Similar Expressions
        getSimilarExprs, putSimilarExprs,

	-- Savings traces
	addToSavings, getSavings, replaceSavings, expensiveSaving, cheapSaving,

        -- Tainted Expressions
        getTaintedExprs, incTaintedExpr, isTaintedExpr,

        -- Logging
        scpLog, scpLogDebug,
    ) where

import Id		( Id, mkSysLocal )
import Type             ( Type )
import Rules		( RuleBase )
import UniqSupply
import DynFlags		( DynFlags(..) ) 
import FastString
import Outputable

import FiniteMap        ( FiniteMap, emptyFM, addToFM, lookupWithDefaultFM )
import Var              ( Var(..) )
import CoreSyn          ( CoreExpr(..) )


-- %************************************************************************
-- %*									*
-- \subsection{Monad plumbing}
-- %*									*
-- %************************************************************************
--
-- For the supercompiler monad, we want to {\em thread} a unique supply and 
-- a counter.
-- (Command-line switches move around through the explicitly-passed ScpEnv.)

data Savings = SZero | SSmall Savings | SBig Savings | SSum [Savings] | SFold | SChoice [Savings] | SSplit [Savings] | SEnables Bool

newtype ScpM s result 
  =  SM { unSM :: SimplTopEnv	-- Envt that does not change much
		-> UniqSupply	-- We thread the unique supply because
				-- constantly splitting it is rather expensive
                -> FiniteMap Var CoreExpr -- Binds
                -> FiniteMap (Var, Int) [[CoreExpr]] -- Similar Exprs
                -> Savings
                -> FiniteMap Var Int -- Tainted exps
		-> s
		-> (result, UniqSupply, FiniteMap Var CoreExpr, FiniteMap (Var, Int) [[CoreExpr]], Savings, FiniteMap Var Int, s)}

data SimplTopEnv = STE	{ st_flags :: DynFlags 
			, st_rules :: RuleBase
                        , st_debug_level :: Int }

initScp :: DynFlags -> RuleBase 
	 -> UniqSupply		
	 -> ScpM [s] result
         -> result

initScp dflags rules us m
  = case unSM m env us emptyFM emptyFM SZero emptyFM [] of
	(result, _us, _bs, _se, _savings, _taint, _store) -> result
  where
    env = STE { st_flags = dflags, st_rules = rules, 
                st_debug_level = scpDebugLevel dflags }

{-# INLINE thenScp #-}
{-# INLINE thenScp_ #-}
{-# INLINE returnScp #-}

instance Monad (ScpM s) where
   (>>)   = thenScp_
   (>>=)  = thenScp
   return = returnScp

returnScp :: a -> ScpM s a 
returnScp e = SM (\_st_env us bs se savings taint sc -> (e, us, bs, se, savings, taint, sc))

thenScp  :: ScpM s a -> (a -> ScpM s b) -> ScpM s b 
thenScp_ :: ScpM s a -> ScpM s b -> ScpM s b

thenScp m k 
  = SM (\ st_env us0 bs0 se0 sv0 t0 sc0 ->
	  case (unSM m st_env us0 bs0 se0 sv0 t0 sc0) of 
		(m_result, us1, bs1, se1, sv1, t1, sc1) -> unSM (k m_result) st_env us1 bs1 se1 sv1 t1 sc1 )

thenScp_ m k 
  = SM (\st_env us0 bs0 se0 sv0 t0 sc0 ->
	 case (unSM m st_env us0 bs0 se0 sv0 t0 sc0) of 
		(_, us1, bs1, se1, sv1, t1, sc1) -> unSM k st_env us1 bs1 se1 sv1 t1 sc1)


-- %************************************************************************
-- %*									*
-- \subsection{The unique supply}
-- %*									*
-- %************************************************************************

instance MonadUnique (ScpM a) where
    getUniqueSupplyM
       = SM (\_st_env us bs se sv t sc -> case splitUniqSupply us of
                                (us1, us2) -> (us1, us2, bs, se, sv, t, sc))

    getUniqueM
       = SM (\_st_env us bs se sv t sc -> case splitUniqSupply us of
                                (us1, us2) -> (uniqFromSupply us1, us2, bs, se,sv, t, sc))

    getUniquesM
        = SM (\_st_env us bs se sv t sc -> case splitUniqSupply us of
                                (us1, us2) -> (uniqsFromSupply us1, us2, bs, se, sv, t, sc))

getDOptsSmpl :: ScpM s DynFlags
getDOptsSmpl = SM (\st_env us bs se sv t sc -> (st_flags st_env, us, bs, se, sv, t, sc))

getSimplRules :: ScpM s RuleBase
getSimplRules = SM (\st_env us bs se sv t sc -> (st_rules st_env, us, bs, se, sv, t, sc))

getDebugLevel :: ScpM s Int
getDebugLevel = SM (\st_env us bs se sv t sc -> (st_debug_level st_env, us, bs, se, sv, t, sc))

newId :: FastString -> Type -> ScpM s Id
newId fs ty = do uniq <- getUniqueM
                 return (mkSysLocal fs uniq ty)



-- %************************************************************************
-- %*									*
-- \subsection{The store}
-- %*									*
-- %************************************************************************


getStore :: ScpM s s
getStore = SM (\_st_env us bs se sv t sc -> (sc, us, bs, se, sv, t, sc))

addToStore :: s -> ScpM [s] ()
-- addToStore x = SM (\_st_env us sc -> ((), us, x:sc))
addToStore x = SM (\_st_env us bs se sv t sc -> x `seq` ((), us, bs, se, sv, t, x:sc))

replaceStore :: s -> ScpM s ()
replaceStore s = SM (\_st_env us bs se sv t _ -> ((), us, bs, se, sv, t, s))


-- %************************************************************************
-- %*									*
-- \subsection{The binds}
-- %*									*
-- %************************************************************************

getBinds :: ScpM s (FiniteMap Var CoreExpr)
getBinds = SM (\_st_env us bs se sv t sc -> (bs, us, bs, se, sv, t, sc))

putBinds :: FiniteMap Var CoreExpr -> ScpM s ()
putBinds bs = SM (\_st_env us _ se sv t sc -> ((), us, bs, se, sv, t, sc))

-- %************************************************************************
-- %*									*
-- \subsection{The Similar Expressions}
-- %*									*
-- %************************************************************************

getSimilarExprs :: ScpM s (FiniteMap (Var, Int) [[CoreExpr]])
getSimilarExprs = SM (\_st_env us bs se sv t sc -> (se, us, bs, se, sv, t, sc))

putSimilarExprs :: FiniteMap (Var, Int) [[CoreExpr]] -> ScpM s ()
putSimilarExprs se = SM (\_st_env us bs _ sv t sc -> ((), us, bs, se, sv, t, sc))

-- %************************************************************************
-- %*									*
-- \subsection{Savings}
-- %*									*
-- %************************************************************************


getSavings :: ScpM s Savings
getSavings = SM (\_st_env us bs se sv t sc -> (sv, us, bs, se, sv, t, sc))

addToSavings :: Savings -> ScpM s ()
addToSavings x = SM (\_st_env us bs se sv t sc -> x `seq` ((), us, bs, se, sv, t, sc))

replaceSavings :: Savings -> ScpM s ()
replaceSavings s = SM (\_st_env us bs se _ t sc -> ((), us, bs, se, s, t, sc))

expensiveSaving :: ScpM s ()
expensiveSaving = SM (\_st_env us bs se sv t sc -> ((), us, bs, se, SBig sv, t, sc))

cheapSaving :: ScpM s ()
cheapSaving = SM (\_st_env us bs se sv t sc -> ((), us, bs, se, SSmall sv, t, sc))

-- %************************************************************************
-- %*									*
-- \subsection{The Tainted Expressions}
-- %*									*
-- %************************************************************************

getTaintedExprs :: ScpM s (FiniteMap Var Int)
getTaintedExprs = SM (\_st_env us bs se sv t sc -> (t, us, bs, se, sv, t, sc))

incTaintedExpr :: Var -> ScpM s ()
incTaintedExpr v = SM (\_st_env us bs se sv t sc -> 
  let e = lookupWithDefaultFM t 0 v
      t' = addToFM t v (e + 1)
  in ((), us, bs, se, sv, t', sc))

isTaintedExpr :: Var -> ScpM s Bool
isTaintedExpr v = SM (\_st_env us bs se sv t sc ->
  let e = lookupWithDefaultFM t 0 v
  in (e >= 3, us, bs, se, sv, t, sc))

-- %************************************************************************
-- %*									*
-- \subsection{Logging}
-- %*									*
-- %************************************************************************


scpLog :: Int -> String -> SDoc -> ScpM s ()
scpLog l tmsg msg = do
  debugLevel <- getDebugLevel
  if l <= debugLevel
--    then liftIO $ putStrLn (tmsg ++ ' ':showSDoc msg)
    then trace (tmsg ++ ' ':showSDoc msg) $ return ()
    else return ()


scpLogDebug :: Int -> String -> SDoc -> ScpM s ()
scpLogDebug l tmsg msg = do
  debugLevel <- getDebugLevel
  if l <= debugLevel
--    then liftIO $ putStrLn (tmsg ++ ' ':showSDocDebug msg)
    then trace (tmsg ++ ' ':showSDocDebug msg) $ return ()
    else return ()