{-# OPTIONS -w #-}
-- The above warning supression flag is a temporary kludge.
-- While working on this module you are encouraged to remove it and fix
-- any warnings in the module. See
--     http://hackage.haskell.org/trac/ghc/wiki/CodingStyle#Warnings
-- for details

module RegSpill (
	regSpill,
	SpillStats(..),
	accSpillSL
)

where

#include "HsVersions.h"

import RegLiveness
import RegAllocInfo
import MachRegs
import MachInstrs
import Cmm

import State
import Unique
import UniqFM
import UniqSet
import UniqSupply
import Outputable

import Data.List
import Data.Maybe


-- | Spill all these virtual regs to memory
--	TODO: 	see if we can split some of the live ranges instead of just globally
--		spilling the virtual reg.
--
--	TODO:	On ciscy x86 and x86_64 we don't nessesarally have to add a mov instruction
--		when making spills. If an instr is using a spilled virtual we may be able to
--		address the spill slot directly.
--
regSpill
	:: [LiveCmmTop]			-- ^ the code
	-> UniqSet Int			-- ^ available stack slots
	-> UniqSet Reg			-- ^ the regs to spill
	-> UniqSM
		([LiveCmmTop]		-- code will spill instructions
		, UniqSet Int		-- left over slots
		, SpillStats )		-- stats about what happened during spilling

regSpill code slotsFree regs

	-- not enough slots to spill these regs
	| sizeUniqSet slotsFree < sizeUniqSet regs
	= pprPanic "regSpill: out of spill slots!"
		(  text "   regs to spill = " <> ppr (sizeUniqSet regs)
		$$ text "   slots left    = " <> ppr (sizeUniqSet slotsFree))

	| otherwise
	= do
		-- allocate a slot for each of the spilled regs
		let slots	= take (sizeUniqSet regs) $ uniqSetToList slotsFree
		let regSlotMap	= listToUFM
				$ zip (uniqSetToList regs) slots

		-- grab the unique supply from the monad
		us	<- getUs

		-- run the spiller on all the blocks
		let (code', state')	=
			runState (mapM (mapBlockTopM (regSpill_block regSlotMap)) code)
				 (initSpillS us)

		return	( code'
			, minusUniqSet slotsFree (mkUniqSet slots)
			, makeSpillStats state')


regSpill_block regSlotMap (BasicBlock i instrs)
 = do	instrss'	<- mapM (regSpill_instr regSlotMap) instrs
 	return	$ BasicBlock i (concat instrss')

regSpill_instr _	li@(Instr _ Nothing)
 = do	return [li]

regSpill_instr regSlotMap
	(Instr instr (Just live))
 = do
	-- work out which regs are read and written in this instr
	let RU rlRead rlWritten	= regUsage instr

	-- sometimes a register is listed as being read more than once,
	--	nub this so we don't end up inserting two lots of spill code.
	let rsRead_		= nub rlRead
	let rsWritten_		= nub rlWritten

	-- if a reg is modified, it appears in both lists, want to undo this..
	let rsRead		= rsRead_    \\ rsWritten_
	let rsWritten		= rsWritten_ \\ rsRead_
	let rsModify		= intersect rsRead_ rsWritten_

	-- work out if any of the regs being used are currently being spilled.
	let rsSpillRead		= filter (\r -> elemUFM r regSlotMap) rsRead
	let rsSpillWritten	= filter (\r -> elemUFM r regSlotMap) rsWritten
	let rsSpillModify	= filter (\r -> elemUFM r regSlotMap) rsModify

	-- rewrite the instr and work out spill code.
	(instr1, prepost1)	<- mapAccumLM (spillRead   regSlotMap) instr  rsSpillRead
	(instr2, prepost2)	<- mapAccumLM (spillWrite  regSlotMap) instr1 rsSpillWritten
	(instr3, prepost3)	<- mapAccumLM (spillModify regSlotMap) instr2 rsSpillModify

	let (mPrefixes, mPostfixes)	= unzip (prepost1 ++ prepost2 ++ prepost3)
	let prefixes			= concat mPrefixes
	let postfixes			= concat mPostfixes

	-- final code
	let instrs'	=  map (\i -> Instr i Nothing) prefixes
			++ [ Instr instr3 Nothing ]
			++ map (\i -> Instr i Nothing) postfixes

	return
{-		$ pprTrace "* regSpill_instr spill"
			(  text "instr  = " <> ppr instr
			$$ text "read   = " <> ppr rsSpillRead
			$$ text "write  = " <> ppr rsSpillWritten
			$$ text "mod    = " <> ppr rsSpillModify
			$$ text "-- out"
			$$ (vcat $ map ppr instrs')
			$$ text " ")
-}
		$ instrs'


spillRead regSlotMap instr reg
	| Just slot	<- lookupUFM regSlotMap reg
	= do 	(instr', nReg)	<- patchInstr reg instr

		modify $ \s -> s
			{ stateSpillSL 	= addToUFM_C accSpillSL (stateSpillSL s) reg (reg, 0, 1) }

	 	return	( instr'
			, ( [RELOAD slot nReg]
			  , []) )

	| otherwise	= panic "RegSpill.spillRead: no slot defined for spilled reg"

spillWrite regSlotMap instr reg
	| Just slot	<- lookupUFM regSlotMap reg
	= do 	(instr', nReg)	<- patchInstr reg instr

		modify $ \s -> s
			{ stateSpillSL 	= addToUFM_C accSpillSL (stateSpillSL s) reg (reg, 1, 0) }

	 	return	( instr'
			, ( []
			  , [SPILL nReg slot]))

	| otherwise	= panic "RegSpill.spillWrite: no slot defined for spilled reg"

spillModify regSlotMap instr reg
	| Just slot	<- lookupUFM regSlotMap reg
	= do	(instr', nReg)	<- patchInstr reg instr

		modify $ \s -> s
			{ stateSpillSL 	= addToUFM_C accSpillSL (stateSpillSL s) reg (reg, 1, 1) }

		return	( instr'
			, ( [RELOAD slot nReg]
			  , [SPILL nReg slot]))

	| otherwise	= panic "RegSpill.spillModify: no slot defined for spilled reg"



-- | rewrite uses of this virtual reg in an instr to use a different virtual reg
patchInstr :: Reg -> Instr -> SpillM (Instr, Reg)
patchInstr reg instr
 = do	nUnique		<- newUnique
 	let nReg	= renameVirtualReg nUnique reg
	let instr'	= patchReg1 reg nReg instr
	return		(instr', nReg)

patchReg1 :: Reg -> Reg -> Instr -> Instr
patchReg1 old new instr
 = let	patchF r
		| r == old	= new
		| otherwise	= r
   in	patchRegs instr patchF


------------------------------------------------------
-- Spiller monad

data SpillS
	= SpillS
	{ stateUS	:: UniqSupply
	, stateSpillSL	:: UniqFM (Reg, Int, Int) } -- ^ spilled reg vs number of times vreg was loaded, stored

initSpillS uniqueSupply
	= SpillS
	{ stateUS	= uniqueSupply
	, stateSpillSL	= emptyUFM }

type SpillM a	= State SpillS a

newUnique :: SpillM Unique
newUnique
 = do	us	<- gets stateUS
 	case splitUniqSupply us of
	 (us1, us2)
	  -> do let uniq = uniqFromSupply us1
	  	modify $ \s -> s { stateUS = us2 }
		return uniq

accSpillSL (r1, s1, l1) (r2, s2, l2)
	= (r1, s1 + s2, l1 + l2)


----------------------------------------------------
-- Spiller stats

data SpillStats
	= SpillStats
	{ spillStoreLoad	:: UniqFM (Reg, Int, Int) }

makeSpillStats :: SpillS -> SpillStats
makeSpillStats s
	= SpillStats
	{ spillStoreLoad	= stateSpillSL s }

instance Outputable SpillStats where
 ppr stats
 	= (vcat $ map (\(r, s, l) -> ppr r <+> int s <+> int l)
			$ eltsUFM (spillStoreLoad stats))