ghc/compiler/nativeGen/RegisterAlloc.hs

   1 -----------------------------------------------------------------------------
   2 --
   3 -- The register allocator
   4 --
   5 -- (c) The University of Glasgow 2004
   6 --
   7 -----------------------------------------------------------------------------
   8
   9 {-
  10 The algorithm is roughly:
  11
  12   1) Compute strongly connected components of the basic block list.
  13
  14   2) Compute liveness (mapping from pseudo register to
  15      point(s) of death?).
  16
  17   3) Walk instructions in each basic block.  We keep track of
  18         (a) Free real registers (a bitmap?)
  19         (b) Current assignment of temporaries to machine registers and/or
  20             spill slots (call this the "assignment").
  21         (c) Partial mapping from basic block ids to a virt-to-loc mapping.
  22             When we first encounter a branch to a basic block,
  23             we fill in its entry in this table with the current mapping.
  24
  25      For each instruction:
  26         (a) For each real register clobbered by this instruction:
  27             If a temporary resides in it,
  28                 If the temporary is live after this instruction,
  29                     Move the temporary to another (non-clobbered & free) reg,
  30                     or spill it to memory.  Mark the temporary as residing
  31                     in both memory and a register if it was spilled (it might
  32                     need to be read by this instruction).
  33             (ToDo: this is wrong for jump instructions?)
  34
  35         (b) For each temporary *read* by the instruction:
  36             If the temporary does not have a real register allocation:
  37                 - Allocate a real register from the free list.  If
  38                   the list is empty:
  39                   - Find a temporary to spill.  Pick one that is
  40                     not used in this instruction (ToDo: not
  41                     used for a while...)
  42                   - generate a spill instruction
  43                 - If the temporary was previously spilled,
  44                   generate an instruction to read the temp from its spill loc.
  45             (optimisation: if we can see that a real register is going to
  46             be used soon, then don't use it for allocation).
  47
  48         (c) Update the current assignment
  49
  50         (d) If the intstruction is a branch:
  51               if the destination block already has a register assignment,
  52                 Generate a new block with fixup code and redirect the
  53                 jump to the new block.
  54               else,
  55                 Update the block id->assignment mapping with the current
  56                 assignment.
  57
  58         (e) Delete all register assignments for temps which are read
  59             (only) and die here.  Update the free register list.
  60
  61         (f) Mark all registers clobbered by this instruction as not free,
  62             and mark temporaries which have been spilled due to clobbering
  63             as in memory (step (a) marks then as in both mem & reg).
  64
  65         (g) For each temporary *written* by this instruction:
  66             Allocate a real register as for (b), spilling something
  67             else if necessary.
  68                 - except when updating the assignment, drop any memory
  69                   locations that the temporary was previously in, since
  70                   they will be no longer valid after this instruction.
  71
  72         (h) Delete all register assignments for temps which are
  73             written and die here (there should rarely be any).  Update
  74             the free register list.
  75
  76         (i) Rewrite the instruction with the new mapping.
  77
  78         (j) For each spilled reg known to be now dead, re-add its stack slot
  79             to the free list.
  80
  81 -}
  82
  83 module RegisterAlloc (
  84         regAlloc
  85   ) where
  86
  87 #include "HsVersions.h"
  88 #include "../includes/ghcconfig.h"
  89
  90 import PprMach
  91 import MachRegs
  92 import MachInstrs
  93 import RegAllocInfo
  94 import Cmm
  95
  96 import Digraph
  97 import Unique           ( Uniquable(..), Unique, getUnique )
  98 import UniqSet
  99 import UniqFM
 100 import Outputable
 101
 102 #ifndef DEBUG
 103 import Maybe            ( fromJust )
 104 #endif
 105 import List             ( nub, partition )
 106 import Monad            ( when )
 107 import DATA_WORD
 108 import DATA_BITS
 109
 110 -- -----------------------------------------------------------------------------
 111 -- Some useful types
 112
 113 type RegSet = UniqSet Reg
 114
 115 type RegMap a = UniqFM a
 116 emptyRegMap = emptyUFM
 117
 118 type BlockMap a = UniqFM a
 119 emptyBlockMap = emptyUFM
 120
 121 -- A basic block where the isntructions are annotated with the registers
 122 -- which are no longer live in the *next* instruction in this sequence.
 123 -- (NB. if the instruction is a jump, these registers might still be live
 124 -- at the jump target(s) - you have to check the liveness at the destination
 125 -- block to find out).
 126 type AnnBasicBlock
 127         = GenBasicBlock (Instr,
 128                          [Reg],         -- registers read (only) which die
 129                          [Reg])         -- registers written which die
 130
 131 -- -----------------------------------------------------------------------------
 132 -- The free register set
 133
 134 -- This needs to be *efficient*
 135
 136 {- Here's an inefficient 'executable specification' of the FreeRegs data type:
 137 type FreeRegs = [RegNo]
 138
 139 noFreeRegs = 0
 140 releaseReg n f = if n `elem` f then f else (n : f)
 141 initFreeRegs = allocatableRegs
 142 getFreeRegs cls f = filter ( (==cls) . regClass . RealReg ) f
 143 allocateReg f r = filter (/= r) f
 144 -}
 145
 146 #if defined(powerpc_TARGET_ARCH)
 147
 148 -- The PowerPC has 32 integer and 32 floating point registers.
 149 -- This is 32bit PowerPC, so Word64 is inefficient - two Word32s are much
 150 -- better.
 151 -- Note that when getFreeRegs scans for free registers, it starts at register
 152 -- 31 and counts down. This is a hack for the PowerPC - the higher-numbered
 153 -- registers are callee-saves, while the lower regs are caller-saves, so it
 154 -- makes sense to start at the high end.
 155 -- Apart from that, the code does nothing PowerPC-specific, so feel free to
 156 -- add your favourite platform to the #if (if you have 64 registers but only
 157 -- 32-bit words).
 158
 159 data FreeRegs = FreeRegs !Word32 !Word32
 160
 161 noFreeRegs = FreeRegs 0 0
 162 releaseReg r (FreeRegs g f)
 163     | r > 31    = FreeRegs g (f .|. (1 `shiftL` (fromIntegral r - 32)))
 164     | otherwise = FreeRegs (g .|. (1 `shiftL` fromIntegral r)) f
 165
 166 initFreeRegs :: FreeRegs
 167 initFreeRegs = foldr releaseReg noFreeRegs allocatableRegs
 168
 169 getFreeRegs cls (FreeRegs g f)
 170     | RcDouble <- cls = go f (0x80000000) 63
 171     | RcInteger <- cls = go g (0x80000000) 31
 172     where
 173         go x 0 i = []
 174         go x m i | x .&. m /= 0 = i : (go x (m `shiftR` 1) $! i-1)
 175                  | otherwise    = go x (m `shiftR` 1) $! i-1
 176
 177 allocateReg (FreeRegs g f) r
 178     | r > 31    = FreeRegs g (f .&. complement (1 `shiftL` (fromIntegral r - 32)))
 179     | otherwise = FreeRegs (g .&. complement (1 `shiftL` fromIntegral r)) f
 180
 181 #else
 182
 183 -- If we have less than 32 registers, or if we have efficient 64-bit words,
 184 -- we will just use a single bitfield.
 185
 186 #if defined(alpha_TARGET_ARCH)
 187 type FreeRegs = Word64
 188 #else
 189 type FreeRegs = Word32
 190 #endif
 191
 192 noFreeRegs :: FreeRegs
 193 noFreeRegs = 0
 194
 195 releaseReg :: RegNo -> FreeRegs -> FreeRegs
 196 releaseReg n f = f .|. (1 `shiftL` n)
 197
 198 initFreeRegs :: FreeRegs
 199 initFreeRegs = foldr releaseReg noFreeRegs allocatableRegs
 200
 201 getFreeRegs :: RegClass -> FreeRegs -> [RegNo]  -- lazilly
 202 getFreeRegs cls f = go f 0
 203   where go 0 m = []
 204         go n m
 205           | n .&. 1 /= 0 && regClass (RealReg m) == cls
 206           = m : (go (n `shiftR` 1) $! (m+1))
 207           | otherwise
 208           = go (n `shiftR` 1) $! (m+1)
 209         -- ToDo: there's no point looking through all the integer registers
 210         -- in order to find a floating-point one.
 211
 212 allocateReg :: FreeRegs -> RegNo -> FreeRegs
 213 allocateReg f r = f .&. complement (1 `shiftL` fromIntegral r)
 214 #endif
 215
 216 -- -----------------------------------------------------------------------------
 217 -- The free list of stack slots
 218
 219 -- This doesn't need to be so efficient.  It also doesn't really need to be
 220 -- maintained as a set, so we just use an ordinary list (lazy, because it
 221 -- contains all the possible stack slots and there are lots :-).
 222
 223 type StackSlot = Int
 224 type FreeStack = [StackSlot]
 225
 226 completelyFreeStack :: FreeStack
 227 completelyFreeStack = [0..maxSpillSlots]
 228
 229 getFreeStackSlot :: FreeStack -> (FreeStack,Int)
 230 getFreeStackSlot (slot:stack) = (stack,slot)
 231
 232 freeStackSlot :: FreeStack -> Int -> FreeStack
 233 freeStackSlot stack slot = slot:stack
 234
 235
 236 -- -----------------------------------------------------------------------------
 237 -- Top level of the register allocator
 238
 239 regAlloc :: NatCmmTop -> NatCmmTop
 240 regAlloc (CmmData sec d) = CmmData sec d
 241 regAlloc (CmmProc info lbl params [])
 242   = CmmProc info lbl params []  -- no blocks to run the regalloc on
 243 regAlloc (CmmProc info lbl params blocks@(first:rest))
 244   = -- pprTrace "Liveness" (ppr block_live) $
 245     CmmProc info lbl params (first':rest')
 246   where
 247     first_id               = blockId first
 248     sccs                   = sccBlocks blocks
 249     (ann_sccs, block_live) = computeLiveness sccs
 250     final_blocks           = linearRegAlloc block_live ann_sccs
 251     ((first':_),rest')     = partition ((== first_id) . blockId) final_blocks
 252
 253
 254 sccBlocks :: [NatBasicBlock] -> [SCC NatBasicBlock]
 255 sccBlocks blocks = stronglyConnComp graph
 256   where
 257         getOutEdges :: [Instr] -> [BlockId]
 258         getOutEdges instrs = foldr jumpDests [] instrs
 259
 260         graph = [ (block, getUnique id, map getUnique (getOutEdges instrs))
 261                 | block@(BasicBlock id instrs) <- blocks ]
 262
 263
 264 -- -----------------------------------------------------------------------------
 265 -- Computing liveness
 266
 267 computeLiveness
 268    :: [SCC NatBasicBlock]
 269    -> ([SCC AnnBasicBlock],     -- instructions annotated with list of registers
 270                                 -- which are "dead after this instruction".
 271        BlockMap RegSet)         -- blocks annontated with set of live registers
 272                                 -- on entry to the block.
 273
 274   -- NOTE: on entry, the SCCs are in "reverse" order: later blocks may transfer
 275   -- control to earlier ones only.  The SCCs returned are in the *opposite*
 276   -- order, which is exactly what we want for the next pass.
 277
 278 computeLiveness sccs
 279   = livenessSCCs emptyBlockMap [] sccs
 280   where
 281   livenessSCCs
 282          :: BlockMap RegSet
 283          -> [SCC AnnBasicBlock]         -- accum
 284          -> [SCC NatBasicBlock]
 285          -> ([SCC AnnBasicBlock], BlockMap RegSet)
 286
 287   livenessSCCs blockmap done [] = (done, blockmap)
 288   livenessSCCs blockmap done
 289         (AcyclicSCC (BasicBlock block_id instrs) : sccs) =
 290           {- pprTrace "live instrs" (ppr (getUnique block_id) $$
 291                                   vcat (map (\(instr,regs) -> docToSDoc (pprInstr instr) $$ ppr regs) instrs')) $
 292           -}
 293           livenessSCCs blockmap'
 294                 (AcyclicSCC (BasicBlock block_id instrs'):done) sccs
 295         where (live,instrs') = liveness emptyUniqSet blockmap []
 296                                         (reverse instrs)
 297               blockmap' = addToUFM blockmap block_id live
 298         -- TODO: cope with recursive blocks
 299
 300   liveness :: RegSet                    -- live regs
 301            -> BlockMap RegSet           -- live regs on entry to other BBs
 302            -> [(Instr,[Reg],[Reg])]     -- instructions (accum)
 303            -> [Instr]                   -- instructions
 304            -> (RegSet, [(Instr,[Reg],[Reg])])
 305
 306   liveness liveregs blockmap done []  = (liveregs, done)
 307   liveness liveregs blockmap done (instr:instrs)
 308         = liveness liveregs2 blockmap ((instr,r_dying,w_dying):done) instrs
 309         where
 310               RU read written = regUsage instr
 311
 312               -- registers that were written here are dead going backwards.
 313               -- registers that were read here are live going backwards.
 314               liveregs1 = (liveregs `delListFromUniqSet` written)
 315                                     `addListToUniqSet` read
 316
 317               -- union in the live regs from all the jump destinations of this
 318               -- instruction.
 319               targets = jumpDests instr [] -- where we go from here
 320               liveregs2 = unionManyUniqSets
 321                             (liveregs1 : map (lookItUp "liveness" blockmap)
 322                                                 targets)
 323
 324               -- registers that are not live beyond this point, are recorded
 325               --  as dying here.
 326               r_dying  = [ reg | reg <- read, reg `notElem` written,
 327                                  not (elementOfUniqSet reg liveregs) ]
 328
 329               w_dying = [ reg | reg <- written,
 330                                 not (elementOfUniqSet reg liveregs) ]
 331
 332 -- -----------------------------------------------------------------------------
 333 -- Linear sweep to allocate registers
 334
 335 data Loc = InReg   {-# UNPACK #-} !RegNo
 336          | InMem   {-# UNPACK #-} !Int          -- stack slot
 337          | InBoth  {-# UNPACK #-} !RegNo
 338                    {-# UNPACK #-} !Int          -- stack slot
 339   deriving (Eq, Show)
 340
 341 {-
 342 A temporary can be marked as living in both a register and memory
 343 (InBoth), for example if it was recently loaded from a spill location.
 344 This makes it cheap to spill (no save instruction required), but we
 345 have to be careful to turn this into InReg if the value in the
 346 register is changed.
 347
 348 This is also useful when a temporary is about to be clobbered.  We
 349 save it in a spill location, but mark it as InBoth because the current
 350 instruction might still want to read it.
 351 -}
 352
 353 #ifdef DEBUG
 354 instance Outputable Loc where
 355   ppr l = text (show l)
 356 #endif
 357
 358 linearRegAlloc
 359    :: BlockMap RegSet           -- live regs on entry to each basic block
 360    -> [SCC AnnBasicBlock]       -- instructions annotated with "deaths"
 361    -> [NatBasicBlock]
 362 linearRegAlloc block_live sccs = linearRA_SCCs emptyBlockMap sccs
 363   where
 364   linearRA_SCCs
 365         :: BlockAssignment
 366         -> [SCC AnnBasicBlock]
 367         -> [NatBasicBlock]
 368   linearRA_SCCs block_assig [] = []
 369   linearRA_SCCs block_assig
 370         (AcyclicSCC (BasicBlock id instrs) : sccs)
 371         = BasicBlock id instrs' : linearRA_SCCs block_assig' sccs
 372     where
 373         (block_assig',(instrs',fixups)) =
 374            case lookupUFM block_assig id of
 375                 -- no prior info about this block: assume everything is
 376                 -- free and the assignment is empty.
 377                 Nothing ->
 378                    runR block_assig initFreeRegs
 379                                 emptyRegMap completelyFreeStack $
 380                         linearRA [] [] instrs
 381                 Just (freeregs,stack,assig) ->
 382                    runR block_assig freeregs assig stack $
 383                         linearRA [] [] instrs
 384
 385   linearRA :: [Instr] -> [NatBasicBlock] -> [(Instr,[Reg],[Reg])]
 386         -> RegM ([Instr], [NatBasicBlock])
 387   linearRA instr_acc fixups [] =
 388     return (reverse instr_acc, fixups)
 389   linearRA instr_acc fixups (instr:instrs) = do
 390     (instr_acc', new_fixups) <- raInsn block_live instr_acc instr
 391     linearRA instr_acc' (new_fixups++fixups) instrs
 392
 393 -- -----------------------------------------------------------------------------
 394 -- Register allocation for a single instruction
 395
 396 type BlockAssignment = BlockMap (FreeRegs, FreeStack, RegMap Loc)
 397
 398 raInsn  :: BlockMap RegSet              -- Live temporaries at each basic block
 399         -> [Instr]                      -- new instructions (accum.)
 400         -> (Instr,[Reg],[Reg])          -- the instruction (with "deaths")
 401         -> RegM (
 402              [Instr],                   -- new instructions
 403              [NatBasicBlock]            -- extra fixup blocks
 404            )
 405
 406 raInsn block_live new_instrs (instr@(DELTA n), _, _) = do
 407     setDeltaR n
 408     return (new_instrs, [])
 409
 410 raInsn block_live new_instrs (instr, r_dying, w_dying) = do
 411     assig    <- getAssigR
 412
 413     -- If we have a reg->reg move between virtual registers, where the
 414     -- src register is not live after this instruction, and the dst
 415     -- register does not already have an assignment, then we can
 416     -- eliminate the instruction.
 417     case isRegRegMove instr of
 418         Just (src,dst)
 419                 | src `elem` r_dying,
 420                   isVirtualReg dst,
 421                   Just loc <- lookupUFM assig src,
 422                   not (dst `elemUFM` assig) -> do
 423                         setAssigR (addToUFM (delFromUFM assig src) dst loc)
 424                         return (new_instrs, [])
 425
 426         other -> genRaInsn block_live new_instrs instr r_dying w_dying
 427
 428
 429 genRaInsn block_live new_instrs instr r_dying w_dying = do
 430     let
 431         RU read written = regUsage instr
 432
 433         (real_written1,virt_written) = partition isRealReg written
 434
 435         real_written = [ r | RealReg r <- real_written1 ]
 436
 437         -- we don't need to do anything with real registers that are
 438         -- only read by this instr.  (the list is typically ~2 elements,
 439         -- so using nub isn't a problem).
 440         virt_read = nub (filter isVirtualReg read)
 441     -- in
 442
 443     -- (a) save any temporaries which will be clobbered by this instruction
 444     clobber_saves <- saveClobberedTemps real_written r_dying
 445
 446     {-
 447     freeregs <- getFreeRegsR
 448     assig <- getAssigR
 449     pprTrace "raInsn" (docToSDoc (pprInstr instr) $$ ppr r_dying <+> ppr w_dying $$ ppr virt_read <+> ppr virt_written $$ text (show freeregs) $$ ppr assig) $ do
 450     -}
 451
 452     -- (b), (c) allocate real regs for all regs read by this instruction.
 453     (r_spills, r_allocd) <-
 454         allocateRegsAndSpill True{-reading-} virt_read [] [] virt_read
 455
 456     -- (d) Update block map for new destinations
 457     -- NB. do this before removing dead regs from the assignment, because
 458     -- these dead regs might in fact be live in the jump targets (they're
 459     -- only dead in the code that follows in the current basic block).
 460     (fixup_blocks, adjusted_instr)
 461         <- joinToTargets block_live [] instr (jumpDests instr [])
 462
 463     -- (e) Delete all register assignments for temps which are read
 464     --     (only) and die here.  Update the free register list.
 465     releaseRegs r_dying
 466
 467     -- (f) Mark regs which are clobbered as unallocatable
 468     clobberRegs real_written
 469
 470     -- (g) Allocate registers for temporaries *written* (only)
 471     (w_spills, w_allocd) <-
 472         allocateRegsAndSpill False{-writing-} virt_written [] [] virt_written
 473
 474     -- (h) Release registers for temps which are written here and not
 475     -- used again.
 476     releaseRegs w_dying
 477
 478     let
 479         -- (i) Patch the instruction
 480         patch_map = listToUFM   [ (t,RealReg r) |
 481                                   (t,r) <- zip virt_read r_allocd
 482                                           ++ zip virt_written w_allocd ]
 483
 484         patched_instr = patchRegs adjusted_instr patchLookup
 485         patchLookup x = case lookupUFM patch_map x of
 486                                 Nothing -> x
 487                                 Just y  -> y
 488     -- in
 489
 490     -- pprTrace "patched" (docToSDoc (pprInstr patched_instr)) $ do
 491
 492     -- (j) free up stack slots for dead spilled regs
 493     -- TODO (can't be bothered right now)
 494
 495     return (patched_instr : w_spills ++ reverse r_spills
 496                  ++ clobber_saves ++ new_instrs,
 497             fixup_blocks)
 498
 499 -- -----------------------------------------------------------------------------
 500 -- releaseRegs
 501
 502 releaseRegs regs = do
 503   assig <- getAssigR
 504   free <- getFreeRegsR
 505   loop assig free regs
 506  where
 507   loop assig free [] = do setAssigR assig; setFreeRegsR free; return ()
 508   loop assig free (RealReg r : rs) = loop assig (releaseReg r free) rs
 509   loop assig free (r:rs) =
 510      case lookupUFM assig r of
 511         Just (InBoth real _) -> loop (delFromUFM assig r) (releaseReg real free) rs
 512         Just (InReg real) -> loop (delFromUFM assig r) (releaseReg real free) rs
 513         _other            -> loop (delFromUFM assig r) free rs
 514
 515 -- -----------------------------------------------------------------------------
 516 -- Clobber real registers
 517
 518 {-
 519 For each temp in a register that is going to be clobbered:
 520   - if the temp dies after this instruction, do nothing
 521   - otherwise, put it somewhere safe (another reg if possible,
 522     otherwise spill and record InBoth in the assignment).
 523
 524 for allocateRegs on the temps *read*,
 525   - clobbered regs are allocatable.
 526
 527 for allocateRegs on the temps *written*,
 528   - clobbered regs are not allocatable.
 529 -}
 530
 531 saveClobberedTemps
 532    :: [RegNo]              -- real registers clobbered by this instruction
 533    -> [Reg]                -- registers which are no longer live after this insn
 534    -> RegM [Instr]         -- return: instructions to spill any temps that will
 535                            -- be clobbered.
 536
 537 saveClobberedTemps [] _ = return [] -- common case
 538 saveClobberedTemps clobbered dying =  do
 539   assig <- getAssigR
 540   let
 541         to_spill  = [ (temp,reg) | (temp, InReg reg) <- ufmToList assig,
 542                                    reg `elem` clobbered,
 543                                    temp `notElem` map getUnique dying  ]
 544   -- in
 545   (instrs,assig') <- clobber assig [] to_spill
 546   setAssigR assig'
 547   return instrs
 548  where
 549   clobber assig instrs [] = return (instrs,assig)
 550   clobber assig instrs ((temp,reg):rest)
 551     = do
 552         --ToDo: copy it to another register if possible
 553       (spill,slot) <- spillR (RealReg reg)
 554       clobber (addToUFM assig temp (InBoth reg slot)) (spill:instrs) rest
 555
 556 clobberRegs :: [RegNo] -> RegM ()
 557 clobberRegs [] = return () -- common case
 558 clobberRegs clobbered = do
 559   freeregs <- getFreeRegsR
 560   setFreeRegsR (foldl allocateReg freeregs clobbered)
 561   assig <- getAssigR
 562   setAssigR $! clobber assig (ufmToList assig)
 563  where
 564     -- if the temp was InReg and clobbered, then we will have
 565     -- saved it in saveClobberedTemps above.  So the only case
 566     -- we have to worry about here is InBoth.  Note that this
 567     -- also catches temps which were loaded up during allocation
 568     -- of read registers, not just those saved in saveClobberedTemps.
 569   clobber assig [] = assig
 570   clobber assig ((temp, InBoth reg slot) : rest)
 571         | reg `elem` clobbered
 572         = clobber (addToUFM assig temp (InMem slot)) rest
 573   clobber assig (entry:rest)
 574         = clobber assig rest
 575
 576 -- -----------------------------------------------------------------------------
 577 -- allocateRegsAndSpill
 578
 579 -- This function does several things:
 580 --   For each temporary referred to by this instruction,
 581 --   we allocate a real register (spilling another temporary if necessary).
 582 --   We load the temporary up from memory if necessary.
 583 --   We also update the register assignment in the process, and
 584 --   the list of free registers and free stack slots.
 585
 586 allocateRegsAndSpill
 587         :: Bool                 -- True <=> reading (load up spilled regs)
 588         -> [Reg]                -- don't push these out
 589         -> [Instr]              -- spill insns
 590         -> [RegNo]              -- real registers allocated (accum.)
 591         -> [Reg]                -- temps to allocate
 592         -> RegM ([Instr], [RegNo])
 593
 594 allocateRegsAndSpill reading keep spills alloc []
 595   = return (spills,reverse alloc)
 596
 597 allocateRegsAndSpill reading keep spills alloc (r:rs) = do
 598   assig <- getAssigR
 599   case lookupUFM assig r of
 600   -- case (1a): already in a register
 601      Just (InReg my_reg) ->
 602         allocateRegsAndSpill reading keep spills (my_reg:alloc) rs
 603
 604   -- case (1b): already in a register (and memory)
 605   -- NB1. if we're writing this register, update its assignemnt to be
 606   -- InReg, because the memory value is no longer valid.
 607   -- NB2. This is why we must process written registers here, even if they
 608   -- are also read by the same instruction.
 609      Just (InBoth my_reg mem) -> do
 610         when (not reading) (setAssigR (addToUFM assig r (InReg my_reg)))
 611         allocateRegsAndSpill reading keep spills (my_reg:alloc) rs
 612
 613   -- Not already in a register, so we need to find a free one...
 614      loc -> do
 615         freeregs <- getFreeRegsR
 616
 617         case getFreeRegs (regClass r) freeregs of
 618
 619         -- case (2): we have a free register
 620           my_reg:_ -> do
 621             spills'   <- do_load reading loc my_reg spills
 622             let new_loc
 623                  | Just (InMem slot) <- loc, reading = InBoth my_reg slot
 624                  | otherwise                         = InReg my_reg
 625             setAssigR (addToUFM assig r $! new_loc)
 626             setFreeRegsR (allocateReg freeregs my_reg)
 627             allocateRegsAndSpill reading keep spills' (my_reg:alloc) rs
 628
 629         -- case (3): we need to push something out to free up a register
 630           [] -> do
 631             let
 632               keep' = map getUnique keep
 633               candidates1 = [ (temp,reg,mem)
 634                             | (temp, InBoth reg mem) <- ufmToList assig,
 635                               temp `notElem` keep', regClass (RealReg reg) == regClass r ]
 636               candidates2 = [ (temp,reg)
 637                             | (temp, InReg reg) <- ufmToList assig,
 638                               temp `notElem` keep', regClass (RealReg reg) == regClass r  ]
 639             -- in
 640             ASSERT2(not (null candidates1 && null candidates2), ppr assig) do
 641
 642             case candidates1 of
 643
 644              -- we have a temporary that is in both register and mem,
 645              -- just free up its register for use.
 646              --
 647              (temp,my_reg,slot):_ -> do
 648                 spills' <- do_load reading loc my_reg spills
 649                 let
 650                   assig1  = addToUFM assig temp (InMem slot)
 651                   assig2  = addToUFM assig1 r (InReg my_reg)
 652                 -- in
 653                 setAssigR assig2
 654                 allocateRegsAndSpill reading keep spills' (my_reg:alloc) rs
 655
 656              -- otherwise, we need to spill a temporary that currently
 657              -- resides in a register.
 658              [] -> do
 659                 let
 660                   (temp_to_push_out, my_reg) = head candidates2
 661                   -- TODO: plenty of room for optimisation in choosing which temp
 662                   -- to spill.  We just pick the first one that isn't used in
 663                   -- the current instruction for now.
 664                 -- in
 665                 (spill_insn,slot) <- spillR (RealReg my_reg)
 666                 let
 667                   assig1  = addToUFM assig temp_to_push_out (InMem slot)
 668                   assig2  = addToUFM assig1 r (InReg my_reg)
 669                 -- in
 670                 setAssigR assig2
 671                 spills' <- do_load reading loc my_reg spills
 672                 allocateRegsAndSpill reading keep (spill_insn:spills')
 673                         (my_reg:alloc) rs
 674   where
 675         -- load up a spilled temporary if we need to
 676         do_load True (Just (InMem slot)) reg spills = do
 677            insn <- loadR (RealReg reg) slot
 678            return (insn : spills)
 679         do_load _ _ _ spills =
 680            return spills
 681
 682 -- -----------------------------------------------------------------------------
 683 -- Joining a jump instruction to its targets
 684
 685 -- The first time we encounter a jump to a particular basic block, we
 686 -- record the assignment of temporaries.  The next time we encounter a
 687 -- jump to the same block, we compare our current assignment to the
 688 -- stored one.  They might be different if spilling has occrred in one
 689 -- branch; so some fixup code will be required to match up the
 690 -- assignments.
 691
 692 joinToTargets
 693         :: BlockMap RegSet
 694         -> [NatBasicBlock]
 695         -> Instr
 696         -> [BlockId]
 697         -> RegM ([NatBasicBlock], Instr)
 698
 699 joinToTargets block_live new_blocks instr []
 700   = return (new_blocks, instr)
 701 joinToTargets block_live new_blocks instr (dest:dests) = do
 702   block_assig <- getBlockAssigR
 703   assig <- getAssigR
 704   let
 705         -- adjust the assignment to remove any registers which are not
 706         -- live on entry to the destination block.
 707         adjusted_assig =
 708           listToUFM [ (reg,loc) | reg <- live,
 709                                   Just loc <- [lookupUFM assig reg] ]
 710   -- in
 711   case lookupUFM block_assig dest of
 712         -- Nothing <=> this is the first time we jumped to this
 713         -- block.
 714         Nothing -> do
 715           freeregs <- getFreeRegsR
 716           stack <- getStackR
 717           setBlockAssigR (addToUFM block_assig dest
 718                                 (freeregs,stack,adjusted_assig))
 719           joinToTargets block_live new_blocks instr dests
 720
 721         Just (freeregs,stack,dest_assig)
 722            | ufmToList dest_assig == ufmToList adjusted_assig
 723            -> -- ok, the assignments match
 724              joinToTargets block_live new_blocks instr dests
 725            | otherwise
 726            -> -- need fixup code
 727              panic "joinToTargets: ToDo: need fixup code"
 728   where
 729         live = uniqSetToList (lookItUp "joinToTargets" block_live dest)
 730
 731 -- -----------------------------------------------------------------------------
 732 -- The register allocator's monad.
 733
 734 -- Here we keep all the state that the register allocator keeps track
 735 -- of as it walks the instructions in a basic block.
 736
 737 data RA_State
 738   = RA_State {
 739         ra_blockassig :: BlockAssignment,
 740                 -- The current mapping from basic blocks to
 741                 -- the register assignments at the beginning of that block.
 742         ra_freeregs   :: FreeRegs,      -- free machine registers
 743         ra_assig      :: RegMap Loc,    -- assignment of temps to locations
 744         ra_delta      :: Int,           -- current stack delta
 745         ra_stack      :: FreeStack      -- free stack slots for spilling
 746   }
 747
 748 newtype RegM a = RegM { unReg :: RA_State -> (# RA_State, a #) }
 749
 750 instance Monad RegM where
 751   m >>= k   =  RegM $ \s -> case unReg m s of { (# s, a #) -> unReg (k a) s }
 752   return a  =  RegM $ \s -> (# s, a #)
 753
 754 runR :: BlockAssignment -> FreeRegs -> RegMap Loc -> FreeStack -> RegM a ->
 755   (BlockAssignment, a)
 756 runR block_assig freeregs assig stack thing =
 757   case unReg thing (RA_State{ ra_blockassig=block_assig, ra_freeregs=freeregs,
 758                         ra_assig=assig, ra_delta=0{-???-}, ra_stack=stack }) of
 759         (# RA_State{ ra_blockassig=block_assig }, returned_thing #)
 760                 -> (block_assig, returned_thing)
 761
 762 spillR :: Reg -> RegM (Instr, Int)
 763 spillR reg = RegM $ \ s@RA_State{ra_delta=delta, ra_stack=stack} ->
 764   let (stack',slot) = getFreeStackSlot stack
 765       instr  = mkSpillInstr reg delta slot
 766   in
 767   (# s{ra_stack=stack'}, (instr,slot) #)
 768
 769 loadR :: Reg -> Int -> RegM Instr
 770 loadR reg slot = RegM $ \ s@RA_State{ra_delta=delta, ra_stack=stack} ->
 771   (# s, mkLoadInstr reg delta slot #)
 772
 773 freeSlotR :: Int -> RegM ()
 774 freeSlotR slot = RegM $ \ s@RA_State{ra_stack=stack} ->
 775   (# s{ra_stack=freeStackSlot stack slot}, () #)
 776
 777 getFreeRegsR :: RegM FreeRegs
 778 getFreeRegsR = RegM $ \ s@RA_State{ra_freeregs = freeregs} ->
 779   (# s, freeregs #)
 780
 781 setFreeRegsR :: FreeRegs -> RegM ()
 782 setFreeRegsR regs = RegM $ \ s ->
 783   (# s{ra_freeregs = regs}, () #)
 784
 785 getAssigR :: RegM (RegMap Loc)
 786 getAssigR = RegM $ \ s@RA_State{ra_assig = assig} ->
 787   (# s, assig #)
 788
 789 setAssigR :: RegMap Loc -> RegM ()
 790 setAssigR assig = RegM $ \ s ->
 791   (# s{ra_assig=assig}, () #)
 792
 793 getStackR :: RegM FreeStack
 794 getStackR = RegM $ \ s@RA_State{ra_stack = stack} ->
 795   (# s, stack #)
 796
 797 setStackR :: FreeStack -> RegM ()
 798 setStackR stack = RegM $ \ s ->
 799   (# s{ra_stack=stack}, () #)
 800
 801 getBlockAssigR :: RegM BlockAssignment
 802 getBlockAssigR = RegM $ \ s@RA_State{ra_blockassig = assig} ->
 803   (# s, assig #)
 804
 805 setBlockAssigR :: BlockAssignment -> RegM ()
 806 setBlockAssigR assig = RegM $ \ s ->
 807   (# s{ra_blockassig = assig}, () #)
 808
 809 setDeltaR :: Int -> RegM ()
 810 setDeltaR n = RegM $ \ s ->
 811   (# s{ra_delta = n}, () #)
 812
 813 -- -----------------------------------------------------------------------------
 814 -- Utils
 815
 816 #ifdef DEBUG
 817 my_fromJust s p Nothing  = pprPanic ("fromJust: " ++ s) p
 818 my_fromJust s p (Just x) = x
 819 #else
 820 my_fromJust _ _ = fromJust
 821 #endif
 822
 823 lookItUp :: Uniquable b => String -> UniqFM a -> b -> a
 824 lookItUp str fm x = my_fromJust str (ppr (getUnique x)) (lookupUFM fm x)