-- | Clean out unneeded spill\/reload instructions.
--
--   Handling of join points
--   ~~~~~~~~~~~~~~~~~~~~~~~
--
--   B1:                          B2:
--    ...                          ...
--       RELOAD SLOT(0), %r1          RELOAD SLOT(0), %r1
--       ... A ...                    ... B ...
--       jump B3                      jump B3
--
--                B3: ... C ...
--                    RELOAD SLOT(0), %r1
--                    ...
--
--   The Plan
--   ~~~~~~~~
--   As long as %r1 hasn't been written to in A, B or C then we don't need
--   the reload in B3.
--
--   What we really care about here is that on the entry to B3, %r1 will
--   always have the same value that is in SLOT(0) (ie, %r1 is _valid_)
--
--   This also works if the reloads in B1\/B2 were spills instead, because
--   spilling %r1 to a slot makes that slot have the same value as %r1.
--
module RegAlloc.Graph.SpillClean (
        cleanSpills
) where
import GhcPrelude

import RegAlloc.Liveness
import Instruction
import Reg

import BlockId
import Cmm
import UniqSet
import UniqFM
import Unique
import State
import Outputable
import Platform
import Hoopl.Collections

import Data.List
import Data.Maybe
import Data.IntSet              (IntSet)
import qualified Data.IntSet    as IntSet

-- | The identification number of a spill slot.
--   A value is stored in a spill slot when we don't have a free
--   register to hold it.
type Slot = Int

-- | Clean out unneeded spill\/reloads from this top level thing.
cleanSpills
        :: Instruction instr
        => Platform
        -> LiveCmmDecl statics instr
        -> LiveCmmDecl statics instr

cleanSpills platform cmm
        = evalState (cleanSpin platform 0 cmm) initCleanS

-- | Do one pass of cleaning.
cleanSpin
        :: Instruction instr
        => Platform
        -> Int                              -- ^ Iteration number for the cleaner.
        -> LiveCmmDecl statics instr        -- ^ Liveness annotated code to clean.
        -> CleanM (LiveCmmDecl statics instr)

cleanSpin platform spinCount code
 = do
        -- Initialise count of cleaned spill and reload instructions.
        modify $ \s -> s
                { sCleanedSpillsAcc     = 0
                , sCleanedReloadsAcc    = 0
                , sReloadedBy           = emptyUFM }

        code_forward    <- mapBlockTopM (cleanBlockForward platform) code
        code_backward   <- cleanTopBackward code_forward

        -- During the cleaning of each block we collected information about
        -- what regs were valid across each jump. Based on this, work out
        -- whether it will be safe to erase reloads after join points for
        -- the next pass.
        collateJoinPoints

        -- Remember how many spill and reload instructions we cleaned in this pass.
        spills          <- gets sCleanedSpillsAcc
        reloads         <- gets sCleanedReloadsAcc
        modify $ \s -> s
                { sCleanedCount = (spills, reloads) : sCleanedCount s }

        -- If nothing was cleaned in this pass or the last one
        --      then we're done and it's time to bail out.
        cleanedCount    <- gets sCleanedCount
        if take 2 cleanedCount == [(0, 0), (0, 0)]
           then return code

        -- otherwise go around again
           else cleanSpin platform (spinCount + 1) code_backward

-------------------------------------------------------------------------------
-- | Clean out unneeded reload instructions,
--   while walking forward over the code.
cleanBlockForward
        :: Instruction instr
        => Platform
        -> LiveBasicBlock instr
        -> CleanM (LiveBasicBlock instr)

cleanBlockForward platform (BasicBlock blockId instrs)
 = do
        -- See if we have a valid association for the entry to this block.
        jumpValid       <- gets sJumpValid
        let assoc       = case lookupUFM jumpValid blockId of
                                Just assoc      -> assoc
                                Nothing         -> emptyAssoc

        instrs_reload   <- cleanForward platform blockId assoc [] instrs
        return  $ BasicBlock blockId instrs_reload

-- | Clean out unneeded reload instructions.
--
--   Walking forwards across the code
--     On a reload, if we know a reg already has the same value as a slot
--     then we don't need to do the reload.
--
cleanForward
        :: Instruction instr
        => Platform
        -> BlockId                  -- ^ the block that we're currently in
        -> Assoc Store              -- ^ two store locations are associated if
                                    --     they have the same value
        -> [LiveInstr instr]        -- ^ acc
        -> [LiveInstr instr]        -- ^ instrs to clean (in backwards order)
        -> CleanM [LiveInstr instr] -- ^ cleaned instrs  (in forward   order)

cleanForward _ _ _ acc []
        = return acc

-- Rewrite live range joins via spill slots to just a spill and a reg-reg move
-- hopefully the spill will be also be cleaned in the next pass
cleanForward platform blockId assoc acc (li1 : li2 : instrs)

        | LiveInstr (SPILL  reg1  slot1) _      <- li1
        , LiveInstr (RELOAD slot2 reg2)  _      <- li2
        , slot1 == slot2
        = do
                modify $ \s -> s { sCleanedReloadsAcc = sCleanedReloadsAcc s + 1 }
                cleanForward platform blockId assoc acc
                 $ li1 : LiveInstr (mkRegRegMoveInstr platform reg1 reg2) Nothing
                       : instrs

cleanForward platform blockId assoc acc (li@(LiveInstr i1 _) : instrs)
        | Just (r1, r2) <- takeRegRegMoveInstr i1
        = if r1 == r2
                -- Erase any left over nop reg reg moves while we're here
                -- this will also catch any nop moves that the previous case
                -- happens to add.
                then cleanForward platform blockId assoc acc instrs

                -- If r1 has the same value as some slots and we copy r1 to r2,
                --      then r2 is now associated with those slots instead
                else do let assoc'      = addAssoc (SReg r1) (SReg r2)
                                        $ delAssoc (SReg r2)
                                        $ assoc

                        cleanForward platform blockId assoc' (li : acc) instrs

cleanForward platform blockId assoc acc (li : instrs)

        -- Update association due to the spill.
        | LiveInstr (SPILL reg slot) _  <- li
        = let   assoc'  = addAssoc (SReg reg)  (SSlot slot)
                        $ delAssoc (SSlot slot)
                        $ assoc
          in    cleanForward platform blockId assoc' (li : acc) instrs

        -- Clean a reload instr.
        | LiveInstr (RELOAD{}) _        <- li
        = do    (assoc', mli)   <- cleanReload platform blockId assoc li
                case mli of
                 Nothing        -> cleanForward platform blockId assoc' acc
                                                instrs

                 Just li'       -> cleanForward platform blockId assoc' (li' : acc)
                                                instrs

        -- Remember the association over a jump.
        | LiveInstr instr _     <- li
        , targets               <- jumpDestsOfInstr instr
        , not $ null targets
        = do    mapM_ (accJumpValid assoc) targets
                cleanForward platform blockId assoc (li : acc) instrs

        -- Writing to a reg changes its value.
        | LiveInstr instr _     <- li
        , RU _ written          <- regUsageOfInstr platform instr
        = let assoc'    = foldr delAssoc assoc (map SReg $ nub written)
          in  cleanForward platform blockId assoc' (li : acc) instrs

-- | Try and rewrite a reload instruction to something more pleasing
cleanReload
        :: Instruction instr
        => Platform
        -> BlockId
        -> Assoc Store
        -> LiveInstr instr
        -> CleanM (Assoc Store, Maybe (LiveInstr instr))

cleanReload platform blockId assoc li@(LiveInstr (RELOAD slot reg) _)

        -- If the reg we're reloading already has the same value as the slot
        --      then we can erase the instruction outright.
        | elemAssoc (SSlot slot) (SReg reg) assoc
        = do    modify  $ \s -> s { sCleanedReloadsAcc = sCleanedReloadsAcc s + 1 }
                return  (assoc, Nothing)

        -- If we can find another reg with the same value as this slot then
        --      do a move instead of a reload.
        | Just reg2     <- findRegOfSlot assoc slot
        = do    modify $ \s -> s { sCleanedReloadsAcc = sCleanedReloadsAcc s + 1 }

                let assoc'      = addAssoc (SReg reg) (SReg reg2)
                                $ delAssoc (SReg reg)
                                $ assoc

                return  ( assoc'
                        , Just $ LiveInstr (mkRegRegMoveInstr platform reg2 reg) Nothing)

        -- Gotta keep this instr.
        | otherwise
        = do    -- Update the association.
                let assoc'
                        = addAssoc (SReg reg)  (SSlot slot)
                                -- doing the reload makes reg and slot the same value
                        $ delAssoc (SReg reg)
                                -- reg value changes on reload
                        $ assoc

                -- Remember that this block reloads from this slot.
                accBlockReloadsSlot blockId slot

                return  (assoc', Just li)

cleanReload _ _ _ _
        = panic "RegSpillClean.cleanReload: unhandled instr"

-------------------------------------------------------------------------------
-- | Clean out unneeded spill instructions,
--   while walking backwards over the code.
--
--      If there were no reloads from a slot between a spill and the last one
--      then the slot was never read and we don't need the spill.
--
--      SPILL   r0 -> s1
--      RELOAD  s1 -> r2
--      SPILL   r3 -> s1        <--- don't need this spill
--      SPILL   r4 -> s1
--      RELOAD  s1 -> r5
--
--      Maintain a set of
--              "slots which were spilled to but not reloaded from yet"
--
--      Walking backwards across the code:
--       a) On a reload from a slot, remove it from the set.
--
--       a) On a spill from a slot
--              If the slot is in set then we can erase the spill,
--               because it won't be reloaded from until after the next spill.
--
--              otherwise
--               keep the spill and add the slot to the set
--
-- TODO: This is mostly inter-block
--       we should really be updating the noReloads set as we cross jumps also.
--
-- TODO: generate noReloads from liveSlotsOnEntry
--
cleanTopBackward
        :: Instruction instr
        => LiveCmmDecl statics instr
        -> CleanM (LiveCmmDecl statics instr)

cleanTopBackward cmm
 = case cmm of
        CmmData{}
         -> return cmm

        CmmProc info label live sccs
         | LiveInfo _ _ _ liveSlotsOnEntry <- info
         -> do  sccs'   <- mapM (mapSCCM (cleanBlockBackward liveSlotsOnEntry)) sccs
                return  $ CmmProc info label live sccs'

cleanBlockBackward
        :: Instruction instr
        => BlockMap IntSet
        -> LiveBasicBlock instr
        -> CleanM (LiveBasicBlock instr)

cleanBlockBackward liveSlotsOnEntry (BasicBlock blockId instrs)
 = do   instrs_spill    <- cleanBackward liveSlotsOnEntry  emptyUniqSet  [] instrs
        return  $ BasicBlock blockId instrs_spill

cleanBackward
        :: Instruction instr
        => BlockMap IntSet          -- ^ Slots live on entry to each block
        -> UniqSet Int              -- ^ Slots that have been spilled, but not reloaded from
        -> [LiveInstr instr]        -- ^ acc
        -> [LiveInstr instr]        -- ^ Instrs to clean (in forwards order)
        -> CleanM [LiveInstr instr] -- ^ Cleaned instrs  (in backwards order)

cleanBackward liveSlotsOnEntry noReloads acc lis
 = do   reloadedBy      <- gets sReloadedBy
        cleanBackward' liveSlotsOnEntry reloadedBy noReloads acc lis

cleanBackward'
        :: Instruction instr
        => BlockMap IntSet
        -> UniqFM [BlockId]
        -> UniqSet Int
        -> [LiveInstr instr]
        -> [LiveInstr instr]
        -> State CleanS [LiveInstr instr]

cleanBackward' _ _ _      acc []
        = return  acc

cleanBackward' liveSlotsOnEntry reloadedBy noReloads acc (li : instrs)

        -- If nothing ever reloads from this slot then we don't need the spill.
        | LiveInstr (SPILL _ slot) _    <- li
        , Nothing       <- lookupUFM reloadedBy (SSlot slot)
        = do    modify $ \s -> s { sCleanedSpillsAcc = sCleanedSpillsAcc s + 1 }
                cleanBackward liveSlotsOnEntry noReloads acc instrs

        | LiveInstr (SPILL _ slot) _    <- li
        = if elementOfUniqSet slot noReloads

           -- We can erase this spill because the slot won't be read until
           -- after the next one
           then do
                modify $ \s -> s { sCleanedSpillsAcc = sCleanedSpillsAcc s + 1 }
                cleanBackward liveSlotsOnEntry noReloads acc instrs

           else do
                -- This slot is being spilled to, but we haven't seen any reloads yet.
                let noReloads'  = addOneToUniqSet noReloads slot
                cleanBackward liveSlotsOnEntry noReloads' (li : acc) instrs

        -- if we reload from a slot then it's no longer unused
        | LiveInstr (RELOAD slot _) _   <- li
        , noReloads'            <- delOneFromUniqSet noReloads slot
        = cleanBackward liveSlotsOnEntry noReloads' (li : acc) instrs

        -- If a slot is live in a jump target then assume it's reloaded there.
        --
        -- TODO: A real dataflow analysis would do a better job here.
        --       If the target block _ever_ used the slot then we assume
        --       it always does, but if those reloads are cleaned the slot
        --       liveness map doesn't get updated.
        | LiveInstr instr _     <- li
        , targets               <- jumpDestsOfInstr instr
        = do
                let slotsReloadedByTargets
                        = IntSet.unions
                        $ catMaybes
                        $ map (flip mapLookup liveSlotsOnEntry)
                        $ targets

                let noReloads'
                        = foldl' delOneFromUniqSet noReloads
                        $ IntSet.toList slotsReloadedByTargets

                cleanBackward liveSlotsOnEntry noReloads' (li : acc) instrs

        -- some other instruction
        | otherwise
        = cleanBackward liveSlotsOnEntry noReloads (li : acc) instrs

-- | Combine the associations from all the inward control flow edges.
--
collateJoinPoints :: CleanM ()
collateJoinPoints
 = modify $ \s -> s
        { sJumpValid    = mapUFM intersects (sJumpValidAcc s)
        , sJumpValidAcc = emptyUFM }

intersects :: [Assoc Store]     -> Assoc Store
intersects []           = emptyAssoc
intersects assocs       = foldl1' intersectAssoc assocs

-- | See if we have a reg with the same value as this slot in the association table.
findRegOfSlot :: Assoc Store -> Int -> Maybe Reg
findRegOfSlot assoc slot
        | close                 <- closeAssoc (SSlot slot) assoc
        , Just (SReg reg)       <- find isStoreReg $ nonDetEltsUniqSet close
           -- See Note [Unique Determinism and code generation]
        = Just reg

        | otherwise
        = Nothing

-------------------------------------------------------------------------------
-- | Cleaner monad.
type CleanM
        = State CleanS

-- | Cleaner state.
data CleanS
        = CleanS
        { -- | Regs which are valid at the start of each block.
          sJumpValid            :: UniqFM (Assoc Store)

          -- | Collecting up what regs were valid across each jump.
          --    in the next pass we can collate these and write the results
          --    to sJumpValid.
        , sJumpValidAcc         :: UniqFM [Assoc Store]

          -- | Map of (slot -> blocks which reload from this slot)
          --    used to decide if whether slot spilled to will ever be
          --    reloaded from on this path.
        , sReloadedBy           :: UniqFM [BlockId]

          -- | Spills and reloads cleaned each pass (latest at front)
        , sCleanedCount         :: [(Int, Int)]

          -- | Spills and reloads that have been cleaned in this pass so far.
        , sCleanedSpillsAcc     :: Int
        , sCleanedReloadsAcc    :: Int }

-- | Construct the initial cleaner state.
initCleanS :: CleanS
initCleanS
        = CleanS
        { sJumpValid            = emptyUFM
        , sJumpValidAcc         = emptyUFM

        , sReloadedBy           = emptyUFM

        , sCleanedCount         = []

        , sCleanedSpillsAcc     = 0
        , sCleanedReloadsAcc    = 0 }

-- | Remember the associations before a jump.
accJumpValid :: Assoc Store -> BlockId -> CleanM ()
accJumpValid assocs target
 = modify $ \s -> s {
        sJumpValidAcc = addToUFM_C (++)
                                (sJumpValidAcc s)
                                target
                                [assocs] }

accBlockReloadsSlot :: BlockId -> Slot -> CleanM ()
accBlockReloadsSlot blockId slot
 = modify $ \s -> s {
        sReloadedBy = addToUFM_C (++)
                                (sReloadedBy s)
                                (SSlot slot)
                                [blockId] }

-------------------------------------------------------------------------------
-- A store location can be a stack slot or a register
data Store
        = SSlot Int
        | SReg  Reg

-- | Check if this is a reg store.
isStoreReg :: Store -> Bool
isStoreReg ss
 = case ss of
        SSlot _ -> False
        SReg  _ -> True

-- Spill cleaning is only done once all virtuals have been allocated to realRegs
instance Uniquable Store where
    getUnique (SReg  r)
        | RegReal (RealRegSingle i)     <- r
        = mkRegSingleUnique i

        | RegReal (RealRegPair r1 r2)   <- r
        = mkRegPairUnique (r1 * 65535 + r2)

        | otherwise
        = error $ "RegSpillClean.getUnique: found virtual reg during spill clean,"
                ++ "only real regs expected."

    getUnique (SSlot i) = mkRegSubUnique i    -- [SLPJ] I hope "SubUnique" is ok

instance Outputable Store where
        ppr (SSlot i)   = text "slot" <> int i
        ppr (SReg  r)   = ppr r

-------------------------------------------------------------------------------
-- Association graphs.
-- In the spill cleaner, two store locations are associated if they are known
-- to hold the same value.
--
type Assoc a    = UniqFM (UniqSet a)

-- | An empty association
emptyAssoc :: Assoc a
emptyAssoc      = emptyUFM

-- | Add an association between these two things.
addAssoc :: Uniquable a
         => a -> a -> Assoc a -> Assoc a

addAssoc a b m
 = let  m1      = addToUFM_C unionUniqSets m  a (unitUniqSet b)
        m2      = addToUFM_C unionUniqSets m1 b (unitUniqSet a)
   in   m2

-- | Delete all associations to a node.
delAssoc :: (Uniquable a)
         => a -> Assoc a -> Assoc a

delAssoc a m
        | Just aSet     <- lookupUFM  m a
        , m1            <- delFromUFM m a
        = nonDetFoldUniqSet (\x m -> delAssoc1 x a m) m1 aSet
          -- It's OK to use nonDetFoldUFM here because deletion is commutative

        | otherwise     = m

-- | Delete a single association edge (a -> b).
delAssoc1 :: Uniquable a
          => a -> a -> Assoc a -> Assoc a

delAssoc1 a b m
        | Just aSet     <- lookupUFM m a
        = addToUFM m a (delOneFromUniqSet aSet b)

        | otherwise     = m

-- | Check if these two things are associated.
elemAssoc :: (Uniquable a)
          => a -> a -> Assoc a -> Bool

elemAssoc a b m
        = elementOfUniqSet b (closeAssoc a m)

-- | Find the refl. trans. closure of the association from this point.
closeAssoc :: (Uniquable a)
        => a -> Assoc a -> UniqSet a

closeAssoc a assoc
 =      closeAssoc' assoc emptyUniqSet (unitUniqSet a)
 where
        closeAssoc' assoc visited toVisit
         = case nonDetEltsUniqSet toVisit of
             -- See Note [Unique Determinism and code generation]

                -- nothing else to visit, we're done
                []      -> visited

                (x:_)
                 -- we've already seen this node
                 |  elementOfUniqSet x visited
                 -> closeAssoc' assoc visited (delOneFromUniqSet toVisit x)

                 -- haven't seen this node before,
                 --     remember to visit all its neighbors
                 |  otherwise
                 -> let neighbors
                         = case lookupUFM assoc x of
                                Nothing         -> emptyUniqSet
                                Just set        -> set

                   in closeAssoc' assoc
                        (addOneToUniqSet visited x)
                        (unionUniqSets   toVisit neighbors)

-- | Intersect two associations.
intersectAssoc :: Assoc a -> Assoc a -> Assoc a
intersectAssoc a b
        = intersectUFM_C (intersectUniqSets) a b


{-# LANGUAGE ScopedTypeVariables #-}
module RegAlloc.Graph.SpillCost (
        SpillCostRecord,
        plusSpillCostRecord,
        pprSpillCostRecord,

        SpillCostInfo,
        zeroSpillCostInfo,
        plusSpillCostInfo,

        slurpSpillCostInfo,
        chooseSpill,

        lifeMapFromSpillCostInfo
) where
import GhcPrelude

import RegAlloc.Liveness
import Instruction
import RegClass
import Reg

import GraphBase

import Hoopl.Collections (mapLookup)
import Cmm
import UniqFM
import UniqSet
import Digraph          (flattenSCCs)
import Outputable
import Platform
import State
import CFG

import Data.List        (nub, minimumBy)
import Data.Maybe
import Control.Monad (join)

-- | Records the expected cost to spill some regster.
type SpillCostRecord
 =      ( VirtualReg    -- register name
        , Int           -- number of writes to this reg
        , Int           -- number of reads from this reg
        , Int)          -- number of instrs this reg was live on entry to

-- | Map of `SpillCostRecord`
type SpillCostInfo
        = UniqFM SpillCostRecord

-- | Block membership in a loop
type LoopMember = Bool

type SpillCostState = State (UniqFM SpillCostRecord) ()

-- | An empty map of spill costs.
zeroSpillCostInfo :: SpillCostInfo
zeroSpillCostInfo       = emptyUFM

-- | Add two spill cost infos.
plusSpillCostInfo :: SpillCostInfo -> SpillCostInfo -> SpillCostInfo
plusSpillCostInfo sc1 sc2
        = plusUFM_C plusSpillCostRecord sc1 sc2

-- | Add two spill cost records.
plusSpillCostRecord :: SpillCostRecord -> SpillCostRecord -> SpillCostRecord
plusSpillCostRecord (r1, a1, b1, c1) (r2, a2, b2, c2)
        | r1 == r2      = (r1, a1 + a2, b1 + b2, c1 + c2)
        | otherwise     = error "RegSpillCost.plusRegInt: regs don't match"

-- | Slurp out information used for determining spill costs.
--
--   For each vreg, the number of times it was written to, read from,
--   and the number of instructions it was live on entry to (lifetime)
--
slurpSpillCostInfo :: forall instr statics. (Outputable instr, Instruction instr)
                   => Platform
                   -> Maybe CFG
                   -> LiveCmmDecl statics instr
                   -> SpillCostInfo

slurpSpillCostInfo platform cfg cmm
        = execState (countCmm cmm) zeroSpillCostInfo
 where
        countCmm CmmData{}              = return ()
        countCmm (CmmProc info _ _ sccs)
                = mapM_ (countBlock info)
                $ flattenSCCs sccs

        -- Lookup the regs that are live on entry to this block in
        --      the info table from the CmmProc.
        countBlock info (BasicBlock blockId instrs)
                | LiveInfo _ _ blockLive _ <- info
                , Just rsLiveEntry  <- mapLookup blockId blockLive
                , rsLiveEntry_virt  <- takeVirtuals rsLiveEntry
                = countLIs (loopMember blockId) rsLiveEntry_virt instrs

                | otherwise
                = error "RegAlloc.SpillCost.slurpSpillCostInfo: bad block"

        countLIs :: LoopMember -> UniqSet VirtualReg -> [LiveInstr instr] -> SpillCostState
        countLIs _      _      []
                = return ()

        -- Skip over comment and delta pseudo instrs.
        countLIs inLoop rsLive (LiveInstr instr Nothing : lis)
                | isMetaInstr instr
                = countLIs inLoop rsLive lis

                | otherwise
                = pprPanic "RegSpillCost.slurpSpillCostInfo"
                $ text "no liveness information on instruction " <> ppr instr

        countLIs inLoop rsLiveEntry (LiveInstr instr (Just live) : lis)
         = do
                -- Increment the lifetime counts for regs live on entry to this instr.
                mapM_ (incLifetime (loopCount inLoop)) $ nonDetEltsUniqSet rsLiveEntry
                    -- This is non-deterministic but we do not
                    -- currently support deterministic code-generation.
                    -- See Note [Unique Determinism and code generation]

                -- Increment counts for what regs were read/written from.
                let (RU read written)   = regUsageOfInstr platform instr
                mapM_ (incUses (loopCount inLoop)) $ catMaybes $ map takeVirtualReg $ nub read
                mapM_ (incDefs (loopCount inLoop)) $ catMaybes $ map takeVirtualReg $ nub written

                -- Compute liveness for entry to next instruction.
                let liveDieRead_virt    = takeVirtuals (liveDieRead  live)
                let liveDieWrite_virt   = takeVirtuals (liveDieWrite live)
                let liveBorn_virt       = takeVirtuals (liveBorn     live)

                let rsLiveAcross
                        = rsLiveEntry `minusUniqSet` liveDieRead_virt

                let rsLiveNext
                        = (rsLiveAcross `unionUniqSets` liveBorn_virt)
                                        `minusUniqSet`  liveDieWrite_virt

                countLIs inLoop rsLiveNext lis

        loopCount inLoop
          | inLoop = 10
          | otherwise = 1
        incDefs     count reg = modify $ \s -> addToUFM_C plusSpillCostRecord s reg (reg, count, 0, 0)
        incUses     count reg = modify $ \s -> addToUFM_C plusSpillCostRecord s reg (reg, 0, count, 0)
        incLifetime count reg = modify $ \s -> addToUFM_C plusSpillCostRecord s reg (reg, 0, 0, count)

        loopBlocks = CFG.loopMembers <$> cfg
        loopMember bid
          | Just isMember <- join (mapLookup bid <$> loopBlocks)
          = isMember
          | otherwise
          = False

-- | Take all the virtual registers from this set.
takeVirtuals :: UniqSet Reg -> UniqSet VirtualReg
takeVirtuals set = mkUniqSet
  [ vr | RegVirtual vr <- nonDetEltsUniqSet set ]
  -- See Note [Unique Determinism and code generation]

-- | Choose a node to spill from this graph
chooseSpill
        :: SpillCostInfo
        -> Graph VirtualReg RegClass RealReg
        -> VirtualReg

chooseSpill info graph
 = let  cost    = spillCost_length info graph
        node    = minimumBy (\n1 n2 -> compare (cost $ nodeId n1) (cost $ nodeId n2))
                $ nonDetEltsUFM $ graphMap graph
                -- See Note [Unique Determinism and code generation]

   in   nodeId node

-------------------------------------------------------------------------------
-- | Chaitins spill cost function is:
--
--   cost =     sum         loadCost * freq (u)  +    sum        storeCost * freq (d)
--          u <- uses (v)                         d <- defs (v)
--
--   There are no loops in our code at the moment, so we can set the freq's to 1.
--
--  If we don't have live range splitting then Chaitins function performs badly
--  if we have lots of nested live ranges and very few registers.
--
--               v1 v2 v3
--      def v1   .
--      use v1   .
--      def v2   .  .
--      def v3   .  .  .
--      use v1   .  .  .
--      use v3   .  .  .
--      use v2   .  .
--      use v1   .
--
--           defs uses degree   cost
--      v1:  1     3     3      1.5
--      v2:  1     2     3      1.0
--      v3:  1     1     3      0.666
--
--   v3 has the lowest cost, but if we only have 2 hardregs and we insert
--   spill code for v3 then this isn't going to improve the colorability of
--   the graph.
--
--  When compiling SHA1, which as very long basic blocks and some vregs
--  with very long live ranges the allocator seems to try and spill from
--  the inside out and eventually run out of stack slots.
--
--  Without live range splitting, its's better to spill from the outside
--  in so set the cost of very long live ranges to zero
--
{-
spillCost_chaitin
        :: SpillCostInfo
        -> Graph Reg RegClass Reg
        -> Reg
        -> Float

spillCost_chaitin info graph reg
        -- Spilling a live range that only lives for 1 instruction
        -- isn't going to help us at all - and we definitely want to avoid
        -- trying to re-spill previously inserted spill code.
        | lifetime <= 1         = 1/0

        -- It's unlikely that we'll find a reg for a live range this long
        -- better to spill it straight up and not risk trying to keep it around
        -- and have to go through the build/color cycle again.
        | lifetime > allocatableRegsInClass (regClass reg) * 10
        = 0

        -- Otherwise revert to chaitin's regular cost function.
        | otherwise     = fromIntegral (uses + defs)
                        / fromIntegral (nodeDegree graph reg)
        where (_, defs, uses, lifetime)
                = fromMaybe (reg, 0, 0, 0) $ lookupUFM info reg
-}

-- Just spill the longest live range.
spillCost_length
        :: SpillCostInfo
        -> Graph VirtualReg RegClass RealReg
        -> VirtualReg
        -> Float

spillCost_length info _ reg
        | lifetime <= 1         = 1/0
        | otherwise             = 1 / fromIntegral lifetime
        where (_, _, _, lifetime)
                = fromMaybe (reg, 0, 0, 0)
                $ lookupUFM info reg

-- | Extract a map of register lifetimes from a `SpillCostInfo`.
lifeMapFromSpillCostInfo :: SpillCostInfo -> UniqFM (VirtualReg, Int)
lifeMapFromSpillCostInfo info
        = listToUFM
        $ map (\(r, _, _, life) -> (r, (r, life)))
        $ nonDetEltsUFM info
        -- See Note [Unique Determinism and code generation]

-- | Determine the degree (number of neighbors) of this node which
--   have the same class.
nodeDegree
        :: (VirtualReg -> RegClass)
        -> Graph VirtualReg RegClass RealReg
        -> VirtualReg
        -> Int

nodeDegree classOfVirtualReg graph reg
        | Just node     <- lookupUFM (graphMap graph) reg

        , virtConflicts
           <- length
           $ filter (\r -> classOfVirtualReg r == classOfVirtualReg reg)
           $ nonDetEltsUniqSet
           -- See Note [Unique Determinism and code generation]
           $ nodeConflicts node

        = virtConflicts + sizeUniqSet (nodeExclusions node)

        | otherwise
        = 0

-- | Show a spill cost record, including the degree from the graph
--   and final calulated spill cost.
pprSpillCostRecord
        :: (VirtualReg -> RegClass)
        -> (Reg -> SDoc)
        -> Graph VirtualReg RegClass RealReg
        -> SpillCostRecord
        -> SDoc

pprSpillCostRecord regClass pprReg graph (reg, uses, defs, life)
        =  hsep
        [ pprReg (RegVirtual reg)
        , ppr uses
        , ppr defs
        , ppr life
        , ppr $ nodeDegree regClass graph reg
        , text $ show $ (fromIntegral (uses + defs)
                       / fromIntegral (nodeDegree regClass graph reg) :: Float) ]


{-# LANGUAGE BangPatterns, CPP #-}

-- | Carries interesting info for debugging / profiling of the
--   graph coloring register allocator.
module RegAlloc.Graph.Stats (
        RegAllocStats (..),

        pprStats,
        pprStatsSpills,
        pprStatsLifetimes,
        pprStatsConflict,
        pprStatsLifeConflict,

        countSRMs, addSRM
) where

#include "nativeGen/NCG.h"

import GhcPrelude

import qualified GraphColor as Color
import RegAlloc.Liveness
import RegAlloc.Graph.Spill
import RegAlloc.Graph.SpillCost
import RegAlloc.Graph.TrivColorable
import Instruction
import RegClass
import Reg
import TargetReg

import PprCmm()
import Outputable
import UniqFM
import UniqSet
import State

-- | Holds interesting statistics from the register allocator.
data RegAllocStats statics instr

        -- Information about the initial conflict graph.
        = RegAllocStatsStart
        { -- | Initial code, with liveness.
          raLiveCmm     :: [LiveCmmDecl statics instr]

          -- | The initial, uncolored graph.
        , raGraph       :: Color.Graph VirtualReg RegClass RealReg

          -- | Information to help choose which regs to spill.
        , raSpillCosts  :: SpillCostInfo }

        -- Information about an intermediate graph.
        -- This is one that we couldn't color, so had to insert spill code
        -- instruction stream.
        | RegAllocStatsSpill
        { -- | Code we tried to allocate registers for.
          raCode        :: [LiveCmmDecl statics instr]

          -- | Partially colored graph.
        , raGraph       :: Color.Graph VirtualReg RegClass RealReg

          -- | The regs that were coalesced.
        , raCoalesced   :: UniqFM VirtualReg

          -- | Spiller stats.
        , raSpillStats  :: SpillStats

          -- | Number of instructions each reg lives for.
        , raSpillCosts  :: SpillCostInfo

          -- | Code with spill instructions added.
        , raSpilled     :: [LiveCmmDecl statics instr] }

        -- a successful coloring
        | RegAllocStatsColored
        { -- | Code we tried to allocate registers for.
          raCode          :: [LiveCmmDecl statics instr]

          -- | Uncolored graph.
        , raGraph         :: Color.Graph VirtualReg RegClass RealReg

          -- | Coalesced and colored graph.
        , raGraphColored  :: Color.Graph VirtualReg RegClass RealReg

          -- | Regs that were coalesced.
        , raCoalesced     :: UniqFM VirtualReg

          -- | Code with coalescings applied.
        , raCodeCoalesced :: [LiveCmmDecl statics instr]

          -- | Code with vregs replaced by hregs.
        , raPatched       :: [LiveCmmDecl statics instr]

          -- | Code with unneeded spill\/reloads cleaned out.
        , raSpillClean    :: [LiveCmmDecl statics instr]

          -- | Final code.
        , raFinal         :: [NatCmmDecl statics instr]

          -- | Spill\/reload\/reg-reg moves present in this code.
        , raSRMs          :: (Int, Int, Int) }

instance (Outputable statics, Outputable instr)
       => Outputable (RegAllocStats statics instr) where

 ppr (s@RegAllocStatsStart{}) = sdocWithPlatform $ \platform ->
           text "#  Start"
        $$ text "#  Native code with liveness information."
        $$ ppr (raLiveCmm s)
        $$ text ""
        $$ text "#  Initial register conflict graph."
        $$ Color.dotGraph
                (targetRegDotColor platform)
                (trivColorable platform
                        (targetVirtualRegSqueeze platform)
                        (targetRealRegSqueeze platform))
                (raGraph s)

 ppr (s@RegAllocStatsSpill{}) =
           text "#  Spill"

        $$ text "#  Code with liveness information."
        $$ ppr (raCode s)
        $$ text ""

        $$ (if (not $ isNullUFM $ raCoalesced s)
                then    text "#  Registers coalesced."
                        $$ pprUFMWithKeys (raCoalesced s) (vcat . map ppr)
                        $$ text ""
                else empty)

        $$ text "#  Spills inserted."
        $$ ppr (raSpillStats s)
        $$ text ""

        $$ text "#  Code with spills inserted."
        $$ ppr (raSpilled s)

 ppr (s@RegAllocStatsColored { raSRMs = (spills, reloads, moves) })
    = sdocWithPlatform $ \platform ->
           text "#  Colored"

        $$ text "#  Code with liveness information."
        $$ ppr (raCode s)
        $$ text ""

        $$ text "#  Register conflict graph (colored)."
        $$ Color.dotGraph
                (targetRegDotColor platform)
                (trivColorable platform
                        (targetVirtualRegSqueeze platform)
                        (targetRealRegSqueeze platform))
                (raGraphColored s)
        $$ text ""

        $$ (if (not $ isNullUFM $ raCoalesced s)
                then    text "#  Registers coalesced."
                        $$ pprUFMWithKeys (raCoalesced s) (vcat . map ppr)
                        $$ text ""
                else empty)

        $$ text "#  Native code after coalescings applied."
        $$ ppr (raCodeCoalesced s)
        $$ text ""

        $$ text "#  Native code after register allocation."
        $$ ppr (raPatched s)
        $$ text ""

        $$ text "#  Clean out unneeded spill/reloads."
        $$ ppr (raSpillClean s)
        $$ text ""

        $$ text "#  Final code, after rewriting spill/rewrite pseudo instrs."
        $$ ppr (raFinal s)
        $$ text ""
        $$  text "#  Score:"
        $$ (text "#          spills  inserted: " <> int spills)
        $$ (text "#          reloads inserted: " <> int reloads)
        $$ (text "#   reg-reg moves remaining: " <> int moves)
        $$ text ""

-- | Do all the different analysis on this list of RegAllocStats
pprStats
        :: [RegAllocStats statics instr]
        -> Color.Graph VirtualReg RegClass RealReg
        -> SDoc

pprStats stats graph
 = let  outSpills       = pprStatsSpills    stats
        outLife         = pprStatsLifetimes stats
        outConflict     = pprStatsConflict  stats
        outScatter      = pprStatsLifeConflict stats graph

  in    vcat [outSpills, outLife, outConflict, outScatter]

-- | Dump a table of how many spill loads \/ stores were inserted for each vreg.
pprStatsSpills
        :: [RegAllocStats statics instr] -> SDoc

pprStatsSpills stats
 = let
        finals  = [ s   | s@RegAllocStatsColored{} <- stats]

        -- sum up how many stores\/loads\/reg-reg-moves were left in the code
        total   = foldl' addSRM (0, 0, 0)
                $ map raSRMs finals

    in  (  text "-- spills-added-total"
        $$ text "--    (stores, loads, reg_reg_moves_remaining)"
        $$ ppr total
        $$ text "")

-- | Dump a table of how long vregs tend to live for in the initial code.
pprStatsLifetimes
        :: [RegAllocStats statics instr] -> SDoc

pprStatsLifetimes stats
 = let  info            = foldl' plusSpillCostInfo zeroSpillCostInfo
                                [ raSpillCosts s
                                        | s@RegAllocStatsStart{} <- stats ]

        lifeBins        = binLifetimeCount $ lifeMapFromSpillCostInfo info

   in   (  text "-- vreg-population-lifetimes"
        $$ text "--   (instruction_count, number_of_vregs_that_lived_that_long)"
        $$ pprUFM lifeBins (vcat . map ppr)
        $$ text "\n")

binLifetimeCount :: UniqFM (VirtualReg, Int) -> UniqFM (Int, Int)
binLifetimeCount fm
 = let  lifes   = map (\l -> (l, (l, 1)))
                $ map snd
                $ nonDetEltsUFM fm
                -- See Note [Unique Determinism and code generation]

   in   addListToUFM_C
                (\(l1, c1) (_, c2) -> (l1, c1 + c2))
                emptyUFM
                lifes

-- | Dump a table of how many conflicts vregs tend to have in the initial code.
pprStatsConflict
        :: [RegAllocStats statics instr] -> SDoc

pprStatsConflict stats
 = let  confMap = foldl' (plusUFM_C (\(c1, n1) (_, n2) -> (c1, n1 + n2)))
                        emptyUFM
                $ map Color.slurpNodeConflictCount
                        [ raGraph s | s@RegAllocStatsStart{} <- stats ]

   in   (  text "-- vreg-conflicts"
        $$ text "--   (conflict_count, number_of_vregs_that_had_that_many_conflicts)"
        $$ pprUFM confMap (vcat . map ppr)
        $$ text "\n")

-- | For every vreg, dump how many conflicts it has, and its lifetime.
--      Good for making a scatter plot.
pprStatsLifeConflict
        :: [RegAllocStats statics instr]
        -> Color.Graph VirtualReg RegClass RealReg -- ^ global register conflict graph
        -> SDoc

pprStatsLifeConflict stats graph
 = let  lifeMap = lifeMapFromSpillCostInfo
                $ foldl' plusSpillCostInfo zeroSpillCostInfo
                $ [ raSpillCosts s | s@RegAllocStatsStart{} <- stats ]

        scatter = map   (\r ->  let lifetime  = case lookupUFM lifeMap r of
                                                      Just (_, l) -> l
                                                      Nothing     -> 0
                                    Just node = Color.lookupNode graph r
                                in parens $ hcat $ punctuate (text ", ")
                                        [ doubleQuotes $ ppr $ Color.nodeId node
                                        , ppr $ sizeUniqSet (Color.nodeConflicts node)
                                        , ppr $ lifetime ])
                $ map Color.nodeId
                $ nonDetEltsUFM
                -- See Note [Unique Determinism and code generation]
                $ Color.graphMap graph

   in   (  text "-- vreg-conflict-lifetime"
        $$ text "--   (vreg, vreg_conflicts, vreg_lifetime)"
        $$ (vcat scatter)
        $$ text "\n")

-- | Count spill/reload/reg-reg moves.
--      Lets us see how well the register allocator has done.
countSRMs
        :: Instruction instr
        => LiveCmmDecl statics instr -> (Int, Int, Int)

countSRMs cmm
        = execState (mapBlockTopM countSRM_block cmm) (0, 0, 0)

countSRM_block
        :: Instruction instr
        => GenBasicBlock (LiveInstr instr)
        -> State (Int, Int, Int) (GenBasicBlock (LiveInstr instr))

countSRM_block (BasicBlock i instrs)
 = do   instrs' <- mapM countSRM_instr instrs
        return  $ BasicBlock i instrs'

countSRM_instr
        :: Instruction instr
        => LiveInstr instr -> State (Int, Int, Int) (LiveInstr instr)

countSRM_instr li
        | LiveInstr SPILL{} _    <- li
        = do    modify  $ \(s, r, m)    -> (s + 1, r, m)
                return li

        | LiveInstr RELOAD{} _  <- li
        = do    modify  $ \(s, r, m)    -> (s, r + 1, m)
                return li

        | LiveInstr instr _     <- li
        , Just _        <- takeRegRegMoveInstr instr
        = do    modify  $ \(s, r, m)    -> (s, r, m + 1)
                return li

        | otherwise
        =       return li

-- sigh..
addSRM :: (Int, Int, Int) -> (Int, Int, Int) -> (Int, Int, Int)
addSRM (s1, r1, m1) (s2, r2, m2)
 = let  !s = s1 + s2
        !r = r1 + r2
        !m = m1 + m2
   in   (s, r, m)


{-# LANGUAGE CPP #-}

module RegAlloc.Graph.TrivColorable (
        trivColorable,
)

where

#include "HsVersions.h"

import GhcPrelude

import RegClass
import Reg

import GraphBase

import UniqSet
import Platform
import Panic

-- trivColorable ---------------------------------------------------------------

-- trivColorable function for the graph coloring allocator
--
--      This gets hammered by scanGraph during register allocation,
--      so needs to be fairly efficient.
--
--      NOTE:   This only works for arcitectures with just RcInteger and RcDouble
--              (which are disjoint) ie. x86, x86_64 and ppc
--
--      The number of allocatable regs is hard coded in here so we can do
--              a fast comparison in trivColorable.
--
--      It's ok if these numbers are _less_ than the actual number of free
--              regs, but they can't be more or the register conflict
--              graph won't color.
--
--      If the graph doesn't color then the allocator will panic, but it won't
--              generate bad object code or anything nasty like that.
--
--      There is an allocatableRegsInClass :: RegClass -> Int, but doing
--      the unboxing is too slow for us here.
--      TODO: Is that still true? Could we use allocatableRegsInClass
--      without losing performance now?
--
--      Look at includes/stg/MachRegs.h to get the numbers.
--

-- Disjoint registers ----------------------------------------------------------
--
--      The definition has been unfolded into individual cases for speed.
--      Each architecture has a different register setup, so we use a
--      different regSqueeze function for each.
--
accSqueeze
        :: Int
        -> Int
        -> (reg -> Int)
        -> UniqSet reg
        -> Int

accSqueeze count maxCount squeeze us = acc count (nonDetEltsUniqSet us)
  -- See Note [Unique Determinism and code generation]
  where acc count [] = count
        acc count _ | count >= maxCount = count
        acc count (r:rs) = acc (count + squeeze r) rs

{- Note [accSqueeze]
~~~~~~~~~~~~~~~~~~~~
BL 2007/09
Doing a nice fold over the UniqSet makes trivColorable use
32% of total compile time and 42% of total alloc when compiling SHA1.hs from darcs.
Therefore the UniqFM is made non-abstract and we use custom fold.

MS 2010/04
When converting UniqFM to use Data.IntMap, the fold cannot use UniqFM internal
representation any more. But it is imperative that the accSqueeze stops
the folding if the count gets greater or equal to maxCount. We thus convert
UniqFM to a (lazy) list, do the fold and stops if necessary, which was
the most efficient variant tried. Benchmark compiling 10-times SHA1.hs follows.
(original = previous implementation, folding = fold of the whole UFM,
 lazyFold = the current implementation,
 hackFold = using internal representation of Data.IntMap)

                                 original  folding   hackFold  lazyFold
 -O -fasm (used everywhere)      31.509s   30.387s   30.791s   30.603s
                                 100.00%   96.44%    97.72%    97.12%
 -fregs-graph                    67.938s   74.875s   62.673s   64.679s
                                 100.00%   110.21%   92.25%    95.20%
 -fregs-iterative                89.761s   143.913s  81.075s   86.912s
                                 100.00%   160.33%   90.32%    96.83%
 -fnew-codegen                   38.225s   37.142s   37.551s   37.119s
                                 100.00%   97.17%    98.24%    97.11%
 -fnew-codegen -fregs-graph      91.786s   91.51s    87.368s   86.88s
                                 100.00%   99.70%    95.19%    94.65%
 -fnew-codegen -fregs-iterative  206.72s   343.632s  194.694s  208.677s
                                 100.00%   166.23%   94.18%    100.95%
-}

trivColorable
        :: Platform
        -> (RegClass -> VirtualReg -> Int)
        -> (RegClass -> RealReg    -> Int)
        -> Triv VirtualReg RegClass RealReg

trivColorable platform virtualRegSqueeze realRegSqueeze RcInteger conflicts exclusions
        | let cALLOCATABLE_REGS_INTEGER
                  =        (case platformArch platform of
                            ArchX86       -> 3
                            ArchX86_64    -> 5
                            ArchPPC       -> 16
                            ArchSPARC     -> 14
                            ArchSPARC64   -> panic "trivColorable ArchSPARC64"
                            ArchPPC_64 _  -> 15
                            ArchARM _ _ _ -> panic "trivColorable ArchARM"
                            ArchARM64     -> panic "trivColorable ArchARM64"
                            ArchAlpha     -> panic "trivColorable ArchAlpha"
                            ArchMipseb    -> panic "trivColorable ArchMipseb"
                            ArchMipsel    -> panic "trivColorable ArchMipsel"
                            ArchJavaScript-> panic "trivColorable ArchJavaScript"
                            ArchUnknown   -> panic "trivColorable ArchUnknown")
        , count2        <- accSqueeze 0 cALLOCATABLE_REGS_INTEGER
                                (virtualRegSqueeze RcInteger)
                                conflicts

        , count3        <- accSqueeze  count2    cALLOCATABLE_REGS_INTEGER
                                (realRegSqueeze   RcInteger)
                                exclusions

        = count3 < cALLOCATABLE_REGS_INTEGER

trivColorable platform virtualRegSqueeze realRegSqueeze RcFloat conflicts exclusions
        | let cALLOCATABLE_REGS_FLOAT
                  =        (case platformArch platform of
                            ArchX86       -> 0
                            ArchX86_64    -> 0
                            ArchPPC       -> 0
                            ArchSPARC     -> 22
                            ArchSPARC64   -> panic "trivColorable ArchSPARC64"
                            ArchPPC_64 _  -> 0
                            ArchARM _ _ _ -> panic "trivColorable ArchARM"
                            ArchARM64     -> panic "trivColorable ArchARM64"
                            ArchAlpha     -> panic "trivColorable ArchAlpha"
                            ArchMipseb    -> panic "trivColorable ArchMipseb"
                            ArchMipsel    -> panic "trivColorable ArchMipsel"
                            ArchJavaScript-> panic "trivColorable ArchJavaScript"
                            ArchUnknown   -> panic "trivColorable ArchUnknown")
        , count2        <- accSqueeze 0 cALLOCATABLE_REGS_FLOAT
                                (virtualRegSqueeze RcFloat)
                                conflicts

        , count3        <- accSqueeze  count2    cALLOCATABLE_REGS_FLOAT
                                (realRegSqueeze   RcFloat)
                                exclusions

        = count3 < cALLOCATABLE_REGS_FLOAT

trivColorable platform virtualRegSqueeze realRegSqueeze RcDouble conflicts exclusions
        | let cALLOCATABLE_REGS_DOUBLE
                  =        (case platformArch platform of
                            ArchX86       -> 6
                            ArchX86_64    -> 0
                            ArchPPC       -> 26
                            ArchSPARC     -> 11
                            ArchSPARC64   -> panic "trivColorable ArchSPARC64"
                            ArchPPC_64 _  -> 20
                            ArchARM _ _ _ -> panic "trivColorable ArchARM"
                            ArchARM64     -> panic "trivColorable ArchARM64"
                            ArchAlpha     -> panic "trivColorable ArchAlpha"
                            ArchMipseb    -> panic "trivColorable ArchMipseb"
                            ArchMipsel    -> panic "trivColorable ArchMipsel"
                            ArchJavaScript-> panic "trivColorable ArchJavaScript"
                            ArchUnknown   -> panic "trivColorable ArchUnknown")
        , count2        <- accSqueeze 0 cALLOCATABLE_REGS_DOUBLE
                                (virtualRegSqueeze RcDouble)
                                conflicts

        , count3        <- accSqueeze  count2    cALLOCATABLE_REGS_DOUBLE
                                (realRegSqueeze   RcDouble)
                                exclusions

        = count3 < cALLOCATABLE_REGS_DOUBLE

trivColorable platform virtualRegSqueeze realRegSqueeze RcDoubleSSE conflicts exclusions
        | let cALLOCATABLE_REGS_SSE
                  =        (case platformArch platform of
                            ArchX86       -> 8
                            ArchX86_64    -> 10
                            ArchPPC       -> 0
                            ArchSPARC     -> 0
                            ArchSPARC64   -> panic "trivColorable ArchSPARC64"
                            ArchPPC_64 _  -> 0
                            ArchARM _ _ _ -> panic "trivColorable ArchARM"
                            ArchARM64     -> panic "trivColorable ArchARM64"
                            ArchAlpha     -> panic "trivColorable ArchAlpha"
                            ArchMipseb    -> panic "trivColorable ArchMipseb"
                            ArchMipsel    -> panic "trivColorable ArchMipsel"
                            ArchJavaScript-> panic "trivColorable ArchJavaScript"
                            ArchUnknown   -> panic "trivColorable ArchUnknown")
        , count2        <- accSqueeze 0 cALLOCATABLE_REGS_SSE
                                (virtualRegSqueeze RcDoubleSSE)
                                conflicts

        , count3        <- accSqueeze  count2    cALLOCATABLE_REGS_SSE
                                (realRegSqueeze   RcDoubleSSE)
                                exclusions

        = count3 < cALLOCATABLE_REGS_SSE

-- Specification Code ----------------------------------------------------------
--
--      The trivColorable function for each particular architecture should
--      implement the following function, but faster.
--

{-
trivColorable :: RegClass -> UniqSet Reg -> UniqSet Reg -> Bool
trivColorable classN conflicts exclusions
 = let

        acc :: Reg -> (Int, Int) -> (Int, Int)
        acc r (cd, cf)
         = case regClass r of
                RcInteger       -> (cd+1, cf)
                RcFloat         -> (cd,   cf+1)
                _               -> panic "Regs.trivColorable: reg class not handled"

        tmp                     = nonDetFoldUFM acc (0, 0) conflicts
        (countInt,  countFloat) = nonDetFoldUFM acc tmp    exclusions

        squeese         = worst countInt   classN RcInteger
                        + worst countFloat classN RcFloat

   in   squeese < allocatableRegsInClass classN

-- | Worst case displacement
--      node N of classN has n neighbors of class C.
--
--      We currently only have RcInteger and RcDouble, which don't conflict at all.
--      This is a bit boring compared to what's in RegArchX86.
--
worst :: Int -> RegClass -> RegClass -> Int
worst n classN classC
 = case classN of
        RcInteger
         -> case classC of
                RcInteger       -> min n (allocatableRegsInClass RcInteger)
                RcFloat         -> 0

        RcDouble
         -> case classC of
                RcFloat         -> min n (allocatableRegsInClass RcFloat)
                RcInteger       -> 0

-- allocatableRegs is allMachRegNos with the fixed-use regs removed.
-- i.e., these are the regs for which we are prepared to allow the
-- register allocator to attempt to map VRegs to.
allocatableRegs :: [RegNo]
allocatableRegs
   = let isFree i = freeReg i
     in  filter isFree allMachRegNos

-- | The number of regs in each class.
--      We go via top level CAFs to ensure that we're not recomputing
--      the length of these lists each time the fn is called.
allocatableRegsInClass :: RegClass -> Int
allocatableRegsInClass cls
 = case cls of
        RcInteger       -> allocatableRegsInteger
        RcFloat         -> allocatableRegsDouble

allocatableRegsInteger :: Int
allocatableRegsInteger
        = length $ filter (\r -> regClass r == RcInteger)
                 $ map RealReg allocatableRegs

allocatableRegsFloat :: Int
allocatableRegsFloat
        = length $ filter (\r -> regClass r == RcFloat
                 $ map RealReg allocatableRegs
-}


-- | Put common type definitions here to break recursive module dependencies.

module RegAlloc.Linear.Base (
        BlockAssignment,

        Loc(..),
        regsOfLoc,

        -- for stats
        SpillReason(..),
        RegAllocStats(..),

        -- the allocator monad
        RA_State(..),
)

where

import GhcPrelude

import RegAlloc.Linear.StackMap
import RegAlloc.Liveness
import Reg

import DynFlags
import Outputable
import Unique
import UniqFM
import UniqSupply
import BlockId

-- | Used to store the register assignment on entry to a basic block.
--      We use this to handle join points, where multiple branch instructions
--      target a particular label. We have to insert fixup code to make
--      the register assignments from the different sources match up.
--
type BlockAssignment freeRegs
        = BlockMap (freeRegs, RegMap Loc)

-- | Where a vreg is currently stored
--      A temporary can be marked as living in both a register and memory
--      (InBoth), for example if it was recently loaded from a spill location.
--      This makes it cheap to spill (no save instruction required), but we
--      have to be careful to turn this into InReg if the value in the
--      register is changed.

--      This is also useful when a temporary is about to be clobbered.  We
--      save it in a spill location, but mark it as InBoth because the current
--      instruction might still want to read it.
--
data Loc
        -- | vreg is in a register
        = InReg   !RealReg

        -- | vreg is held in a stack slot
        | InMem   {-# UNPACK #-}  !StackSlot

        -- | vreg is held in both a register and a stack slot
        | InBoth   !RealReg
                   {-# UNPACK #-} !StackSlot
        deriving (Eq, Show, Ord)

instance Outputable Loc where
        ppr l = text (show l)

-- | Get the reg numbers stored in this Loc.
regsOfLoc :: Loc -> [RealReg]
regsOfLoc (InReg r)    = [r]
regsOfLoc (InBoth r _) = [r]
regsOfLoc (InMem _)    = []

-- | Reasons why instructions might be inserted by the spiller.
--      Used when generating stats for -ddrop-asm-stats.
--
data SpillReason
        -- | vreg was spilled to a slot so we could use its
        --      current hreg for another vreg
        = SpillAlloc    !Unique

        -- | vreg was moved because its hreg was clobbered
        | SpillClobber  !Unique

        -- | vreg was loaded from a spill slot
        | SpillLoad     !Unique

        -- | reg-reg move inserted during join to targets
        | SpillJoinRR   !Unique

        -- | reg-mem move inserted during join to targets
        | SpillJoinRM   !Unique

-- | Used to carry interesting stats out of the register allocator.
data RegAllocStats
        = RegAllocStats
        { ra_spillInstrs        :: UniqFM [Int]
        , ra_fixupList     :: [(BlockId,BlockId,BlockId)]
        -- ^ (from,fixup,to) : We inserted fixup code between from and to
        }

-- | The register allocator state
data RA_State freeRegs
        = RA_State

        {
        -- | the current mapping from basic blocks to
        --      the register assignments at the beginning of that block.
          ra_blockassig :: BlockAssignment freeRegs

        -- | free machine registers
        , ra_freeregs   :: !freeRegs

        -- | assignment of temps to locations
        , ra_assig      :: RegMap Loc

        -- | current stack delta
        , ra_delta      :: Int

        -- | free stack slots for spilling
        , ra_stack      :: StackMap

        -- | unique supply for generating names for join point fixup blocks.
        , ra_us         :: UniqSupply

        -- | Record why things were spilled, for -ddrop-asm-stats.
        --      Just keep a list here instead of a map of regs -> reasons.
        --      We don't want to slow down the allocator if we're not going to emit the stats.
        , ra_spills     :: [SpillReason]
        , ra_DynFlags   :: DynFlags

        -- | (from,fixup,to) : We inserted fixup code between from and to
        , ra_fixups     :: [(BlockId,BlockId,BlockId)] }


{-# LANGUAGE CPP #-}

module RegAlloc.Linear.FreeRegs (
    FR(..),
    maxSpillSlots
)

#include "HsVersions.h"

where

import GhcPrelude

import Reg
import RegClass

import DynFlags
import Panic
import Platform

-- -----------------------------------------------------------------------------
-- The free register set
-- This needs to be *efficient*
-- Here's an inefficient 'executable specification' of the FreeRegs data type:
--
--      type FreeRegs = [RegNo]
--      noFreeRegs = 0
--      releaseReg n f = if n `elem` f then f else (n : f)
--      initFreeRegs = allocatableRegs
--      getFreeRegs cls f = filter ( (==cls) . regClass . RealReg ) f
--      allocateReg f r = filter (/= r) f

import qualified RegAlloc.Linear.PPC.FreeRegs    as PPC
import qualified RegAlloc.Linear.SPARC.FreeRegs  as SPARC
import qualified RegAlloc.Linear.X86.FreeRegs    as X86
import qualified RegAlloc.Linear.X86_64.FreeRegs as X86_64

import qualified PPC.Instr
import qualified SPARC.Instr
import qualified X86.Instr

class Show freeRegs => FR freeRegs where
    frAllocateReg :: Platform -> RealReg -> freeRegs -> freeRegs
    frGetFreeRegs :: Platform -> RegClass -> freeRegs -> [RealReg]
    frInitFreeRegs :: Platform -> freeRegs
    frReleaseReg :: Platform -> RealReg -> freeRegs -> freeRegs

instance FR X86.FreeRegs where
    frAllocateReg  = \_ -> X86.allocateReg
    frGetFreeRegs  = X86.getFreeRegs
    frInitFreeRegs = X86.initFreeRegs
    frReleaseReg   = \_ -> X86.releaseReg

instance FR X86_64.FreeRegs where
    frAllocateReg  = \_ -> X86_64.allocateReg
    frGetFreeRegs  = X86_64.getFreeRegs
    frInitFreeRegs = X86_64.initFreeRegs
    frReleaseReg   = \_ -> X86_64.releaseReg

instance FR PPC.FreeRegs where
    frAllocateReg  = \_ -> PPC.allocateReg
    frGetFreeRegs  = \_ -> PPC.getFreeRegs
    frInitFreeRegs = PPC.initFreeRegs
    frReleaseReg   = \_ -> PPC.releaseReg

instance FR SPARC.FreeRegs where
    frAllocateReg  = SPARC.allocateReg
    frGetFreeRegs  = \_ -> SPARC.getFreeRegs
    frInitFreeRegs = SPARC.initFreeRegs
    frReleaseReg   = SPARC.releaseReg

maxSpillSlots :: DynFlags -> Int
maxSpillSlots dflags
              = case platformArch (targetPlatform dflags) of
                ArchX86       -> X86.Instr.maxSpillSlots dflags
                ArchX86_64    -> X86.Instr.maxSpillSlots dflags
                ArchPPC       -> PPC.Instr.maxSpillSlots dflags
                ArchSPARC     -> SPARC.Instr.maxSpillSlots dflags
                ArchSPARC64   -> panic "maxSpillSlots ArchSPARC64"
                ArchARM _ _ _ -> panic "maxSpillSlots ArchARM"
                ArchARM64     -> panic "maxSpillSlots ArchARM64"
                ArchPPC_64 _  -> PPC.Instr.maxSpillSlots dflags
                ArchAlpha     -> panic "maxSpillSlots ArchAlpha"
                ArchMipseb    -> panic "maxSpillSlots ArchMipseb"
                ArchMipsel    -> panic "maxSpillSlots ArchMipsel"
                ArchJavaScript-> panic "maxSpillSlots ArchJavaScript"
                ArchUnknown   -> panic "maxSpillSlots ArchUnknown"


-- | Handles joining of a jump instruction to its targets.

--      The first time we encounter a jump to a particular basic block, we
--      record the assignment of temporaries.  The next time we encounter a
--      jump to the same block, we compare our current assignment to the
--      stored one.  They might be different if spilling has occurred in one
--      branch; so some fixup code will be required to match up the assignments.
--
module RegAlloc.Linear.JoinToTargets (joinToTargets) where

import GhcPrelude

import RegAlloc.Linear.State
import RegAlloc.Linear.Base
import RegAlloc.Linear.FreeRegs
import RegAlloc.Liveness
import Instruction
import Reg

import BlockId
import Hoopl.Collections
import Digraph
import DynFlags
import Outputable
import Unique
import UniqFM
import UniqSet

-- | For a jump instruction at the end of a block, generate fixup code so its
--      vregs are in the correct regs for its destination.
--
joinToTargets
        :: (FR freeRegs, Instruction instr, Outputable instr)
        => BlockMap RegSet              -- ^ maps the unique of the blockid to the set of vregs
                                        --      that are known to be live on the entry to each block.

        -> BlockId                      -- ^ id of the current block
        -> instr                        -- ^ branch instr on the end of the source block.

        -> RegM freeRegs ([NatBasicBlock instr] -- fresh blocks of fixup code.
                         , instr)               -- the original branch
                                                -- instruction, but maybe
                                                -- patched to jump
                                                -- to a fixup block first.

joinToTargets block_live id instr

        -- we only need to worry about jump instructions.
        | not $ isJumpishInstr instr
        = return ([], instr)

        | otherwise
        = joinToTargets' block_live [] id instr (jumpDestsOfInstr instr)

-----
joinToTargets'
        :: (FR freeRegs, Instruction instr, Outputable instr)
        => BlockMap RegSet              -- ^ maps the unique of the blockid to the set of vregs
                                        --      that are known to be live on the entry to each block.

        -> [NatBasicBlock instr]        -- ^ acc blocks of fixup code.

        -> BlockId                      -- ^ id of the current block
        -> instr                        -- ^ branch instr on the end of the source block.

        -> [BlockId]                    -- ^ branch destinations still to consider.

        -> RegM freeRegs ([NatBasicBlock instr], instr)

-- no more targets to consider. all done.
joinToTargets' _          new_blocks _ instr []
        = return (new_blocks, instr)

-- handle a branch target.
joinToTargets' block_live new_blocks block_id instr (dest:dests)
 = do
        -- get the map of where the vregs are stored on entry to each basic block.
        block_assig     <- getBlockAssigR

        -- get the assignment on entry to the branch instruction.
        assig           <- getAssigR

        -- adjust the current assignment to remove any vregs that are not live
        -- on entry to the destination block.
        let Just live_set       = mapLookup dest block_live
        let still_live uniq _   = uniq `elemUniqSet_Directly` live_set
        let adjusted_assig      = filterUFM_Directly still_live assig

        -- and free up those registers which are now free.
        let to_free =
                [ r     | (reg, loc) <- nonDetUFMToList assig
                        -- This is non-deterministic but we do not
                        -- currently support deterministic code-generation.
                        -- See Note [Unique Determinism and code generation]
                        , not (elemUniqSet_Directly reg live_set)
                        , r          <- regsOfLoc loc ]

        case mapLookup dest block_assig of
         Nothing
          -> joinToTargets_first
                        block_live new_blocks block_id instr dest dests
                        block_assig adjusted_assig to_free

         Just (_, dest_assig)
          -> joinToTargets_again
                        block_live new_blocks block_id instr dest dests
                        adjusted_assig dest_assig

-- this is the first time we jumped to this block.
joinToTargets_first :: (FR freeRegs, Instruction instr, Outputable instr)
                    => BlockMap RegSet
                    -> [NatBasicBlock instr]
                    -> BlockId
                    -> instr
                    -> BlockId
                    -> [BlockId]
                    -> BlockAssignment freeRegs
                    -> RegMap Loc
                    -> [RealReg]
                    -> RegM freeRegs ([NatBasicBlock instr], instr)
joinToTargets_first block_live new_blocks block_id instr dest dests
        block_assig src_assig
        to_free

 = do   dflags <- getDynFlags
        let platform = targetPlatform dflags

        -- free up the regs that are not live on entry to this block.
        freeregs        <- getFreeRegsR
        let freeregs' = foldl' (flip $ frReleaseReg platform) freeregs to_free

        -- remember the current assignment on entry to this block.
        setBlockAssigR (mapInsert dest (freeregs', src_assig) block_assig)

        joinToTargets' block_live new_blocks block_id instr dests

-- we've jumped to this block before
joinToTargets_again :: (Instruction instr, FR freeRegs, Outputable instr)
                    => BlockMap RegSet
                    -> [NatBasicBlock instr]
                    -> BlockId
                    -> instr
                    -> BlockId
                    -> [BlockId]
                    -> UniqFM Loc
                    -> UniqFM Loc
                    -> RegM freeRegs ([NatBasicBlock instr], instr)
joinToTargets_again
    block_live new_blocks block_id instr dest dests
    src_assig dest_assig

        -- the assignments already match, no problem.
        | nonDetUFMToList dest_assig == nonDetUFMToList src_assig
        -- This is non-deterministic but we do not
        -- currently support deterministic code-generation.
        -- See Note [Unique Determinism and code generation]
        = joinToTargets' block_live new_blocks block_id instr dests

        -- assignments don't match, need fixup code
        | otherwise
        = do

                -- make a graph of what things need to be moved where.
                let graph = makeRegMovementGraph src_assig dest_assig

                -- look for cycles in the graph. This can happen if regs need to be swapped.
                -- Note that we depend on the fact that this function does a
                --      bottom up traversal of the tree-like portions of the graph.
                --
                --  eg, if we have
                --      R1 -> R2 -> R3
                --
                --  ie move value in R1 to R2 and value in R2 to R3.
                --
                -- We need to do the R2 -> R3 move before R1 -> R2.
                --
                let sccs  = stronglyConnCompFromEdgedVerticesOrdR graph

              -- debugging
                {-
                pprTrace
                        ("joinToTargets: making fixup code")
                        (vcat   [ text "        in block: "     <> ppr block_id
                                , text " jmp instruction: "     <> ppr instr
                                , text "  src assignment: "     <> ppr src_assig
                                , text " dest assignment: "     <> ppr dest_assig
                                , text "  movement graph: "     <> ppr graph
                                , text "   sccs of graph: "     <> ppr sccs
                                , text ""])
                        (return ())
                -}
                delta           <- getDeltaR
                fixUpInstrs_    <- mapM (handleComponent delta instr) sccs
                let fixUpInstrs = concat fixUpInstrs_

                -- make a new basic block containing the fixup code.
                --      A the end of the current block we will jump to the fixup one,
                --      then that will jump to our original destination.
                fixup_block_id <- mkBlockId <$> getUniqueR
                let block = BasicBlock fixup_block_id
                                $ fixUpInstrs ++ mkJumpInstr dest

                -- if we didn't need any fixups, then don't include the block
                case fixUpInstrs of
                 []     -> joinToTargets' block_live new_blocks block_id instr dests

                 -- patch the original branch instruction so it goes to our
                 --     fixup block instead.
                 _      -> let  instr'  =  patchJumpInstr instr
                                            (\bid -> if bid == dest
                                                        then fixup_block_id
                                                        else bid) -- no change!

                           in do
                                {- --debugging
                                pprTrace "FixUpEdge info:"
                                    (
                                    text "inBlock:" <> ppr block_id $$
                                    text "instr:" <> ppr instr $$
                                    text "instr':" <> ppr instr' $$
                                    text "fixup_block_id':" <>
                                        ppr fixup_block_id $$
                                    text "dest:" <> ppr dest
                                    ) (return ())
                                -}
                                recordFixupBlock block_id fixup_block_id dest
                                joinToTargets' block_live (block : new_blocks)
                                               block_id instr' dests

-- | Construct a graph of register\/spill movements.
--
--      Cyclic components seem to occur only very rarely.
--
--      We cut some corners by not handling memory-to-memory moves.
--      This shouldn't happen because every temporary gets its own stack slot.
--
makeRegMovementGraph :: RegMap Loc -> RegMap Loc -> [Node Loc Unique]
makeRegMovementGraph adjusted_assig dest_assig
 = [ node       | (vreg, src) <- nonDetUFMToList adjusted_assig
                    -- This is non-deterministic but we do not
                    -- currently support deterministic code-generation.
                    -- See Note [Unique Determinism and code generation]
                    -- source reg might not be needed at the dest:
                , Just loc <- [lookupUFM_Directly dest_assig vreg]
                , node <- expandNode vreg src loc ]

-- | Expand out the destination, so InBoth destinations turn into
--      a combination of InReg and InMem.

--      The InBoth handling is a little tricky here.  If the destination is
--      InBoth, then we must ensure that the value ends up in both locations.
--      An InBoth  destination must conflict with an InReg or InMem source, so
--      we expand an InBoth destination as necessary.
--
--      An InBoth source is slightly different: we only care about the register
--      that the source value is in, so that we can move it to the destinations.
--
expandNode
        :: a
        -> Loc                  -- ^ source of move
        -> Loc                  -- ^ destination of move
        -> [Node Loc a ]

expandNode vreg loc@(InReg src) (InBoth dst mem)
        | src == dst = [DigraphNode vreg loc [InMem mem]]
        | otherwise  = [DigraphNode vreg loc [InReg dst, InMem mem]]

expandNode vreg loc@(InMem src) (InBoth dst mem)
        | src == mem = [DigraphNode vreg loc [InReg dst]]
        | otherwise  = [DigraphNode vreg loc [InReg dst, InMem mem]]

expandNode _        (InBoth _ src) (InMem dst)
        | src == dst = [] -- guaranteed to be true

expandNode _        (InBoth src _) (InReg dst)
        | src == dst = []

expandNode vreg     (InBoth src _) dst
        = expandNode vreg (InReg src) dst

expandNode vreg src dst
        | src == dst = []
        | otherwise  = [DigraphNode vreg src [dst]]

-- | Generate fixup code for a particular component in the move graph
--      This component tells us what values need to be moved to what
--      destinations. We have eliminated any possibility of single-node
--      cycles in expandNode above.
--
handleComponent
        :: Instruction instr
        => Int -> instr -> SCC (Node Loc Unique)
        -> RegM freeRegs [instr]

-- If the graph is acyclic then we won't get the swapping problem below.
--      In this case we can just do the moves directly, and avoid having to
--      go via a spill slot.
--
handleComponent delta _  (AcyclicSCC (DigraphNode vreg src dsts))
        = mapM (makeMove delta vreg src) dsts

-- Handle some cyclic moves.
--      This can happen if we have two regs that need to be swapped.
--      eg:
--           vreg   source loc   dest loc
--          (vreg1, InReg r1,    [InReg r2])
--          (vreg2, InReg r2,    [InReg r1])
--
--      To avoid needing temp register, we just spill all the source regs, then
--      reaload them into their destination regs.
--
--      Note that we can not have cycles that involve memory locations as
--      sources as single destination because memory locations (stack slots)
--      are allocated exclusively for a virtual register and therefore can not
--      require a fixup.
--
handleComponent delta instr
        (CyclicSCC ((DigraphNode vreg (InReg sreg) ((InReg dreg: _))) : rest))
        -- dest list may have more than one element, if the reg is also InMem.
 = do
        -- spill the source into its slot
        (instrSpill, slot)
                        <- spillR (RegReal sreg) vreg

        -- reload into destination reg
        instrLoad       <- loadR (RegReal dreg) slot

        remainingFixUps <- mapM (handleComponent delta instr)
                                (stronglyConnCompFromEdgedVerticesOrdR rest)

        -- make sure to do all the reloads after all the spills,
        --      so we don't end up clobbering the source values.
        return ([instrSpill] ++ concat remainingFixUps ++ [instrLoad])

handleComponent _ _ (CyclicSCC _)
 = panic "Register Allocator: handleComponent cyclic"

-- | Move a vreg between these two locations.
--
makeMove
    :: Instruction instr
    => Int      -- ^ current C stack delta.
    -> Unique   -- ^ unique of the vreg that we're moving.
    -> Loc      -- ^ source location.
    -> Loc      -- ^ destination location.
    -> RegM freeRegs instr  -- ^ move instruction.

makeMove delta vreg src dst
 = do dflags <- getDynFlags
      let platform = targetPlatform dflags

      case (src, dst) of
          (InReg s, InReg d) ->
              do recordSpill (SpillJoinRR vreg)
                 return $ mkRegRegMoveInstr platform (RegReal s) (RegReal d)
          (InMem s, InReg d) ->
              do recordSpill (SpillJoinRM vreg)
                 return $ mkLoadInstr dflags (RegReal d) delta s
          (InReg s, InMem d) ->
              do recordSpill (SpillJoinRM vreg)
                 return $ mkSpillInstr dflags (RegReal s) delta d
          _ ->
              -- we don't handle memory to memory moves.
              -- they shouldn't happen because we don't share
              -- stack slots between vregs.
              panic ("makeMove " ++ show vreg ++ " (" ++ show src ++ ") ("
                  ++ show dst ++ ")"
                  ++ " we don't handle mem->mem moves.")