    {-# OPTIONS -fno-warn-missing-signatures #-}

    -- | Handles joining of a jump instruction to its targets.

    -- 	The first time we encounter a jump to a particular basic block, we
    -- 	record the assignment of temporaries.  The next time we encounter a
    -- 	jump to the same block, we compare our current assignment to the
    -- 	stored one.  They might be different if spilling has occrred in one
    -- 	branch; so some fixup code will be required to match up the assignments.
    --
    module RegAlloc.Linear.JoinToTargets (
    	joinToTargets
    )

    where

    import RegAlloc.Linear.State
    import RegAlloc.Linear.Base
    import RegAlloc.Linear.FreeRegs
    import RegAlloc.Liveness
    import Instruction
    import Reg

    import BlockId
    import Cmm	hiding (RegSet)
    import Digraph
    import Outputable
    import Unique
    import UniqFM
    import UniqSet

    -- | For a jump instruction at the end of a block, generate fixup code so its
    --	vregs are in the correct regs for its destination.
    --
    joinToTargets
    	:: Instruction instr
    	=> BlockMap RegSet		-- ^ maps the unique of the blockid to the set of vregs
    					--	that are known to be live on the entry to each block.

    	-> BlockId			-- ^ id of the current block
    	-> instr			-- ^ branch instr on the end of the source block.

    	-> RegM ([NatBasicBlock instr]	--   fresh blocks of fixup code.
    		, instr)		--   the original branch instruction, but maybe patched to jump
    					--	to a fixup block first.

    joinToTargets block_live id instr

    	-- we only need to worry about jump instructions.
    	| not $ isJumpishInstr instr
    	= return ([], instr)

    	| otherwise
    	= joinToTargets' block_live [] id instr (jumpDestsOfInstr instr)

    -----
    joinToTargets'
    	:: Instruction instr
    	=> BlockMap RegSet		-- ^ maps the unique of the blockid to the set of vregs
    					--	that are known to be live on the entry to each block.

    	-> [NatBasicBlock instr]	-- ^ acc blocks of fixup code.

    	-> BlockId			-- ^ id of the current block
    	-> instr			-- ^ branch instr on the end of the source block.

    	-> [BlockId]			-- ^ branch destinations still to consider.

    	-> RegM ( [NatBasicBlock instr]
    		, instr)

    -- no more targets to consider. all done.
    joinToTargets' _          new_blocks _ instr []
    	= return (new_blocks, instr)

    -- handle a branch target.
    joinToTargets' block_live new_blocks block_id instr (dest:dests)
     = do
     	-- get the map of where the vregs are stored on entry to each basic block.
    	block_assig 	<- getBlockAssigR

    	-- get the assignment on entry to the branch instruction.
    	assig 		<- getAssigR

    	-- adjust the current assignment to remove any vregs that are not live
    	-- on entry to the destination block.
    	let Just live_set 	= lookupBlockEnv block_live dest
    	let still_live uniq _ 	= uniq `elemUniqSet_Directly` live_set
    	let adjusted_assig	= filterUFM_Directly still_live assig

    	-- and free up those registers which are now free.
    	let to_free =
    		[ r 	| (reg, loc) <- ufmToList assig
    			, not (elemUniqSet_Directly reg live_set)
    			, r 	     <- regsOfLoc loc ]

    	case lookupBlockEnv block_assig dest of
    	 Nothing
    	  -> joinToTargets_first
    	  		block_live new_blocks block_id instr dest dests
    			block_assig adjusted_assig to_free

    	 Just (_, dest_assig)
    	  -> joinToTargets_again
    	  		block_live new_blocks block_id instr dest dests
    	  		adjusted_assig dest_assig

    -- this is the first time we jumped to this block.
    joinToTargets_first block_live new_blocks block_id instr dest dests
    	block_assig src_assig
    	(to_free :: [RealReg])

     = do	-- free up the regs that are not live on entry to this block.
      	freeregs 	<- getFreeRegsR
    	let freeregs' 	= foldr releaseReg freeregs to_free

    	-- remember the current assignment on entry to this block.
    	setBlockAssigR (extendBlockEnv block_assig dest
    				(freeregs', src_assig))

    	joinToTargets' block_live new_blocks block_id instr dests

    -- we've jumped to this block before
    joinToTargets_again
    	block_live new_blocks block_id instr dest dests
    	src_assig dest_assig

    	-- the assignments already match, no problem.
    	| ufmToList dest_assig == ufmToList src_assig
    	= joinToTargets' block_live new_blocks block_id instr dests

     	-- assignments don't match, need fixup code
    	| otherwise
    	= do

    		-- make a graph of what things need to be moved where.
    		let graph = makeRegMovementGraph src_assig dest_assig

    		-- look for cycles in the graph. This can happen if regs need to be swapped.
    		-- Note that we depend on the fact that this function does a
    		--	bottom up traversal of the tree-like portions of the graph.
    		--
    		--  eg, if we have
    		--	R1 -> R2 -> R3
    		--
    		--  ie move value in R1 to R2 and value in R2 to R3.
    		--
    		-- We need to do the R2 -> R3 move before R1 -> R2.
    		--
    		let sccs  = stronglyConnCompFromEdgedVerticesR graph

    {-		-- debugging
    		pprTrace
    			("joinToTargets: making fixup code")
    			(vcat	[ text "        in block: "	<> ppr block_id
    				, text " jmp instruction: "	<> ppr instr
    				, text "  src assignment: "	<> ppr src_assig
    				, text " dest assignment: " 	<> ppr dest_assig
    				, text "  movement graph: "	<> ppr graph
    				, text "   sccs of graph: "	<> ppr sccs
    				, text ""])
    			(return ())
    -}
    		delta 		<- getDeltaR
    		fixUpInstrs_ 	<- mapM (handleComponent delta instr) sccs
    		let fixUpInstrs	= concat fixUpInstrs_

    		-- make a new basic block containing the fixup code.
    		--	A the end of the current block we will jump to the fixup one,
    		--	then that will jump to our original destination.
    		fixup_block_id <- getUniqueR
    		let block = BasicBlock (BlockId fixup_block_id)
    				$ fixUpInstrs ++ mkJumpInstr dest

    {-		pprTrace
    			("joinToTargets: fixup code is:")
    			(vcat	[ ppr block
    				, text ""])
    			(return ())
    -}
    		-- if we didn't need any fixups, then don't include the block
    		case fixUpInstrs of
    		 []	-> joinToTargets' block_live new_blocks block_id instr dests

    		 -- patch the original branch instruction so it goes to our
    		 --	fixup block instead.
    		 _	-> let	instr'	=  patchJumpInstr instr
    		 				(\bid -> if bid == dest
    								then BlockId fixup_block_id
    								else dest)

    		 	   in	joinToTargets' block_live (block : new_blocks) block_id instr' dests

    -- | Construct a graph of register\/spill movements.
    --
    -- 	Cyclic components seem to occur only very rarely.
    --
    -- 	We cut some corners by not handling memory-to-memory moves.
    --	This shouldn't happen because every temporary gets its own stack slot.
    --
    makeRegMovementGraph :: RegMap Loc -> RegMap Loc -> [(Unique, Loc, [Loc])]
    makeRegMovementGraph adjusted_assig dest_assig
     = let
     	mkNodes src vreg
    	 = expandNode vreg src
    	 $ lookupWithDefaultUFM_Directly
               	dest_assig
                    (panic "RegAllocLinear.makeRegMovementGraph")
    		vreg

       in	[ node 	| (vreg, src) <- ufmToList adjusted_assig
     		, node <- mkNodes src vreg ]

    -- | Expand out the destination, so InBoth destinations turn into
    --	a combination of InReg and InMem.

    --	The InBoth handling is a little tricky here.  If the destination is
    --	InBoth, then we must ensure that the value ends up in both locations.
    --	An InBoth  destination must conflict with an InReg or InMem source, so
    --	we expand an InBoth destination as necessary.
    --
    --	An InBoth source is slightly different: we only care about the register
    --	that the source value is in, so that we can move it to the destinations.
    --
    expandNode
    	:: a
    	-> Loc 			-- ^ source of move
    	-> Loc 			-- ^ destination of move
    	-> [(a, Loc, [Loc])]

    expandNode vreg loc@(InReg src) (InBoth dst mem)
    	| src == dst = [(vreg, loc, [InMem mem])]
    	| otherwise  = [(vreg, loc, [InReg dst, InMem mem])]

    expandNode vreg loc@(InMem src) (InBoth dst mem)
    	| src == mem = [(vreg, loc, [InReg dst])]
    	| otherwise  = [(vreg, loc, [InReg dst, InMem mem])]

    expandNode _        (InBoth _ src) (InMem dst)
    	| src == dst = [] -- guaranteed to be true

    expandNode _        (InBoth src _) (InReg dst)
    	| src == dst = []

    expandNode vreg     (InBoth src _) dst
    	= expandNode vreg (InReg src) dst

    expandNode vreg src dst
    	| src == dst = []
    	| otherwise  = [(vreg, src, [dst])]

    -- | Generate fixup code for a particular component in the move graph
    --	This component tells us what values need to be moved to what
    --	destinations. We have eliminated any possibility of single-node
    --	cycles in expandNode above.
    --
    handleComponent
    	:: Instruction instr
    	=> Int -> instr -> SCC (Unique, Loc, [Loc]) -> RegM [instr]

    -- If the graph is acyclic then we won't get the swapping problem below.
    --	In this case we can just do the moves directly, and avoid having to
    --	go via a spill slot.
    --
    handleComponent delta _  (AcyclicSCC (vreg, src, dsts))
    	 = mapM (makeMove delta vreg src) dsts

    -- Handle some cyclic moves.
    --	This can happen if we have two regs that need to be swapped.
    --	eg:
    --	     vreg   source loc   dest loc
    --	    (vreg1, InReg r1,    [InReg r2])
    --	    (vreg2, InReg r2,    [InReg r1])
    --
    --	To avoid needing temp register, we just spill all the source regs, then
    --	reaload them into their destination regs.
    --
    --	Note that we can not have cycles that involve memory locations as
    --	sources as single destination because memory locations (stack slots)
    --	are allocated exclusively for a virtual register and therefore can not
    --	require a fixup.
    --
    handleComponent delta instr
    	(CyclicSCC 	( (vreg, InReg sreg, (InReg dreg: _)) : rest))
            -- dest list may have more than one element, if the reg is also InMem.
     = do
    	-- spill the source into its slot
    	(instrSpill, slot)
    			<- spillR (RegReal sreg) vreg

    	-- reload into destination reg
    	instrLoad	<- loadR  (RegReal dreg) slot

    	remainingFixUps <- mapM (handleComponent delta instr)
    				(stronglyConnCompFromEdgedVerticesR rest)

    	-- make sure to do all the reloads after all the spills,
    	--	so we don't end up clobbering the source values.
    	return ([instrSpill] ++ concat remainingFixUps ++ [instrLoad])

    handleComponent _ _ (CyclicSCC _)
     = panic "Register Allocator: handleComponent cyclic"

    -- | Move a vreg between these two locations.
    --
    makeMove
    	:: Instruction instr
    	=> Int 		-- ^ current C stack delta.
    	-> Unique 	-- ^ unique of the vreg that we're moving.
    	-> Loc 		-- ^ source location.
    	-> Loc 		-- ^ destination location.
    	-> RegM instr	-- ^ move instruction.

    makeMove _     vreg (InReg src) (InReg dst)
     = do	recordSpill (SpillJoinRR vreg)
     	return	$ mkRegRegMoveInstr (RegReal src) (RegReal dst)

    makeMove delta vreg (InMem src) (InReg dst)
     = do	recordSpill (SpillJoinRM vreg)
     	return	$ mkLoadInstr  (RegReal dst) delta src

    makeMove delta vreg (InReg src) (InMem dst)
     = do	recordSpill (SpillJoinRM vreg)
    	return	$ mkSpillInstr (RegReal src) delta dst

    -- we don't handle memory to memory moves.
    --	they shouldn't happen because we don't share stack slots between vregs.
    makeMove _     vreg src dst
    	= panic $ "makeMove " ++ show vreg ++ " (" ++ show src ++ ") ("
    		++ show dst ++ ")"
    		++ " we don't handle mem->mem moves."