{-# LANGUAGE UnboxedTuples #-}

-- | State monad for the linear register allocator.

--      Here we keep all the state that the register allocator keeps track
--      of as it walks the instructions in a basic block.

module RegAlloc.Linear.State (
        RA_State(..),
        RegM,
        runR,

        spillR,
        loadR,

        getFreeRegsR,
        setFreeRegsR,

        getAssigR,
        setAssigR,

        getBlockAssigR,
        setBlockAssigR,

        setDeltaR,
        getDeltaR,

        getUniqueR,

        recordSpill,
        recordFixupBlock
)
where

import GhcPrelude

import RegAlloc.Linear.Stats
import RegAlloc.Linear.StackMap
import RegAlloc.Linear.Base
import RegAlloc.Liveness
import Instruction
import Reg
import BlockId

import DynFlags
import Unique
import UniqSupply

import Control.Monad (liftM, ap)

-- | The register allocator monad type.
newtype RegM freeRegs a
        = RegM { RegM freeRegs a -> RA_State freeRegs -> (# RA_State freeRegs, a #)
unReg :: RA_State freeRegs -> (# RA_State freeRegs, a #) }

instance Functor (RegM freeRegs) where
      fmap :: (a -> b) -> RegM freeRegs a -> RegM freeRegs b
fmap = (a -> b) -> RegM freeRegs a -> RegM freeRegs b
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM

instance Applicative (RegM freeRegs) where
      pure :: a -> RegM freeRegs a
pure a :: a
a  =  (RA_State freeRegs -> (# RA_State freeRegs, a #))
-> RegM freeRegs a
forall freeRegs a.
(RA_State freeRegs -> (# RA_State freeRegs, a #))
-> RegM freeRegs a
RegM ((RA_State freeRegs -> (# RA_State freeRegs, a #))
 -> RegM freeRegs a)
-> (RA_State freeRegs -> (# RA_State freeRegs, a #))
-> RegM freeRegs a
forall a b. (a -> b) -> a -> b
$ \s :: RA_State freeRegs
s -> (# RA_State freeRegs
s, a
a #)
      <*> :: RegM freeRegs (a -> b) -> RegM freeRegs a -> RegM freeRegs b
(<*>) = RegM freeRegs (a -> b) -> RegM freeRegs a -> RegM freeRegs b
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance Monad (RegM freeRegs) where
  m :: RegM freeRegs a
m >>= :: RegM freeRegs a -> (a -> RegM freeRegs b) -> RegM freeRegs b
>>= k :: a -> RegM freeRegs b
k   =  (RA_State freeRegs -> (# RA_State freeRegs, b #))
-> RegM freeRegs b
forall freeRegs a.
(RA_State freeRegs -> (# RA_State freeRegs, a #))
-> RegM freeRegs a
RegM ((RA_State freeRegs -> (# RA_State freeRegs, b #))
 -> RegM freeRegs b)
-> (RA_State freeRegs -> (# RA_State freeRegs, b #))
-> RegM freeRegs b
forall a b. (a -> b) -> a -> b
$ \s :: RA_State freeRegs
s -> case RegM freeRegs a -> RA_State freeRegs -> (# RA_State freeRegs, a #)
forall freeRegs a.
RegM freeRegs a -> RA_State freeRegs -> (# RA_State freeRegs, a #)
unReg RegM freeRegs a
m RA_State freeRegs
s of { (# s :: RA_State freeRegs
s, a :: a
a #) -> RegM freeRegs b -> RA_State freeRegs -> (# RA_State freeRegs, b #)
forall freeRegs a.
RegM freeRegs a -> RA_State freeRegs -> (# RA_State freeRegs, a #)
unReg (a -> RegM freeRegs b
k a
a) RA_State freeRegs
s }

instance HasDynFlags (RegM a) where
    getDynFlags :: RegM a DynFlags
getDynFlags = (RA_State a -> (# RA_State a, DynFlags #)) -> RegM a DynFlags
forall freeRegs a.
(RA_State freeRegs -> (# RA_State freeRegs, a #))
-> RegM freeRegs a
RegM ((RA_State a -> (# RA_State a, DynFlags #)) -> RegM a DynFlags)
-> (RA_State a -> (# RA_State a, DynFlags #)) -> RegM a DynFlags
forall a b. (a -> b) -> a -> b
$ \s :: RA_State a
s -> (# RA_State a
s, RA_State a -> DynFlags
forall freeRegs. RA_State freeRegs -> DynFlags
ra_DynFlags RA_State a
s #)


-- | Run a computation in the RegM register allocator monad.
runR    :: DynFlags
        -> BlockAssignment freeRegs
        -> freeRegs
        -> RegMap Loc
        -> StackMap
        -> UniqSupply
        -> RegM freeRegs a
        -> (BlockAssignment freeRegs, StackMap, RegAllocStats, a)

runR :: DynFlags
-> BlockAssignment freeRegs
-> freeRegs
-> RegMap Loc
-> StackMap
-> UniqSupply
-> RegM freeRegs a
-> (BlockAssignment freeRegs, StackMap, RegAllocStats, a)
runR dflags :: DynFlags
dflags block_assig :: BlockAssignment freeRegs
block_assig freeregs :: freeRegs
freeregs assig :: RegMap Loc
assig stack :: StackMap
stack us :: UniqSupply
us thing :: RegM freeRegs a
thing =
  case RegM freeRegs a -> RA_State freeRegs -> (# RA_State freeRegs, a #)
forall freeRegs a.
RegM freeRegs a -> RA_State freeRegs -> (# RA_State freeRegs, a #)
unReg RegM freeRegs a
thing
        ($WRA_State :: forall freeRegs.
BlockAssignment freeRegs
-> freeRegs
-> RegMap Loc
-> Int
-> StackMap
-> UniqSupply
-> [SpillReason]
-> DynFlags
-> [(BlockId, BlockId, BlockId)]
-> RA_State freeRegs
RA_State
                { ra_blockassig :: BlockAssignment freeRegs
ra_blockassig = BlockAssignment freeRegs
block_assig
                , ra_freeregs :: freeRegs
ra_freeregs   = freeRegs
freeregs
                , ra_assig :: RegMap Loc
ra_assig      = RegMap Loc
assig
                , ra_delta :: Int
ra_delta      = 0{-???-}
                , ra_stack :: StackMap
ra_stack      = StackMap
stack
                , ra_us :: UniqSupply
ra_us         = UniqSupply
us
                , ra_spills :: [SpillReason]
ra_spills     = []
                , ra_DynFlags :: DynFlags
ra_DynFlags   = DynFlags
dflags
                , ra_fixups :: [(BlockId, BlockId, BlockId)]
ra_fixups     = [] })
   of
        (# state' :: RA_State freeRegs
state'@RA_State
                { ra_blockassig :: forall freeRegs. RA_State freeRegs -> BlockAssignment freeRegs
ra_blockassig = BlockAssignment freeRegs
block_assig
                , ra_stack :: forall freeRegs. RA_State freeRegs -> StackMap
ra_stack      = StackMap
stack' }
                , returned_thing :: a
returned_thing #)

         ->     (BlockAssignment freeRegs
block_assig, StackMap
stack', RA_State freeRegs -> RegAllocStats
forall freeRegs. RA_State freeRegs -> RegAllocStats
makeRAStats RA_State freeRegs
state', a
returned_thing)


-- | Make register allocator stats from its final state.
makeRAStats :: RA_State freeRegs -> RegAllocStats
makeRAStats :: RA_State freeRegs -> RegAllocStats
makeRAStats state :: RA_State freeRegs
state
        = RegAllocStats :: UniqFM [Int] -> [(BlockId, BlockId, BlockId)] -> RegAllocStats
RegAllocStats
        { ra_spillInstrs :: UniqFM [Int]
ra_spillInstrs        = [SpillReason] -> UniqFM [Int]
binSpillReasons (RA_State freeRegs -> [SpillReason]
forall freeRegs. RA_State freeRegs -> [SpillReason]
ra_spills RA_State freeRegs
state)
        , ra_fixupList :: [(BlockId, BlockId, BlockId)]
ra_fixupList          = RA_State freeRegs -> [(BlockId, BlockId, BlockId)]
forall freeRegs. RA_State freeRegs -> [(BlockId, BlockId, BlockId)]
ra_fixups RA_State freeRegs
state }


spillR :: Instruction instr
       => Reg -> Unique -> RegM freeRegs (instr, Int)

spillR :: Reg -> Unique -> RegM freeRegs (instr, Int)
spillR reg :: Reg
reg temp :: Unique
temp = (RA_State freeRegs -> (# RA_State freeRegs, (instr, Int) #))
-> RegM freeRegs (instr, Int)
forall freeRegs a.
(RA_State freeRegs -> (# RA_State freeRegs, a #))
-> RegM freeRegs a
RegM ((RA_State freeRegs -> (# RA_State freeRegs, (instr, Int) #))
 -> RegM freeRegs (instr, Int))
-> (RA_State freeRegs -> (# RA_State freeRegs, (instr, Int) #))
-> RegM freeRegs (instr, Int)
forall a b. (a -> b) -> a -> b
$ \ s :: RA_State freeRegs
s@RA_State{ra_delta :: forall freeRegs. RA_State freeRegs -> Int
ra_delta=Int
delta, ra_stack :: forall freeRegs. RA_State freeRegs -> StackMap
ra_stack=StackMap
stack} ->
  let dflags :: DynFlags
dflags = RA_State freeRegs -> DynFlags
forall freeRegs. RA_State freeRegs -> DynFlags
ra_DynFlags RA_State freeRegs
s
      (stack' :: StackMap
stack',slot :: Int
slot) = StackMap -> Unique -> (StackMap, Int)
getStackSlotFor StackMap
stack Unique
temp
      instr :: instr
instr  = DynFlags -> Reg -> Int -> Int -> instr
forall instr.
Instruction instr =>
DynFlags -> Reg -> Int -> Int -> instr
mkSpillInstr DynFlags
dflags Reg
reg Int
delta Int
slot
  in
  (# RA_State freeRegs
s{ra_stack :: StackMap
ra_stack=StackMap
stack'}, (instr
instr,Int
slot) #)


loadR :: Instruction instr
      => Reg -> Int -> RegM freeRegs instr

loadR :: Reg -> Int -> RegM freeRegs instr
loadR reg :: Reg
reg slot :: Int
slot = (RA_State freeRegs -> (# RA_State freeRegs, instr #))
-> RegM freeRegs instr
forall freeRegs a.
(RA_State freeRegs -> (# RA_State freeRegs, a #))
-> RegM freeRegs a
RegM ((RA_State freeRegs -> (# RA_State freeRegs, instr #))
 -> RegM freeRegs instr)
-> (RA_State freeRegs -> (# RA_State freeRegs, instr #))
-> RegM freeRegs instr
forall a b. (a -> b) -> a -> b
$ \ s :: RA_State freeRegs
s@RA_State{ra_delta :: forall freeRegs. RA_State freeRegs -> Int
ra_delta=Int
delta} ->
  let dflags :: DynFlags
dflags = RA_State freeRegs -> DynFlags
forall freeRegs. RA_State freeRegs -> DynFlags
ra_DynFlags RA_State freeRegs
s
  in (# RA_State freeRegs
s, DynFlags -> Reg -> Int -> Int -> instr
forall instr.
Instruction instr =>
DynFlags -> Reg -> Int -> Int -> instr
mkLoadInstr DynFlags
dflags Reg
reg Int
delta Int
slot #)

getFreeRegsR :: RegM freeRegs freeRegs
getFreeRegsR :: RegM freeRegs freeRegs
getFreeRegsR = (RA_State freeRegs -> (# RA_State freeRegs, freeRegs #))
-> RegM freeRegs freeRegs
forall freeRegs a.
(RA_State freeRegs -> (# RA_State freeRegs, a #))
-> RegM freeRegs a
RegM ((RA_State freeRegs -> (# RA_State freeRegs, freeRegs #))
 -> RegM freeRegs freeRegs)
-> (RA_State freeRegs -> (# RA_State freeRegs, freeRegs #))
-> RegM freeRegs freeRegs
forall a b. (a -> b) -> a -> b
$ \ s :: RA_State freeRegs
s@RA_State{ra_freeregs :: forall freeRegs. RA_State freeRegs -> freeRegs
ra_freeregs = freeRegs
freeregs} ->
  (# RA_State freeRegs
s, freeRegs
freeregs #)

setFreeRegsR :: freeRegs -> RegM freeRegs ()
setFreeRegsR :: freeRegs -> RegM freeRegs ()
setFreeRegsR regs :: freeRegs
regs = (RA_State freeRegs -> (# RA_State freeRegs, () #))
-> RegM freeRegs ()
forall freeRegs a.
(RA_State freeRegs -> (# RA_State freeRegs, a #))
-> RegM freeRegs a
RegM ((RA_State freeRegs -> (# RA_State freeRegs, () #))
 -> RegM freeRegs ())
-> (RA_State freeRegs -> (# RA_State freeRegs, () #))
-> RegM freeRegs ()
forall a b. (a -> b) -> a -> b
$ \ s :: RA_State freeRegs
s ->
  (# RA_State freeRegs
s{ra_freeregs :: freeRegs
ra_freeregs = freeRegs
regs}, () #)

getAssigR :: RegM freeRegs (RegMap Loc)
getAssigR :: RegM freeRegs (RegMap Loc)
getAssigR = (RA_State freeRegs -> (# RA_State freeRegs, RegMap Loc #))
-> RegM freeRegs (RegMap Loc)
forall freeRegs a.
(RA_State freeRegs -> (# RA_State freeRegs, a #))
-> RegM freeRegs a
RegM ((RA_State freeRegs -> (# RA_State freeRegs, RegMap Loc #))
 -> RegM freeRegs (RegMap Loc))
-> (RA_State freeRegs -> (# RA_State freeRegs, RegMap Loc #))
-> RegM freeRegs (RegMap Loc)
forall a b. (a -> b) -> a -> b
$ \ s :: RA_State freeRegs
s@RA_State{ra_assig :: forall freeRegs. RA_State freeRegs -> RegMap Loc
ra_assig = RegMap Loc
assig} ->
  (# RA_State freeRegs
s, RegMap Loc
assig #)

setAssigR :: RegMap Loc -> RegM freeRegs ()
setAssigR :: RegMap Loc -> RegM freeRegs ()
setAssigR assig :: RegMap Loc
assig = (RA_State freeRegs -> (# RA_State freeRegs, () #))
-> RegM freeRegs ()
forall freeRegs a.
(RA_State freeRegs -> (# RA_State freeRegs, a #))
-> RegM freeRegs a
RegM ((RA_State freeRegs -> (# RA_State freeRegs, () #))
 -> RegM freeRegs ())
-> (RA_State freeRegs -> (# RA_State freeRegs, () #))
-> RegM freeRegs ()
forall a b. (a -> b) -> a -> b
$ \ s :: RA_State freeRegs
s ->
  (# RA_State freeRegs
s{ra_assig :: RegMap Loc
ra_assig=RegMap Loc
assig}, () #)

getBlockAssigR :: RegM freeRegs (BlockAssignment freeRegs)
getBlockAssigR :: RegM freeRegs (BlockAssignment freeRegs)
getBlockAssigR = (RA_State freeRegs
 -> (# RA_State freeRegs, BlockAssignment freeRegs #))
-> RegM freeRegs (BlockAssignment freeRegs)
forall freeRegs a.
(RA_State freeRegs -> (# RA_State freeRegs, a #))
-> RegM freeRegs a
RegM ((RA_State freeRegs
  -> (# RA_State freeRegs, BlockAssignment freeRegs #))
 -> RegM freeRegs (BlockAssignment freeRegs))
-> (RA_State freeRegs
    -> (# RA_State freeRegs, BlockAssignment freeRegs #))
-> RegM freeRegs (BlockAssignment freeRegs)
forall a b. (a -> b) -> a -> b
$ \ s :: RA_State freeRegs
s@RA_State{ra_blockassig :: forall freeRegs. RA_State freeRegs -> BlockAssignment freeRegs
ra_blockassig = BlockAssignment freeRegs
assig} ->
  (# RA_State freeRegs
s, BlockAssignment freeRegs
assig #)

setBlockAssigR :: BlockAssignment freeRegs -> RegM freeRegs ()
setBlockAssigR :: BlockAssignment freeRegs -> RegM freeRegs ()
setBlockAssigR assig :: BlockAssignment freeRegs
assig = (RA_State freeRegs -> (# RA_State freeRegs, () #))
-> RegM freeRegs ()
forall freeRegs a.
(RA_State freeRegs -> (# RA_State freeRegs, a #))
-> RegM freeRegs a
RegM ((RA_State freeRegs -> (# RA_State freeRegs, () #))
 -> RegM freeRegs ())
-> (RA_State freeRegs -> (# RA_State freeRegs, () #))
-> RegM freeRegs ()
forall a b. (a -> b) -> a -> b
$ \ s :: RA_State freeRegs
s ->
  (# RA_State freeRegs
s{ra_blockassig :: BlockAssignment freeRegs
ra_blockassig = BlockAssignment freeRegs
assig}, () #)

setDeltaR :: Int -> RegM freeRegs ()
setDeltaR :: Int -> RegM freeRegs ()
setDeltaR n :: Int
n = (RA_State freeRegs -> (# RA_State freeRegs, () #))
-> RegM freeRegs ()
forall freeRegs a.
(RA_State freeRegs -> (# RA_State freeRegs, a #))
-> RegM freeRegs a
RegM ((RA_State freeRegs -> (# RA_State freeRegs, () #))
 -> RegM freeRegs ())
-> (RA_State freeRegs -> (# RA_State freeRegs, () #))
-> RegM freeRegs ()
forall a b. (a -> b) -> a -> b
$ \ s :: RA_State freeRegs
s ->
  (# RA_State freeRegs
s{ra_delta :: Int
ra_delta = Int
n}, () #)

getDeltaR :: RegM freeRegs Int
getDeltaR :: RegM freeRegs Int
getDeltaR = (RA_State freeRegs -> (# RA_State freeRegs, Int #))
-> RegM freeRegs Int
forall freeRegs a.
(RA_State freeRegs -> (# RA_State freeRegs, a #))
-> RegM freeRegs a
RegM ((RA_State freeRegs -> (# RA_State freeRegs, Int #))
 -> RegM freeRegs Int)
-> (RA_State freeRegs -> (# RA_State freeRegs, Int #))
-> RegM freeRegs Int
forall a b. (a -> b) -> a -> b
$ \s :: RA_State freeRegs
s -> (# RA_State freeRegs
s, RA_State freeRegs -> Int
forall freeRegs. RA_State freeRegs -> Int
ra_delta RA_State freeRegs
s #)

getUniqueR :: RegM freeRegs Unique
getUniqueR :: RegM freeRegs Unique
getUniqueR = (RA_State freeRegs -> (# RA_State freeRegs, Unique #))
-> RegM freeRegs Unique
forall freeRegs a.
(RA_State freeRegs -> (# RA_State freeRegs, a #))
-> RegM freeRegs a
RegM ((RA_State freeRegs -> (# RA_State freeRegs, Unique #))
 -> RegM freeRegs Unique)
-> (RA_State freeRegs -> (# RA_State freeRegs, Unique #))
-> RegM freeRegs Unique
forall a b. (a -> b) -> a -> b
$ \s :: RA_State freeRegs
s ->
  case UniqSupply -> (Unique, UniqSupply)
takeUniqFromSupply (RA_State freeRegs -> UniqSupply
forall freeRegs. RA_State freeRegs -> UniqSupply
ra_us RA_State freeRegs
s) of
    (uniq :: Unique
uniq, us :: UniqSupply
us) -> (# RA_State freeRegs
s{ra_us :: UniqSupply
ra_us = UniqSupply
us}, Unique
uniq #)


-- | Record that a spill instruction was inserted, for profiling.
recordSpill :: SpillReason -> RegM freeRegs ()
recordSpill :: SpillReason -> RegM freeRegs ()
recordSpill spill :: SpillReason
spill
    = (RA_State freeRegs -> (# RA_State freeRegs, () #))
-> RegM freeRegs ()
forall freeRegs a.
(RA_State freeRegs -> (# RA_State freeRegs, a #))
-> RegM freeRegs a
RegM ((RA_State freeRegs -> (# RA_State freeRegs, () #))
 -> RegM freeRegs ())
-> (RA_State freeRegs -> (# RA_State freeRegs, () #))
-> RegM freeRegs ()
forall a b. (a -> b) -> a -> b
$ \s :: RA_State freeRegs
s -> (# RA_State freeRegs
s { ra_spills :: [SpillReason]
ra_spills = SpillReason
spill SpillReason -> [SpillReason] -> [SpillReason]
forall a. a -> [a] -> [a]
: RA_State freeRegs -> [SpillReason]
forall freeRegs. RA_State freeRegs -> [SpillReason]
ra_spills RA_State freeRegs
s}, () #)

-- | Record a created fixup block
recordFixupBlock :: BlockId -> BlockId -> BlockId -> RegM freeRegs ()
recordFixupBlock :: BlockId -> BlockId -> BlockId -> RegM freeRegs ()
recordFixupBlock from :: BlockId
from between :: BlockId
between to :: BlockId
to
    = (RA_State freeRegs -> (# RA_State freeRegs, () #))
-> RegM freeRegs ()
forall freeRegs a.
(RA_State freeRegs -> (# RA_State freeRegs, a #))
-> RegM freeRegs a
RegM ((RA_State freeRegs -> (# RA_State freeRegs, () #))
 -> RegM freeRegs ())
-> (RA_State freeRegs -> (# RA_State freeRegs, () #))
-> RegM freeRegs ()
forall a b. (a -> b) -> a -> b
$ \s :: RA_State freeRegs
s -> (# RA_State freeRegs
s { ra_fixups :: [(BlockId, BlockId, BlockId)]
ra_fixups = (BlockId
from,BlockId
between,BlockId
to) (BlockId, BlockId, BlockId)
-> [(BlockId, BlockId, BlockId)] -> [(BlockId, BlockId, BlockId)]
forall a. a -> [a] -> [a]
: RA_State freeRegs -> [(BlockId, BlockId, BlockId)]
forall freeRegs. RA_State freeRegs -> [(BlockId, BlockId, BlockId)]
ra_fixups RA_State freeRegs
s}, () #)