module CgStackery (
	spRel, getVirtSp, getRealSp, setRealSp,
	setRealAndVirtualSp, getSpRelOffset,

	allocPrimStack, allocStackTop, deAllocStackTop,
	adjustStackHW, getFinalStackHW, 
	setStackFrame, getStackFrame,
	mkVirtStkOffsets, mkStkAmodes,
	freeStackSlots, 
	pushUpdateFrame, pushBHUpdateFrame, emitPushUpdateFrame,
    ) where

#include "HsVersions.h"

import CgMonad
import CgUtils
import CgProf
import SMRep
import Cmm
import CmmUtils
import CLabel
import Constants
import Util
import OrdList
import Outputable

import Control.Monad

spRel :: VirtualSpOffset 	-- virtual offset of Sp
      -> VirtualSpOffset 	-- virtual offset of The Thing
      -> WordOff		-- integer offset
spRel sp off = sp - off

setRealAndVirtualSp :: VirtualSpOffset 	-- New real Sp
		     -> Code

setRealAndVirtualSp new_sp 
  = do	{ stk_usg <- getStkUsage
	; setStkUsage (stk_usg {virtSp = new_sp, 
				realSp = new_sp, 
				hwSp   = new_sp}) }

getVirtSp :: FCode VirtualSpOffset
getVirtSp
  = do	{ stk_usg <- getStkUsage
	; return (virtSp stk_usg) }

getRealSp :: FCode VirtualSpOffset
getRealSp
  = do	{ stk_usg <- getStkUsage
	; return (realSp stk_usg) }

setRealSp :: VirtualSpOffset -> Code
setRealSp new_real_sp
  = do	{ stk_usg <- getStkUsage
	; setStkUsage (stk_usg {realSp = new_real_sp}) }

getSpRelOffset :: VirtualSpOffset -> FCode CmmExpr
getSpRelOffset virtual_offset
  = do	{ real_sp <- getRealSp
	; return (cmmRegOffW spReg (spRel real_sp virtual_offset)) }

mkVirtStkOffsets
	  :: VirtualSpOffset 	-- Offset of the last allocated thing
	  -> [(CgRep,a)]		-- things to make offsets for
	  -> (VirtualSpOffset,		-- OUTPUTS: Topmost allocated word
	      [(a, VirtualSpOffset)])	-- things with offsets (voids filtered out)

mkVirtStkOffsets init_Sp_offset things
    = loop init_Sp_offset [] (reverse things)
  where
    loop offset offs [] = (offset,offs)
    loop offset offs ((VoidArg,_):things) = loop offset offs things
	-- ignore Void arguments
    loop offset offs ((rep,t):things)
	= loop thing_slot ((t,thing_slot):offs) things
	where
	  thing_slot = offset + cgRepSizeW rep
	    -- offset of thing is offset+size, because we're 
	    -- growing the stack *downwards* as the offsets increase.

-- | 'mkStkAmodes' is a higher-level version of
-- 'mkVirtStkOffsets'.  It starts from the tail-call locations.
-- It returns a single list of addressing modes for the stack
-- locations, and therefore is in the monad.  It /doesn't/ adjust the
-- high water mark.

mkStkAmodes 
	:: VirtualSpOffset	    -- Tail call positions
	-> [(CgRep,CmmExpr)]	    -- things to make offsets for
	-> FCode (VirtualSpOffset,  -- OUTPUTS: Topmost allocated word
	          CmmStmts)	    -- Assignments to appropriate stk slots

mkStkAmodes tail_Sp things
  = do	{ rSp <- getRealSp
	; let (last_Sp_offset, offsets) = mkVirtStkOffsets tail_Sp things
	      abs_cs = [ CmmStore (cmmRegOffW spReg (spRel rSp offset)) amode
		       | (amode, offset) <- offsets
		       ]
	; returnFC (last_Sp_offset, toOL abs_cs) }

allocPrimStack :: CgRep -> FCode VirtualSpOffset
allocPrimStack rep
  = do	{ stk_usg <- getStkUsage
	; let free_stk = freeStk stk_usg
	; case find_block free_stk of
	     Nothing -> do 
		{ let push_virt_sp = virtSp stk_usg + size
		; setStkUsage (stk_usg { virtSp = push_virt_sp,
					 hwSp   = hwSp stk_usg `max` push_virt_sp })
						-- Adjust high water mark
		; return push_virt_sp }
	     Just slot -> do
		{ setStkUsage (stk_usg { freeStk = delete_block free_stk slot }) 
		; return slot }
	}
  where
    size :: WordOff
    size = cgRepSizeW rep

	-- Find_block looks for a contiguous chunk of free slots
	-- returning the offset of its topmost word
    find_block :: [VirtualSpOffset] -> Maybe VirtualSpOffset
    find_block [] = Nothing
    find_block (slot:slots)
	| take size (slot:slots) == [slot..top_slot]
	= Just top_slot
	| otherwise
	= find_block slots
	where	-- The stack grows downwards, with increasing virtual offsets.
		-- Therefore, the address of a multi-word object is the *highest*
		-- virtual offset it occupies (top_slot below).
	    top_slot = slot+size-1

    delete_block free_stk slot = [ s | s <- free_stk, 
				       (s<=slot-size) || (s>slot) ]
		      -- Retain slots which are not in the range
		      -- slot-size+1..slot

allocStackTop :: WordOff -> FCode ()
allocStackTop size
  = do	{ stk_usg <- getStkUsage
	; let push_virt_sp = virtSp stk_usg + size
	; setStkUsage (stk_usg { virtSp = push_virt_sp,
				 hwSp   = hwSp stk_usg `max` push_virt_sp }) }

deAllocStackTop :: WordOff -> FCode ()
deAllocStackTop size
  = do	{ stk_usg <- getStkUsage
	; let pop_virt_sp = virtSp stk_usg - size
	; setStkUsage (stk_usg { virtSp = pop_virt_sp }) }

adjustStackHW :: VirtualSpOffset -> Code
adjustStackHW offset
  = do	{ stk_usg <- getStkUsage
	; setStkUsage (stk_usg { hwSp = hwSp stk_usg `max` offset }) }

getFinalStackHW :: (VirtualSpOffset -> Code) -> Code
getFinalStackHW fcode
  = do	{ fixC_ (\hw_sp -> do
		{ fcode hw_sp
		; stk_usg <- getStkUsage
		; return (hwSp stk_usg) })
	; return () }

setStackFrame :: VirtualSpOffset -> Code
setStackFrame offset
  = do	{ stk_usg <- getStkUsage
	; setStkUsage (stk_usg { frameSp = offset }) }

getStackFrame :: FCode VirtualSpOffset
getStackFrame
  = do	{ stk_usg <- getStkUsage
	; return (frameSp stk_usg) }

pushUpdateFrame :: CmmExpr -> Code -> Code
pushUpdateFrame updatee code
  = pushSpecUpdateFrame mkUpdInfoLabel updatee code

pushBHUpdateFrame :: CmmExpr -> Code -> Code
pushBHUpdateFrame updatee code
  = pushSpecUpdateFrame mkBHUpdInfoLabel updatee code

pushSpecUpdateFrame :: CLabel -> CmmExpr -> Code -> Code
pushSpecUpdateFrame lbl updatee code
  = do	{
      when debugIsOn $ do
    	{ EndOfBlockInfo _ sequel <- getEndOfBlockInfo ;
    	; MASSERT(case sequel of { OnStack -> True; _ -> False}) }
	; allocStackTop (fixedHdrSize + 
			   sIZEOF_StgUpdateFrame_NoHdr `quot` wORD_SIZE)
	; vsp <- getVirtSp
	; setStackFrame vsp
	; frame_addr <- getSpRelOffset vsp
		-- The location of the lowest-address
		-- word of the update frame itself

                -- NB. we used to set the Sequel to 'UpdateCode' so
                -- that we could jump directly to the update code if
                -- we know that the next frame on the stack is an
                -- update frame.  However, the RTS can sometimes
                -- change an update frame into something else (see
                -- e.g. Note [upd-black-hole] in rts/sm/Scav.c), so we
                -- no longer make this assumption.
	; setEndOfBlockInfo (EndOfBlockInfo vsp OnStack) $
	    do	{ emitSpecPushUpdateFrame lbl frame_addr updatee
		; code }
	}

emitPushUpdateFrame :: CmmExpr -> CmmExpr -> Code
emitPushUpdateFrame = emitSpecPushUpdateFrame mkUpdInfoLabel

emitSpecPushUpdateFrame :: CLabel -> CmmExpr -> CmmExpr -> Code
emitSpecPushUpdateFrame lbl frame_addr updatee = do
	stmtsC [  -- Set the info word
		  CmmStore frame_addr (mkLblExpr lbl)
		, -- And the updatee
		  CmmStore (cmmOffsetB frame_addr off_updatee) updatee ]
	initUpdFrameProf frame_addr

off_updatee :: ByteOff
off_updatee = fixedHdrSize*wORD_SIZE + oFFSET_StgUpdateFrame_updatee

freeStackSlots :: [VirtualSpOffset] -> Code
freeStackSlots extra_free
  = do	{ stk_usg <- getStkUsage
	; let all_free = addFreeSlots (freeStk stk_usg) (sortLe (<=) extra_free)
	; let (new_vsp, new_free) = trim (virtSp stk_usg) all_free
	; setStkUsage (stk_usg { virtSp = new_vsp, freeStk = new_free }) }

addFreeSlots :: [VirtualSpOffset] -> [VirtualSpOffset] -> [VirtualSpOffset]
-- Merge the two, assuming both are in increasing order
addFreeSlots cs [] = cs
addFreeSlots [] ns = ns
addFreeSlots (c:cs) (n:ns)
  | c < n     = c : addFreeSlots cs (n:ns)
  | otherwise = n : addFreeSlots (c:cs) ns

trim :: VirtualSpOffset -> [VirtualSpOffset] -> (VirtualSpOffset, [VirtualSpOffset])
-- Try to trim back the virtual stack pointer, where there is a
-- continuous bunch of free slots at the end of the free list
trim vsp [] = (vsp, [])
trim vsp (slot:slots)
  = case trim vsp slots of
      (vsp', []) 
	| vsp' < slot  -> pprTrace "trim: strange" (ppr vsp <+> ppr (slot:slots))
			  (vsp',   [])
	| vsp' == slot -> (vsp'-1, [])
	| otherwise    -> (vsp',   [slot])
      (vsp', slots')   -> (vsp',   slot:slots')