module DFMonad
    ( DataflowLattice(..) , DataflowAnalysis
    , markFactsUnchanged, factsStatus, getFact, setFact, getExitFact, setExitFact
                        , forgetFact, botFact, setAllFacts, getAllFacts, factsEnv
    , addLastOutFact, bareLastOutFacts, forgetLastOutFacts, checkFactMatch
    , subAnalysis

    , DFM, runDFM, liftToDFM
    , markGraphRewritten, graphWasRewritten
    , freshBlockId
    , module OptimizationFuel
    )
where

import BlockId
import CmmTx
import PprCmm()
import OptimizationFuel

import Control.Monad
import Maybes
import Outputable
import UniqFM
import UniqSupply

{-

A dataflow monad maintains a mapping from BlockIds to dataflow facts,
where a dataflow fact is a value of type [[a]].  Values of type [[a]]
must form a lattice, as described by type [[Fact a]].

The dataflow engine uses the lattice structure to compute a least
solution to a set of dataflow equations.  To compute a greatest
solution, flip the lattice over.

The engine works by starting at the bottom and iterating to a fixed
point, so in principle we require the bottom element, a join (least
upper bound) operation, and a comparison to find out if a value has
changed (grown).  In practice, the comparison is only ever used in
conjunction with the join, so we have [[fact_add_to]]:

  fact_add_to new old =
     let j = join new old in
     if j <= old then noTx old -- nothing changed
     else aTx j                -- the fact changed

-}

data DataflowLattice a = DataflowLattice  { 
  fact_name    :: String,                 -- documentation
  fact_bot     :: a,                      -- lattice bottom element
  fact_add_to  :: a -> a -> TxRes a,      -- lattice join and compare
    -- ^ compute join of two args; something changed iff join is greater than 2nd arg
  fact_do_logging :: Bool  -- log changes
}


-- DFM is the monad of combined analysis and transformation,
-- which needs a UniqSupply and may consume optimization fuel
-- DFM is defined using a monad transformer, DFM', which is the general
-- case of DFM, parameterized over any monad.
-- In practice, we apply DFM' to the FuelMonad, which provides optimization fuel and
-- the unique supply.
data DFState f = DFState { df_rewritten    :: ChangeFlag
                         , df_facts        :: BlockEnv f
                         , df_exit_fact    :: f
                         , df_last_outs    :: [(BlockId, f)]
                         , df_facts_change :: ChangeFlag
                         }

newtype DFM' m fact a = DFM' (DataflowLattice fact -> DFState  fact
                                                   -> m (a, DFState  fact))
type DFM fact a = DFM' FuelMonad fact a


runDFM :: Monad m => DataflowLattice f -> DFM' m f a -> m a
runDFM lattice (DFM' f) =
  (f lattice $ DFState NoChange emptyBlockEnv (fact_bot lattice)[] NoChange)
  >>= return . fst

class DataflowAnalysis m where
  markFactsUnchanged :: m f ()   -- ^ Useful for starting a new iteration
  factsStatus :: m f ChangeFlag
  subAnalysis :: m f a -> m f a  -- ^ Do a new analysis and then throw away
                                 -- /all/ the related state.

  getFact :: BlockId -> m f f
  setFact :: Outputable f => BlockId -> f -> m f ()
  getExitFact :: m f f
  setExitFact :: Outputable f => f -> m f  ()
  checkFactMatch :: Outputable f =>
                    BlockId -> f -> m f () -- ^ assert fact already at this val
  botFact :: m f f
  forgetFact :: BlockId -> m f ()
  -- | It might be surprising these next two are needed in a pure analysis,
  -- but for some problems we do a 'shallow' rewriting in which a rewritten
  -- graph is not itself considered for further rewriting but merely undergoes 
  -- an analysis.  In this case the results of a forward analysis might produce
  -- new facts that go on BlockId's that reside outside the graph being analyzed.
  -- Thus these 'lastOutFacts' need to be available even in a pure analysis. 
  addLastOutFact :: (BlockId, f) -> m f ()
  bareLastOutFacts :: m f [(BlockId, f)]
  forgetLastOutFacts :: m f ()
  getAllFacts :: m f (BlockEnv f)
  setAllFacts :: BlockEnv f -> m f ()
  factsEnv :: Monad (m f) => m f (BlockId -> f)

  lattice :: m f (DataflowLattice f)
  factsEnv = do { map <- getAllFacts
                ; bot <- botFact
                ; return $ \id -> lookupBlockEnv map id `orElse` bot }

instance Monad m => DataflowAnalysis (DFM' m) where
  markFactsUnchanged = DFM' f
    where f _ s = return ((), s {df_facts_change = NoChange}) 
  factsStatus = DFM' f'
    where f' _ s = return (df_facts_change s, s)
  subAnalysis (DFM' f) = DFM' f'
    where f' l s = do (a, _) <- f l (subAnalysisState s)
                      return (a, s)
  getFact id = DFM' get
    where get lattice s =
            return (lookupBlockEnv (df_facts s) id `orElse` fact_bot lattice, s)
  setFact id a = DFM' set
    where set (DataflowLattice name bot add_fact log) s =
            case add_fact a old of
                 TxRes NoChange _ -> if initialized then return ((), s) else update old old
                 TxRes SomeChange join -> update join old
              where (old, initialized) =
                      case lookupBlockEnv (df_facts s) id of
                        Just f  -> (f,   True)
                        Nothing -> (bot, False)
                    update join old =
                      let facts' = extendBlockEnv (df_facts s) id join
                          debug = if log then pprTrace else \_ _ a -> a
                      in  debug name (pprSetFact id old a join) $
                          return ((), s { df_facts = facts', df_facts_change = SomeChange })
  getExitFact = DFM' get
    where get _ s = return (df_exit_fact s, s)
  setExitFact a =
    do old <- getExitFact
       DataflowLattice { fact_add_to = add_fact
                       , fact_name = name, fact_do_logging = log } <- lattice
       case add_fact a old of
         TxRes NoChange _ -> return ()
         TxRes SomeChange join -> DFM' $ \_ s ->
             let debug = if log then pprTrace else \_ _ a -> a
             in  debug name (pprSetFact "exit" old a join) $
                 return ((), s { df_exit_fact = join, df_facts_change = SomeChange })
  getAllFacts = DFM' f
    where f _ s = return (df_facts s, s)
  setAllFacts env = DFM' f
    where f _ s = return ((), s { df_facts = env})
  botFact = DFM' f
    where f lattice s = return (fact_bot lattice, s)
  forgetFact id = DFM' f 
    where f _ s = return ((), s { df_facts = delFromUFM (df_facts s) id })
  addLastOutFact pair = DFM' f
    where f _ s = return ((), s { df_last_outs = pair : df_last_outs s })
  bareLastOutFacts = DFM' f
    where f _ s = return (df_last_outs s, s)
  forgetLastOutFacts = DFM' f
    where f _ s = return ((), s { df_last_outs = [] })
  checkFactMatch id a =
      do { fact <- lattice
         ; old_a <- getFact id
         ; case fact_add_to fact a old_a of
             TxRes NoChange _ -> return ()
             TxRes SomeChange new ->
               do { facts <- getAllFacts
                  ; pprPanic "checkFactMatch"
                            (f4sep [text (fact_name fact), text "at id" <+> ppr id,
                                    text "changed from", nest 4 (ppr old_a), text "to",
                                    nest 4 (ppr new),
                                    text "after supposedly reaching fixed point;",
                                    text "env is", pprFacts facts]) 
                  ; setFact id a }
         }
    where pprFacts env = vcat (map pprFact (ufmToList env))
          pprFact (id, a) = hang (ppr id <> colon) 4 (ppr a)

  lattice = DFM' f
    where f l s = return (l, s)

subAnalysisState :: DFState f -> DFState f
subAnalysisState s = s {df_facts_change = NoChange}


markGraphRewritten :: Monad m => DFM' m f ()
markGraphRewritten = DFM' f
    where f _ s = return ((), s {df_rewritten = SomeChange})

graphWasRewritten :: DFM f ChangeFlag
graphWasRewritten = DFM' f
    where f _ s = return (df_rewritten s, s)
                    
freshBlockId :: String -> DFM f BlockId
freshBlockId _s = getUniqueM >>= return . BlockId

instance Monad m => Monad (DFM' m f) where
  DFM' f >>= k = DFM' (\l s -> do (a, s') <- f l s
                                  let DFM' f' = k a in f' l s')
  return a = DFM' (\_ s -> return (a, s))

instance FuelUsingMonad (DFM' FuelMonad f) where
  fuelRemaining = liftToDFM' fuelRemaining
  lastFuelPass  = liftToDFM' lastFuelPass
  fuelExhausted = liftToDFM' fuelExhausted
  fuelDecrement p f f' = liftToDFM' (fuelDecrement p f f')
  fuelDec1      = liftToDFM' fuelDec1
instance MonadUnique (DFM' FuelMonad f) where
    getUniqueSupplyM = liftToDFM' getUniqueSupplyM
    getUniqueM       = liftToDFM' getUniqueM
    getUniquesM      = liftToDFM' getUniquesM

liftToDFM' :: Monad m => m x -> DFM' m f x
liftToDFM' m = DFM' (\ _ s -> m >>= (\a -> return (a, s)))
liftToDFM :: FuelMonad x -> DFM f x
liftToDFM m = DFM' (\ _ s -> m >>= (\a -> return (a, s)))


pprSetFact :: (Show a, Outputable f) => a -> f -> f -> f -> SDoc
pprSetFact id old a join =
    f4sep [text "at" <+> text (show id),
           text "added" <+> ppr a, text "to" <+> ppr old,
           text "yielding" <+> ppr join]

f4sep :: [SDoc] -> SDoc
f4sep [] = fsep []
f4sep (d:ds) = fsep (d : map (nest 4) ds)