{-# LANGUAGE BangPatterns, CPP, ScopedTypeVariables #-}
module RegAlloc.Linear.Main (
        regAlloc,
        module  RegAlloc.Linear.Base,
        module  RegAlloc.Linear.Stats
  ) where
#include "HsVersions.h"
import GhcPrelude
import RegAlloc.Linear.State
import RegAlloc.Linear.Base
import RegAlloc.Linear.StackMap
import RegAlloc.Linear.FreeRegs
import RegAlloc.Linear.Stats
import RegAlloc.Linear.JoinToTargets
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 TargetReg
import RegAlloc.Liveness
import Instruction
import Reg
import BlockId
import Hoopl.Collections
import Cmm hiding (RegSet)
import Digraph
import DynFlags
import Unique
import UniqSet
import UniqFM
import UniqSupply
import Outputable
import Platform
import Data.Maybe
import Data.List
import Control.Monad
regAlloc
        :: (Outputable instr, Instruction instr)
        => DynFlags
        -> LiveCmmDecl statics instr
        -> UniqSM ( NatCmmDecl statics instr
                  , Maybe Int  
                               
                  , Maybe RegAllocStats
                  )
regAlloc _ (CmmData sec d)
        = return
                ( CmmData sec d
                , Nothing
                , Nothing )
regAlloc _ (CmmProc (LiveInfo info _ _ _) lbl live [])
        = return ( CmmProc info lbl live (ListGraph [])
                 , Nothing
                 , Nothing )
regAlloc dflags (CmmProc static lbl live sccs)
        | LiveInfo info entry_ids@(first_id:_) block_live _ <- static
        = do
                
                (final_blocks, stats, stack_use)
                        <- linearRegAlloc dflags entry_ids block_live sccs
                
                
                let ((first':_), rest')
                                = partition ((== first_id) . blockId) final_blocks
                let max_spill_slots = maxSpillSlots dflags
                    extra_stack
                      | stack_use > max_spill_slots
                      = Just (stack_use - max_spill_slots)
                      | otherwise
                      = Nothing
                return  ( CmmProc info lbl live (ListGraph (first' : rest'))
                        , extra_stack
                        , Just stats)
regAlloc _ (CmmProc _ _ _ _)
        = panic "RegAllocLinear.regAlloc: no match"
linearRegAlloc
        :: (Outputable instr, Instruction instr)
        => DynFlags
        -> [BlockId] 
        -> BlockMap RegSet
              
        -> [SCC (LiveBasicBlock instr)]
              
        -> UniqSM ([NatBasicBlock instr], RegAllocStats, Int)
linearRegAlloc dflags entry_ids block_live sccs
 = case platformArch platform of
      ArchX86        -> go $ (frInitFreeRegs platform :: X86.FreeRegs)
      ArchX86_64     -> go $ (frInitFreeRegs platform :: X86_64.FreeRegs)
      ArchSPARC      -> go $ (frInitFreeRegs platform :: SPARC.FreeRegs)
      ArchSPARC64    -> panic "linearRegAlloc ArchSPARC64"
      ArchPPC        -> go $ (frInitFreeRegs platform :: PPC.FreeRegs)
      ArchARM _ _ _  -> panic "linearRegAlloc ArchARM"
      ArchARM64      -> panic "linearRegAlloc ArchARM64"
      ArchPPC_64 _   -> go $ (frInitFreeRegs platform :: PPC.FreeRegs)
      ArchAlpha      -> panic "linearRegAlloc ArchAlpha"
      ArchMipseb     -> panic "linearRegAlloc ArchMipseb"
      ArchMipsel     -> panic "linearRegAlloc ArchMipsel"
      ArchJavaScript -> panic "linearRegAlloc ArchJavaScript"
      ArchUnknown    -> panic "linearRegAlloc ArchUnknown"
 where
  go f = linearRegAlloc' dflags f entry_ids block_live sccs
  platform = targetPlatform dflags
linearRegAlloc'
        :: (FR freeRegs, Outputable instr, Instruction instr)
        => DynFlags
        -> freeRegs
        -> [BlockId]                    
        -> BlockMap RegSet              
        -> [SCC (LiveBasicBlock instr)] 
        -> UniqSM ([NatBasicBlock instr], RegAllocStats, Int)
linearRegAlloc' dflags initFreeRegs entry_ids block_live sccs
 = do   us      <- getUniqueSupplyM
        let (_, stack, stats, blocks) =
                runR dflags mapEmpty initFreeRegs emptyRegMap (emptyStackMap dflags) us
                    $ linearRA_SCCs entry_ids block_live [] sccs
        return  (blocks, stats, getStackUse stack)
linearRA_SCCs :: (FR freeRegs, Instruction instr, Outputable instr)
              => [BlockId]
              -> BlockMap RegSet
              -> [NatBasicBlock instr]
              -> [SCC (LiveBasicBlock instr)]
              -> RegM freeRegs [NatBasicBlock instr]
linearRA_SCCs _ _ blocksAcc []
        = return $ reverse blocksAcc
linearRA_SCCs entry_ids block_live blocksAcc (AcyclicSCC block : sccs)
 = do   blocks' <- processBlock block_live block
        linearRA_SCCs entry_ids block_live
                ((reverse blocks') ++ blocksAcc)
                sccs
linearRA_SCCs entry_ids block_live blocksAcc (CyclicSCC blocks : sccs)
 = do
        blockss' <- process entry_ids block_live blocks [] (return []) False
        linearRA_SCCs entry_ids block_live
                (reverse (concat blockss') ++ blocksAcc)
                sccs
process :: (FR freeRegs, Instruction instr, Outputable instr)
        => [BlockId]
        -> BlockMap RegSet
        -> [GenBasicBlock (LiveInstr instr)]
        -> [GenBasicBlock (LiveInstr instr)]
        -> [[NatBasicBlock instr]]
        -> Bool
        -> RegM freeRegs [[NatBasicBlock instr]]
process _ _ [] []         accum _
        = return $ reverse accum
process entry_ids block_live [] next_round accum madeProgress
        | not madeProgress
          
        = return $ reverse accum
        | otherwise
        = process entry_ids block_live
                  next_round [] accum False
process entry_ids block_live (b@(BasicBlock id _) : blocks)
        next_round accum madeProgress
 = do
        block_assig <- getBlockAssigR
        if isJust (mapLookup id block_assig)
             || id `elem` entry_ids
         then do
                b'  <- processBlock block_live b
                process entry_ids block_live blocks
                        next_round (b' : accum) True
         else   process entry_ids block_live blocks
                        (b : next_round) accum madeProgress
processBlock
        :: (FR freeRegs, Outputable instr, Instruction instr)
        => BlockMap RegSet              
        -> LiveBasicBlock instr         
        -> RegM freeRegs [NatBasicBlock instr]   
processBlock block_live (BasicBlock id instrs)
 = do   initBlock id block_live
        (instrs', fixups)
                <- linearRA block_live [] [] id instrs
        return  $ BasicBlock id instrs' : fixups
initBlock :: FR freeRegs
          => BlockId -> BlockMap RegSet -> RegM freeRegs ()
initBlock id block_live
 = do   dflags <- getDynFlags
        let platform = targetPlatform dflags
        block_assig     <- getBlockAssigR
        case mapLookup id block_assig of
                
                
                
                
                
                Nothing
                 -> do  
                        case mapLookup id block_live of
                          Nothing ->
                            setFreeRegsR    (frInitFreeRegs platform)
                          Just live ->
                            setFreeRegsR $ foldl' (flip $ frAllocateReg platform) (frInitFreeRegs platform)
                                                  [ r | RegReal r <- nonDetEltsUniqSet live ]
                            
                        setAssigR       emptyRegMap
                
                Just (freeregs, assig)
                 -> do  setFreeRegsR    freeregs
                        setAssigR       assig
linearRA
        :: (FR freeRegs, Outputable instr, Instruction instr)
        => BlockMap RegSet                      
        -> [instr]                              
        -> [NatBasicBlock instr]                
        -> BlockId                              
        -> [LiveInstr instr]                    
        -> RegM freeRegs
                ( [instr]                       
                , [NatBasicBlock instr])        
linearRA _          accInstr accFixup _ []
        = return
                ( reverse accInstr              
                , accFixup)                     
linearRA block_live accInstr accFixups id (instr:instrs)
 = do
        (accInstr', new_fixups) <- raInsn block_live accInstr id instr
        linearRA block_live accInstr' (new_fixups ++ accFixups) id instrs
raInsn
        :: (FR freeRegs, Outputable instr, Instruction instr)
        => BlockMap RegSet                      
        -> [instr]                              
        -> BlockId                              
        -> LiveInstr instr                      
        -> RegM freeRegs
                ( [instr]                       
                , [NatBasicBlock instr])        
raInsn _     new_instrs _ (LiveInstr ii Nothing)
        | Just n        <- takeDeltaInstr ii
        = do    setDeltaR n
                return (new_instrs, [])
raInsn _     new_instrs _ (LiveInstr ii@(Instr i) Nothing)
        | isMetaInstr ii
        = return (i : new_instrs, [])
raInsn block_live new_instrs id (LiveInstr (Instr instr) (Just live))
 = do
    assig    <- getAssigR
    
    
    
    
    
    
    
    case takeRegRegMoveInstr instr of
        Just (src,dst)  | src `elementOfUniqSet` (liveDieRead live),
                          isVirtualReg dst,
                          not (dst `elemUFM` assig),
                          isRealReg src || isInReg src assig -> do
           case src of
              (RegReal rr) -> setAssigR (addToUFM assig dst (InReg rr))
                
                
                
              _virt -> case lookupUFM assig src of
                         Nothing -> panic "raInsn"
                         Just loc ->
                           setAssigR (addToUFM (delFromUFM assig src) dst loc)
           
          
           return (new_instrs, [])
        _ -> genRaInsn block_live new_instrs id instr
                        (nonDetEltsUniqSet $ liveDieRead live)
                        (nonDetEltsUniqSet $ liveDieWrite live)
                        
raInsn _ _ _ instr
        = pprPanic "raInsn" (text "no match for:" <> ppr instr)
isInReg :: Reg -> RegMap Loc -> Bool
isInReg src assig | Just (InReg _) <- lookupUFM assig src = True
                  | otherwise = False
genRaInsn :: (FR freeRegs, Instruction instr, Outputable instr)
          => BlockMap RegSet
          -> [instr]
          -> BlockId
          -> instr
          -> [Reg]
          -> [Reg]
          -> RegM freeRegs ([instr], [NatBasicBlock instr])
genRaInsn block_live new_instrs block_id instr r_dying w_dying = do
  dflags <- getDynFlags
  let platform = targetPlatform dflags
  case regUsageOfInstr platform instr of { RU read written ->
    do
    let real_written    = [ rr  | (RegReal     rr) <- written ]
    let virt_written    = [ vr  | (RegVirtual  vr) <- written ]
    
    
    
    let virt_read       = nub [ vr      | (RegVirtual vr) <- read ]
    
    
    (r_spills, r_allocd) <-
        allocateRegsAndSpill True virt_read [] [] virt_read
    
    clobber_saves <- saveClobberedTemps real_written r_dying
    
    
    
    
    (fixup_blocks, adjusted_instr)
        <- joinToTargets block_live block_id instr
    
    
    
    
    
    
    releaseRegs r_dying
    
    clobberRegs real_written
    
    (w_spills, w_allocd) <-
        allocateRegsAndSpill False virt_written [] [] virt_written
    
    
    releaseRegs w_dying
    let
        
        patch_map
                = listToUFM
                        [ (t, RegReal r)
                                | (t, r) <- zip virt_read    r_allocd
                                         ++ zip virt_written w_allocd ]
        patched_instr
                = patchRegsOfInstr adjusted_instr patchLookup
        patchLookup x
                = case lookupUFM patch_map x of
                        Nothing -> x
                        Just y  -> y
    
    
    
    
    
    let squashed_instr  = case takeRegRegMoveInstr patched_instr of
                                Just (src, dst)
                                 | src == dst   -> []
                                _               -> [patched_instr]
    let code = squashed_instr ++ w_spills ++ reverse r_spills
                ++ clobber_saves ++ new_instrs
    return (code, fixup_blocks)
  }
releaseRegs :: FR freeRegs => [Reg] -> RegM freeRegs ()
releaseRegs regs = do
  dflags <- getDynFlags
  let platform = targetPlatform dflags
  assig <- getAssigR
  free <- getFreeRegsR
  let loop assig !free [] = do setAssigR assig; setFreeRegsR free; return ()
      loop assig !free (RegReal rr : rs) = loop assig (frReleaseReg platform rr free) rs
      loop assig !free (r:rs) =
         case lookupUFM assig r of
         Just (InBoth real _) -> loop (delFromUFM assig r)
                                      (frReleaseReg platform real free) rs
         Just (InReg real)    -> loop (delFromUFM assig r)
                                      (frReleaseReg platform real free) rs
         _                    -> loop (delFromUFM assig r) free rs
  loop assig free regs
saveClobberedTemps
        :: (Instruction instr, FR freeRegs)
        => [RealReg]            
        -> [Reg]                
        -> RegM freeRegs [instr]         
                                
saveClobberedTemps [] _
        = return []
saveClobberedTemps clobbered dying
 = do
        assig   <- getAssigR
        let to_spill
                = [ (temp,reg)
                        | (temp, InReg reg) <- nonDetUFMToList assig
                        
                        
                        
                        , any (realRegsAlias reg) clobbered
                        , temp `notElem` map getUnique dying  ]
        (instrs,assig') <- clobber assig [] to_spill
        setAssigR assig'
        return instrs
   where
     clobber assig instrs []
            = return (instrs, assig)
     clobber assig instrs ((temp, reg) : rest)
       = do dflags <- getDynFlags
            let platform = targetPlatform dflags
            freeRegs <- getFreeRegsR
            let regclass = targetClassOfRealReg platform reg
                freeRegs_thisClass = frGetFreeRegs platform regclass freeRegs
            case filter (`notElem` clobbered) freeRegs_thisClass of
              
              
              
              (my_reg : _) -> do
                  setFreeRegsR (frAllocateReg platform my_reg freeRegs)
                  let new_assign = addToUFM assig temp (InReg my_reg)
                  let instr = mkRegRegMoveInstr platform
                                  (RegReal reg) (RegReal my_reg)
                  clobber new_assign (instr : instrs) rest
              
              [] -> do
                  (spill, slot)   <- spillR (RegReal reg) temp
                  
                  recordSpill (SpillClobber temp)
                  let new_assign  = addToUFM assig temp (InBoth reg slot)
                  clobber new_assign (spill : instrs) rest
clobberRegs :: FR freeRegs => [RealReg] -> RegM freeRegs ()
clobberRegs []
        = return ()
clobberRegs clobbered
 = do   dflags <- getDynFlags
        let platform = targetPlatform dflags
        freeregs        <- getFreeRegsR
        setFreeRegsR $! foldl' (flip $ frAllocateReg platform) freeregs clobbered
        assig           <- getAssigR
        setAssigR $! clobber assig (nonDetUFMToList assig)
          
          
          
   where
        
        
        
        
        
        clobber assig []
                = assig
        clobber assig ((temp, InBoth reg slot) : rest)
                | any (realRegsAlias reg) clobbered
                = clobber (addToUFM assig temp (InMem slot)) rest
        clobber assig (_:rest)
                = clobber assig rest
data SpillLoc = ReadMem StackSlot  
              | WriteNew           
              | WriteMem           
allocateRegsAndSpill
        :: (FR freeRegs, Outputable instr, Instruction instr)
        => Bool                 
        -> [VirtualReg]         
        -> [instr]              
        -> [RealReg]            
        -> [VirtualReg]         
        -> RegM freeRegs ( [instr] , [RealReg])
allocateRegsAndSpill _       _    spills alloc []
        = return (spills, reverse alloc)
allocateRegsAndSpill reading keep spills alloc (r:rs)
 = do   assig <- getAssigR
        let doSpill = allocRegsAndSpill_spill reading keep spills alloc r rs assig
        case lookupUFM assig r of
                
                Just (InReg my_reg) ->
                        allocateRegsAndSpill reading keep spills (my_reg:alloc) rs
                
                
                
                
                
                Just (InBoth my_reg _)
                 -> do  when (not reading) (setAssigR (addToUFM assig r (InReg my_reg)))
                        allocateRegsAndSpill reading keep spills (my_reg:alloc) rs
                
                Just (InMem slot) | reading   -> doSpill (ReadMem slot)
                                  | otherwise -> doSpill WriteMem
                Nothing | reading   ->
                   pprPanic "allocateRegsAndSpill: Cannot read from uninitialized register" (ppr r)
                   
                   
                   
                   
                   
                   
                        | otherwise -> doSpill WriteNew
allocRegsAndSpill_spill :: (FR freeRegs, Instruction instr, Outputable instr)
                        => Bool
                        -> [VirtualReg]
                        -> [instr]
                        -> [RealReg]
                        -> VirtualReg
                        -> [VirtualReg]
                        -> UniqFM Loc
                        -> SpillLoc
                        -> RegM freeRegs ([instr], [RealReg])
allocRegsAndSpill_spill reading keep spills alloc r rs assig spill_loc
 = do   dflags <- getDynFlags
        let platform = targetPlatform dflags
        freeRegs                <- getFreeRegsR
        let freeRegs_thisClass  = frGetFreeRegs platform (classOfVirtualReg r) freeRegs
        case freeRegs_thisClass of
         
         (my_reg : _) ->
           do   spills'   <- loadTemp r spill_loc my_reg spills
                setAssigR       (addToUFM assig r $! newLocation spill_loc my_reg)
                setFreeRegsR $  frAllocateReg platform my_reg freeRegs
                allocateRegsAndSpill reading keep spills' (my_reg : alloc) rs
          
         [] ->
           do   let inRegOrBoth (InReg _) = True
                    inRegOrBoth (InBoth _ _) = True
                    inRegOrBoth _ = False
                let candidates' =
                      flip delListFromUFM keep $
                      filterUFM inRegOrBoth $
                      assig
                      
                      
                      
                let candidates = nonDetUFMToList candidates'
                
                let candidates_inBoth
                        = [ (temp, reg, mem)
                          | (temp, InBoth reg mem) <- candidates
                          , targetClassOfRealReg platform reg == classOfVirtualReg r ]
                
                
                let candidates_inReg
                        = [ (temp, reg)
                          | (temp, InReg reg) <- candidates
                          , targetClassOfRealReg platform reg == classOfVirtualReg r ]
                let result
                        
                        
                        | (temp, my_reg, slot) : _      <- candidates_inBoth
                        = do    spills' <- loadTemp r spill_loc my_reg spills
                                let assig1  = addToUFM assig temp (InMem slot)
                                let assig2  = addToUFM assig1 r $! newLocation spill_loc my_reg
                                setAssigR assig2
                                allocateRegsAndSpill reading keep spills' (my_reg:alloc) rs
                        
                        
                        | (temp_to_push_out, (my_reg :: RealReg)) : _
                                        <- candidates_inReg
                        = do
                                (spill_insn, slot) <- spillR (RegReal my_reg) temp_to_push_out
                                let spill_store  = (if reading then id else reverse)
                                                        [ 
                                                           spill_insn ]
                                
                                recordSpill (SpillAlloc temp_to_push_out)
                                
                                let assig1  = addToUFM assig temp_to_push_out   (InMem slot)
                                let assig2  = addToUFM assig1 r                 $! newLocation spill_loc my_reg
                                setAssigR assig2
                                
                                spills' <- loadTemp r spill_loc my_reg spills
                                allocateRegsAndSpill reading keep
                                        (spill_store ++ spills')
                                        (my_reg:alloc) rs
                        
                        | otherwise
                        = pprPanic ("RegAllocLinear.allocRegsAndSpill: no spill candidates\n")
                        $ vcat
                                [ text "allocating vreg:  " <> text (show r)
                                , text "assignment:       " <> ppr assig
                                , text "freeRegs:         " <> text (show freeRegs)
                                , text "initFreeRegs:     " <> text (show (frInitFreeRegs platform `asTypeOf` freeRegs)) ]
                result
newLocation :: SpillLoc -> RealReg -> Loc
newLocation (ReadMem slot) my_reg = InBoth my_reg slot
newLocation _ my_reg = InReg my_reg
loadTemp
        :: (Instruction instr)
        => VirtualReg   
        -> SpillLoc     
        -> RealReg      
        -> [instr]
        -> RegM freeRegs [instr]
loadTemp vreg (ReadMem slot) hreg spills
 = do
        insn <- loadR (RegReal hreg) slot
        recordSpill (SpillLoad $ getUnique vreg)
        return  $   insn : spills
loadTemp _ _ _ spills =
   return spills