{-# LANGUAGE CPP, PatternSynonyms #-}

#if !defined(GHC_LOADED_INTO_GHCI)
{-# LANGUAGE UnboxedTuples #-}
#endif

-- | 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)

-- Avoids using unboxed tuples when loading into GHCi
#if !defined(GHC_LOADED_INTO_GHCI)

type RA_Result freeRegs a = (# RA_State freeRegs, a #)

pattern RA_Result :: a -> b -> (# a, b #)
pattern $bRA_Result :: a -> b -> (# a, b #)
$mRA_Result :: forall r a b. (# a, b #) -> (a -> b -> r) -> (Void# -> r) -> r
RA_Result a b = (# a, b #)
{-# COMPLETE RA_Result #-}
#else

data RA_Result freeRegs a = RA_Result {-# UNPACK #-} !(RA_State freeRegs) !a

#endif

-- | The register allocator monad type.
newtype RegM freeRegs a
        = RegM { RegM freeRegs a -> RA_State freeRegs -> RA_Result freeRegs a
unReg :: RA_State freeRegs -> RA_Result 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_Result freeRegs a) -> RegM freeRegs a
forall freeRegs a.
(RA_State freeRegs -> RA_Result freeRegs a) -> RegM freeRegs a
RegM ((RA_State freeRegs -> RA_Result freeRegs a) -> RegM freeRegs a)
-> (RA_State freeRegs -> RA_Result freeRegs a) -> RegM freeRegs a
forall a b. (a -> b) -> a -> b
$ \s :: RA_State freeRegs
s -> RA_State freeRegs -> a -> RA_Result freeRegs a
forall a b. a -> b -> (# a, b #)
RA_Result 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_Result freeRegs b) -> RegM freeRegs b
forall freeRegs a.
(RA_State freeRegs -> RA_Result freeRegs a) -> RegM freeRegs a
RegM ((RA_State freeRegs -> RA_Result freeRegs b) -> RegM freeRegs b)
-> (RA_State freeRegs -> RA_Result 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_Result freeRegs a
forall freeRegs a.
RegM freeRegs a -> RA_State freeRegs -> RA_Result freeRegs a
unReg RegM freeRegs a
m RA_State freeRegs
s of { RA_Result s :: RA_State freeRegs
s a :: a
a -> RegM freeRegs b -> RA_State freeRegs -> RA_Result freeRegs b
forall freeRegs a.
RegM freeRegs a -> RA_State freeRegs -> RA_Result 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_Result a DynFlags) -> RegM a DynFlags
forall freeRegs a.
(RA_State freeRegs -> RA_Result freeRegs a) -> RegM freeRegs a
RegM ((RA_State a -> RA_Result a DynFlags) -> RegM a DynFlags)
-> (RA_State a -> RA_Result a DynFlags) -> RegM a DynFlags
forall a b. (a -> b) -> a -> b
$ \s :: RA_State a
s -> RA_State a -> DynFlags -> RA_Result a DynFlags
forall a b. a -> b -> (# a, b #)
RA_Result 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_Result freeRegs a
forall freeRegs a.
RegM freeRegs a -> RA_State freeRegs -> RA_Result 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
        RA_Result state :: RA_State freeRegs
state returned_thing :: a
returned_thing
         ->     (RA_State freeRegs -> BlockAssignment freeRegs
forall freeRegs. RA_State freeRegs -> BlockAssignment freeRegs
ra_blockassig RA_State freeRegs
state, RA_State freeRegs -> StackMap
forall freeRegs. RA_State freeRegs -> StackMap
ra_stack RA_State freeRegs
state, 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_Result freeRegs (instr, Int))
-> RegM freeRegs (instr, Int)
forall freeRegs a.
(RA_State freeRegs -> RA_Result freeRegs a) -> RegM freeRegs a
RegM ((RA_State freeRegs -> RA_Result freeRegs (instr, Int))
 -> RegM freeRegs (instr, Int))
-> (RA_State freeRegs -> RA_Result 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
stack0} ->
  let dflags :: DynFlags
dflags = RA_State freeRegs -> DynFlags
forall freeRegs. RA_State freeRegs -> DynFlags
ra_DynFlags RA_State freeRegs
s
      (stack1 :: StackMap
stack1,slot :: Int
slot) = StackMap -> Unique -> (StackMap, Int)
getStackSlotFor StackMap
stack0 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
-> (instr, Int) -> RA_Result freeRegs (instr, Int)
forall a b. a -> b -> (# a, b #)
RA_Result RA_State freeRegs
s{ra_stack :: StackMap
ra_stack=StackMap
stack1} (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_Result freeRegs instr)
-> RegM freeRegs instr
forall freeRegs a.
(RA_State freeRegs -> RA_Result freeRegs a) -> RegM freeRegs a
RegM ((RA_State freeRegs -> RA_Result freeRegs instr)
 -> RegM freeRegs instr)
-> (RA_State freeRegs -> RA_Result 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 -> instr -> RA_Result freeRegs instr
forall a b. a -> b -> (# a, b #)
RA_Result 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_Result freeRegs freeRegs)
-> RegM freeRegs freeRegs
forall freeRegs a.
(RA_State freeRegs -> RA_Result freeRegs a) -> RegM freeRegs a
RegM ((RA_State freeRegs -> RA_Result freeRegs freeRegs)
 -> RegM freeRegs freeRegs)
-> (RA_State freeRegs -> RA_Result 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 -> freeRegs -> RA_Result freeRegs freeRegs
forall a b. a -> b -> (# a, b #)
RA_Result RA_State freeRegs
s freeRegs
freeregs

setFreeRegsR :: freeRegs -> RegM freeRegs ()
setFreeRegsR :: freeRegs -> RegM freeRegs ()
setFreeRegsR regs :: freeRegs
regs = (RA_State freeRegs -> RA_Result freeRegs ()) -> RegM freeRegs ()
forall freeRegs a.
(RA_State freeRegs -> RA_Result freeRegs a) -> RegM freeRegs a
RegM ((RA_State freeRegs -> RA_Result freeRegs ()) -> RegM freeRegs ())
-> (RA_State freeRegs -> RA_Result freeRegs ()) -> RegM freeRegs ()
forall a b. (a -> b) -> a -> b
$ \ s :: RA_State freeRegs
s ->
  RA_State freeRegs -> () -> RA_Result freeRegs ()
forall a b. a -> b -> (# a, b #)
RA_Result 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_Result freeRegs (RegMap Loc))
-> RegM freeRegs (RegMap Loc)
forall freeRegs a.
(RA_State freeRegs -> RA_Result freeRegs a) -> RegM freeRegs a
RegM ((RA_State freeRegs -> RA_Result freeRegs (RegMap Loc))
 -> RegM freeRegs (RegMap Loc))
-> (RA_State freeRegs -> RA_Result 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 -> RegMap Loc -> RA_Result freeRegs (RegMap Loc)
forall a b. a -> b -> (# a, b #)
RA_Result 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_Result freeRegs ()) -> RegM freeRegs ()
forall freeRegs a.
(RA_State freeRegs -> RA_Result freeRegs a) -> RegM freeRegs a
RegM ((RA_State freeRegs -> RA_Result freeRegs ()) -> RegM freeRegs ())
-> (RA_State freeRegs -> RA_Result freeRegs ()) -> RegM freeRegs ()
forall a b. (a -> b) -> a -> b
$ \ s :: RA_State freeRegs
s ->
  RA_State freeRegs -> () -> RA_Result freeRegs ()
forall a b. a -> b -> (# a, b #)
RA_Result 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_Result freeRegs (BlockAssignment freeRegs))
-> RegM freeRegs (BlockAssignment freeRegs)
forall freeRegs a.
(RA_State freeRegs -> RA_Result freeRegs a) -> RegM freeRegs a
RegM ((RA_State freeRegs
  -> RA_Result freeRegs (BlockAssignment freeRegs))
 -> RegM freeRegs (BlockAssignment freeRegs))
-> (RA_State freeRegs
    -> RA_Result 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
-> BlockAssignment freeRegs
-> RA_Result freeRegs (BlockAssignment freeRegs)
forall a b. a -> b -> (# a, b #)
RA_Result 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_Result freeRegs ()) -> RegM freeRegs ()
forall freeRegs a.
(RA_State freeRegs -> RA_Result freeRegs a) -> RegM freeRegs a
RegM ((RA_State freeRegs -> RA_Result freeRegs ()) -> RegM freeRegs ())
-> (RA_State freeRegs -> RA_Result freeRegs ()) -> RegM freeRegs ()
forall a b. (a -> b) -> a -> b
$ \ s :: RA_State freeRegs
s ->
  RA_State freeRegs -> () -> RA_Result freeRegs ()
forall a b. a -> b -> (# a, b #)
RA_Result 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_Result freeRegs ()) -> RegM freeRegs ()
forall freeRegs a.
(RA_State freeRegs -> RA_Result freeRegs a) -> RegM freeRegs a
RegM ((RA_State freeRegs -> RA_Result freeRegs ()) -> RegM freeRegs ())
-> (RA_State freeRegs -> RA_Result freeRegs ()) -> RegM freeRegs ()
forall a b. (a -> b) -> a -> b
$ \ s :: RA_State freeRegs
s ->
  RA_State freeRegs -> () -> RA_Result freeRegs ()
forall a b. a -> b -> (# a, b #)
RA_Result RA_State freeRegs
s{ra_delta :: Int
ra_delta = Int
n} ()

getDeltaR :: RegM freeRegs Int
getDeltaR :: RegM freeRegs Int
getDeltaR = (RA_State freeRegs -> RA_Result freeRegs Int) -> RegM freeRegs Int
forall freeRegs a.
(RA_State freeRegs -> RA_Result freeRegs a) -> RegM freeRegs a
RegM ((RA_State freeRegs -> RA_Result freeRegs Int)
 -> RegM freeRegs Int)
-> (RA_State freeRegs -> RA_Result freeRegs Int)
-> RegM freeRegs Int
forall a b. (a -> b) -> a -> b
$ \s :: RA_State freeRegs
s -> RA_State freeRegs -> Int -> RA_Result freeRegs Int
forall a b. a -> b -> (# a, b #)
RA_Result 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_Result freeRegs Unique)
-> RegM freeRegs Unique
forall freeRegs a.
(RA_State freeRegs -> RA_Result freeRegs a) -> RegM freeRegs a
RegM ((RA_State freeRegs -> RA_Result freeRegs Unique)
 -> RegM freeRegs Unique)
-> (RA_State freeRegs -> RA_Result 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 -> Unique -> RA_Result freeRegs Unique
forall a b. a -> b -> (# a, b #)
RA_Result 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_Result freeRegs ()) -> RegM freeRegs ()
forall freeRegs a.
(RA_State freeRegs -> RA_Result freeRegs a) -> RegM freeRegs a
RegM ((RA_State freeRegs -> RA_Result freeRegs ()) -> RegM freeRegs ())
-> (RA_State freeRegs -> RA_Result freeRegs ()) -> RegM freeRegs ()
forall a b. (a -> b) -> a -> b
$ \s :: RA_State freeRegs
s -> RA_State freeRegs -> () -> RA_Result freeRegs ()
forall a b. a -> b -> (# a, b #)
RA_Result (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_Result freeRegs ()) -> RegM freeRegs ()
forall freeRegs a.
(RA_State freeRegs -> RA_Result freeRegs a) -> RegM freeRegs a
RegM ((RA_State freeRegs -> RA_Result freeRegs ()) -> RegM freeRegs ())
-> (RA_State freeRegs -> RA_Result freeRegs ()) -> RegM freeRegs ()
forall a b. (a -> b) -> a -> b
$ \s :: RA_State freeRegs
s -> RA_State freeRegs -> () -> RA_Result freeRegs ()
forall a b. a -> b -> (# a, b #)
RA_Result (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 }) ()