{-# LANGUAGE BangPatterns, RecordWildCards, GADTs #-}
module CmmLayoutStack (
       cmmLayoutStack, setInfoTableStackMap
  ) where

import GhcPrelude hiding ((<*>))

import StgCmmUtils      ( callerSaveVolatileRegs ) -- XXX layering violation
import StgCmmForeign    ( saveThreadState, loadThreadState ) -- XXX layering violation

import BasicTypes
import Cmm
import CmmInfo
import BlockId
import CLabel
import CmmUtils
import MkGraph
import ForeignCall
import CmmLive
import CmmProcPoint
import SMRep
import Hoopl.Block
import Hoopl.Collections
import Hoopl.Dataflow
import Hoopl.Graph
import Hoopl.Label
import UniqSupply
import StgCmmUtils      ( newTemp )
import Maybes
import UniqFM
import Util

import DynFlags
import FastString
import Outputable hiding ( isEmpty )
import qualified Data.Set as Set
import Control.Monad.Fix
import Data.Array as Array
import Data.Bits
import Data.List (nub)

{- Note [Stack Layout]

The job of this pass is to

 - replace references to abstract stack Areas with fixed offsets from Sp.

 - replace the CmmHighStackMark constant used in the stack check with
   the maximum stack usage of the proc.

 - save any variables that are live across a call, and reload them as
   necessary.

Before stack allocation, local variables remain live across native
calls (CmmCall{ cmm_cont = Just _ }), and after stack allocation local
variables are clobbered by native calls.

We want to do stack allocation so that as far as possible
 - stack use is minimized, and
 - unnecessary stack saves and loads are avoided.

The algorithm we use is a variant of linear-scan register allocation,
where the stack is our register file.

We proceed in two passes, see Note [Two pass approach] for why they are not easy
to merge into one.

Pass 1:

 - First, we do a liveness analysis, which annotates every block with
   the variables live on entry to the block.

 - We traverse blocks in reverse postorder DFS; that is, we visit at
   least one predecessor of a block before the block itself.  The
   stack layout flowing from the predecessor of the block will
   determine the stack layout on entry to the block.

 - We maintain a data structure

     Map Label StackMap

   which describes the contents of the stack and the stack pointer on
   entry to each block that is a successor of a block that we have
   visited.

 - For each block we visit:

    - Look up the StackMap for this block.

    - If this block is a proc point (or a call continuation, if we aren't
      splitting proc points), we need to reload all the live variables from the
      stack - but this is done in Pass 2, which calculates more precise liveness
      information (see description of Pass 2).

    - Walk forwards through the instructions:
      - At an assignment  x = Sp[loc]
        - Record the fact that Sp[loc] contains x, so that we won't
          need to save x if it ever needs to be spilled.
      - At an assignment  x = E
        - If x was previously on the stack, it isn't any more
      - At the last node, if it is a call or a jump to a proc point
        - Lay out the stack frame for the call (see setupStackFrame)
        - emit instructions to save all the live variables
        - Remember the StackMaps for all the successors
        - emit an instruction to adjust Sp
      - If the last node is a branch, then the current StackMap is the
        StackMap for the successors.

    - Manifest Sp: replace references to stack areas in this block
      with real Sp offsets. We cannot do this until we have laid out
      the stack area for the successors above.

      In this phase we also eliminate redundant stores to the stack;
      see elimStackStores.

  - There is one important gotcha: sometimes we'll encounter a control
    transfer to a block that we've already processed (a join point),
    and in that case we might need to rearrange the stack to match
    what the block is expecting. (exactly the same as in linear-scan
    register allocation, except here we have the luxury of an infinite
    supply of temporary variables).

  - Finally, we update the magic CmmHighStackMark constant with the
    stack usage of the function, and eliminate the whole stack check
    if there was no stack use. (in fact this is done as part of the
    main traversal, by feeding the high-water-mark output back in as
    an input. I hate cyclic programming, but it's just too convenient
    sometimes.)

  There are plenty of tricky details: update frames, proc points, return
  addresses, foreign calls, and some ad-hoc optimisations that are
  convenient to do here and effective in common cases.  Comments in the
  code below explain these.

Pass 2:

- Calculate live registers, but taking into account that nothing is live at the
  entry to a proc point.

- At each proc point and call continuation insert reloads of live registers from
  the stack (they were saved by Pass 1).

Note [Two pass approach]

The main reason for Pass 2 is being able to insert only the reloads that are
needed and the fact that the two passes need different liveness information.
Let's consider an example:

  .....
   \ /
    D   <- proc point
   / \
  E   F
   \ /
    G   <- proc point
    |
    X

Pass 1 needs liveness assuming that local variables are preserved across calls.
This is important because it needs to save any local registers to the stack
(e.g., if register a is used in block X, it must be saved before any native
call).
However, for Pass 2, where we want to reload registers from stack (in a proc
point), this is overly conservative and would lead us to generate reloads in D
for things used in X, even though we're going to generate reloads in G anyway
(since it's also a proc point).
So Pass 2 calculates liveness knowing that nothing is live at the entry to a
proc point. This means that in D we only need to reload things used in E or F.
This can be quite important, for an extreme example see testcase for #3294.

Merging the two passes is not trivial - Pass 2 is a backward rewrite and Pass 1
is a forward one. Furthermore, Pass 1 is creating code that uses local registers
(saving them before a call), which the liveness analysis for Pass 2 must see to
be correct.

-}

-- All stack locations are expressed as positive byte offsets from the
-- "base", which is defined to be the address above the return address
-- on the stack on entry to this CmmProc.
--
-- Lower addresses have higher StackLocs.
--
type StackLoc = ByteOff

{-
 A StackMap describes the stack at any given point.  At a continuation
 it has a particular layout, like this:

         |             | <- base
         |-------------|
         |     ret0    | <- base + 8
         |-------------|
         .  upd frame  . <- base + sm_ret_off
         |-------------|
         |             |
         .    vars     .
         . (live/dead) .
         |             | <- base + sm_sp - sm_args
         |-------------|
         |    ret1     |
         .  ret vals   . <- base + sm_sp    (<--- Sp points here)
         |-------------|

Why do we include the final return address (ret0) in our stack map?  I
have absolutely no idea, but it seems to be done that way consistently
in the rest of the code generator, so I played along here. --SDM

Note that we will be constructing an info table for the continuation
(ret1), which needs to describe the stack down to, but not including,
the update frame (or ret0, if there is no update frame).
-}

data StackMap = StackMap
 {  sm_sp   :: StackLoc
       -- ^ the offset of Sp relative to the base on entry
       -- to this block.
 ,  sm_args :: ByteOff
       -- ^ the number of bytes of arguments in the area for this block
       -- Defn: the offset of young(L) relative to the base is given by
       -- (sm_sp - sm_args) of the StackMap for block L.
 ,  sm_ret_off :: ByteOff
       -- ^ Number of words of stack that we do not describe with an info
       -- table, because it contains an update frame.
 ,  sm_regs :: UniqFM (LocalReg,StackLoc)
       -- ^ regs on the stack
 }

instance Outputable StackMap where
  ppr StackMap{..} =
     text "Sp = " <> int sm_sp $$
     text "sm_args = " <> int sm_args $$
     text "sm_ret_off = " <> int sm_ret_off $$
     text "sm_regs = " <> pprUFM sm_regs ppr

cmmLayoutStack :: DynFlags -> ProcPointSet -> ByteOff -> CmmGraph
               -> UniqSM (CmmGraph, LabelMap StackMap)
cmmLayoutStack dflags procpoints entry_args
               graph@(CmmGraph { g_entry = entry })
  = do
    -- We need liveness info. Dead assignments are removed later
    -- by the sinking pass.
    let liveness = cmmLocalLiveness dflags graph
        blocks = revPostorder graph

    (final_stackmaps, _final_high_sp, new_blocks) <-
          mfix $ \ ~(rec_stackmaps, rec_high_sp, _new_blocks) ->
            layout dflags procpoints liveness entry entry_args
                   rec_stackmaps rec_high_sp blocks

    blocks_with_reloads <-
        insertReloadsAsNeeded dflags procpoints final_stackmaps entry new_blocks
    new_blocks' <- mapM (lowerSafeForeignCall dflags) blocks_with_reloads
    return (ofBlockList entry new_blocks', final_stackmaps)

-- -----------------------------------------------------------------------------
-- Pass 1
-- -----------------------------------------------------------------------------

layout :: DynFlags
       -> LabelSet                      -- proc points
       -> LabelMap CmmLocalLive         -- liveness
       -> BlockId                       -- entry
       -> ByteOff                       -- stack args on entry

       -> LabelMap StackMap             -- [final] stack maps
       -> ByteOff                       -- [final] Sp high water mark

       -> [CmmBlock]                    -- [in] blocks

       -> UniqSM
          ( LabelMap StackMap           -- [out] stack maps
          , ByteOff                     -- [out] Sp high water mark
          , [CmmBlock]                  -- [out] new blocks
          )

layout dflags procpoints liveness entry entry_args final_stackmaps final_sp_high blocks
  = go blocks init_stackmap entry_args []
  where
    (updfr, cont_info)  = collectContInfo blocks

    init_stackmap = mapSingleton entry StackMap{ sm_sp   = entry_args
                                               , sm_args = entry_args
                                               , sm_ret_off = updfr
                                               , sm_regs = emptyUFM
                                               }

    go [] acc_stackmaps acc_hwm acc_blocks
      = return (acc_stackmaps, acc_hwm, acc_blocks)

    go (b0 : bs) acc_stackmaps acc_hwm acc_blocks
      = do
       let (entry0@(CmmEntry entry_lbl tscope), middle0, last0) = blockSplit b0

       let stack0@StackMap { sm_sp = sp0 }
               = mapFindWithDefault
                     (pprPanic "no stack map for" (ppr entry_lbl))
                     entry_lbl acc_stackmaps

       -- (a) Update the stack map to include the effects of
       --     assignments in this block
       let stack1 = foldBlockNodesF (procMiddle acc_stackmaps) middle0 stack0

       -- (b) Look at the last node and if we are making a call or
       --     jumping to a proc point, we must save the live
       --     variables, adjust Sp, and construct the StackMaps for
       --     each of the successor blocks.  See handleLastNode for
       --     details.
       (middle1, sp_off, last1, fixup_blocks, out)
           <- handleLastNode dflags procpoints liveness cont_info
                             acc_stackmaps stack1 tscope middle0 last0

       -- (c) Manifest Sp: run over the nodes in the block and replace
       --     CmmStackSlot with CmmLoad from Sp with a concrete offset.
       --
       -- our block:
       --    middle0          -- the original middle nodes
       --    middle1          -- live variable saves from handleLastNode
       --    Sp = Sp + sp_off -- Sp adjustment goes here
       --    last1            -- the last node
       --
       let middle_pre = blockToList $ foldl' blockSnoc middle0 middle1

       let final_blocks =
               manifestSp dflags final_stackmaps stack0 sp0 final_sp_high
                          entry0 middle_pre sp_off last1 fixup_blocks

       let acc_stackmaps' = mapUnion acc_stackmaps out

           -- If this block jumps to the GC, then we do not take its
           -- stack usage into account for the high-water mark.
           -- Otherwise, if the only stack usage is in the stack-check
           -- failure block itself, we will do a redundant stack
           -- check.  The stack has a buffer designed to accommodate
           -- the largest amount of stack needed for calling the GC.
           --
           this_sp_hwm | isGcJump last0 = 0
                       | otherwise      = sp0 - sp_off

           hwm' = maximum (acc_hwm : this_sp_hwm : map sm_sp (mapElems out))

       go bs acc_stackmaps' hwm' (final_blocks ++ acc_blocks)

-- -----------------------------------------------------------------------------

-- Not foolproof, but GCFun is the culprit we most want to catch
isGcJump :: CmmNode O C -> Bool
isGcJump (CmmCall { cml_target = CmmReg (CmmGlobal l) })
  = l == GCFun || l == GCEnter1
isGcJump _something_else = False

-- -----------------------------------------------------------------------------

-- This doesn't seem right somehow.  We need to find out whether this
-- proc will push some update frame material at some point, so that we
-- can avoid using that area of the stack for spilling.  The
-- updfr_space field of the CmmProc *should* tell us, but it doesn't
-- (I think maybe it gets filled in later when we do proc-point
-- splitting).
--
-- So we'll just take the max of all the cml_ret_offs.  This could be
-- unnecessarily pessimistic, but probably not in the code we
-- generate.

collectContInfo :: [CmmBlock] -> (ByteOff, LabelMap ByteOff)
collectContInfo blocks
  = (maximum ret_offs, mapFromList (catMaybes mb_argss))
 where
  (mb_argss, ret_offs) = mapAndUnzip get_cont blocks

  get_cont :: Block CmmNode x C -> (Maybe (Label, ByteOff), ByteOff)
  get_cont b =
     case lastNode b of
        CmmCall { cml_cont = Just l, .. }
           -> (Just (l, cml_ret_args), cml_ret_off)
        CmmForeignCall { .. }
           -> (Just (succ, ret_args), ret_off)
        _other -> (Nothing, 0)

-- -----------------------------------------------------------------------------
-- Updating the StackMap from middle nodes

-- Look for loads from stack slots, and update the StackMap.  This is
-- purely for optimisation reasons, so that we can avoid saving a
-- variable back to a different stack slot if it is already on the
-- stack.
--
-- This happens a lot: for example when function arguments are passed
-- on the stack and need to be immediately saved across a call, we
-- want to just leave them where they are on the stack.
--
procMiddle :: LabelMap StackMap -> CmmNode e x -> StackMap -> StackMap
procMiddle stackmaps node sm
  = case node of
     CmmAssign (CmmLocal r) (CmmLoad (CmmStackSlot area off) _)
       -> sm { sm_regs = addToUFM (sm_regs sm) r (r,loc) }
        where loc = getStackLoc area off stackmaps
     CmmAssign (CmmLocal r) _other
       -> sm { sm_regs = delFromUFM (sm_regs sm) r }
     _other
       -> sm

getStackLoc :: Area -> ByteOff -> LabelMap StackMap -> StackLoc
getStackLoc Old       n _         = n
getStackLoc (Young l) n stackmaps =
  case mapLookup l stackmaps of
    Nothing -> pprPanic "getStackLoc" (ppr l)
    Just sm -> sm_sp sm - sm_args sm + n

-- -----------------------------------------------------------------------------
-- Handling stack allocation for a last node

-- We take a single last node and turn it into:
--
--    C1 (some statements)
--    Sp = Sp + N
--    C2 (some more statements)
--    call f()          -- the actual last node
--
-- plus possibly some more blocks (we may have to add some fixup code
-- between the last node and the continuation).
--
-- C1: is the code for saving the variables across this last node onto
-- the stack, if the continuation is a call or jumps to a proc point.
--
-- C2: if the last node is a safe foreign call, we have to inject some
-- extra code that goes *after* the Sp adjustment.

handleLastNode
   :: DynFlags -> ProcPointSet -> LabelMap CmmLocalLive -> LabelMap ByteOff
   -> LabelMap StackMap -> StackMap -> CmmTickScope
   -> Block CmmNode O O
   -> CmmNode O C
   -> UniqSM
      ( [CmmNode O O]      -- nodes to go *before* the Sp adjustment
      , ByteOff            -- amount to adjust Sp
      , CmmNode O C        -- new last node
      , [CmmBlock]         -- new blocks
      , LabelMap StackMap  -- stackmaps for the continuations
      )

handleLastNode dflags procpoints liveness cont_info stackmaps
               stack0@StackMap { sm_sp = sp0 } tscp middle last
 = case last of
    --  At each return / tail call,
    --  adjust Sp to point to the last argument pushed, which
    --  is cml_args, after popping any other junk from the stack.
    CmmCall{ cml_cont = Nothing, .. } -> do
      let sp_off = sp0 - cml_args
      return ([], sp_off, last, [], mapEmpty)

    --  At each CmmCall with a continuation:
    CmmCall{ cml_cont = Just cont_lbl, .. } ->
       return $ lastCall cont_lbl cml_args cml_ret_args cml_ret_off

    CmmForeignCall{ succ = cont_lbl, .. } -> do
       return $ lastCall cont_lbl (wORD_SIZE dflags) ret_args ret_off
            -- one word of args: the return address

    CmmBranch {}     ->  handleBranches
    CmmCondBranch {} ->  handleBranches
    CmmSwitch {}     ->  handleBranches

  where
     -- Calls and ForeignCalls are handled the same way:
     lastCall :: BlockId -> ByteOff -> ByteOff -> ByteOff
              -> ( [CmmNode O O]
                 , ByteOff
                 , CmmNode O C
                 , [CmmBlock]
                 , LabelMap StackMap
                 )
     lastCall lbl cml_args cml_ret_args cml_ret_off
      =  ( assignments
         , spOffsetForCall sp0 cont_stack cml_args
         , last
         , [] -- no new blocks
         , mapSingleton lbl cont_stack )
      where
         (assignments, cont_stack) = prepareStack lbl cml_ret_args cml_ret_off

     prepareStack lbl cml_ret_args cml_ret_off
       | Just cont_stack <- mapLookup lbl stackmaps
             -- If we have already seen this continuation before, then
             -- we just have to make the stack look the same:
       = (fixupStack stack0 cont_stack, cont_stack)
             -- Otherwise, we have to allocate the stack frame
       | otherwise
       = (save_assignments, new_cont_stack)
       where
        (new_cont_stack, save_assignments)
           = setupStackFrame dflags lbl liveness cml_ret_off cml_ret_args stack0

     -- For other last nodes (branches), if any of the targets is a
     -- proc point, we have to set up the stack to match what the proc
     -- point is expecting.
     --
     handleBranches :: UniqSM ( [CmmNode O O]
                                , ByteOff
                                , CmmNode O C
                                , [CmmBlock]
                                , LabelMap StackMap )

     handleBranches
         -- Note [diamond proc point]
       | Just l <- futureContinuation middle
       , (nub $ filter (`setMember` procpoints) $ successors last) == [l]
       = do
         let cont_args = mapFindWithDefault 0 l cont_info
             (assigs, cont_stack) = prepareStack l cont_args (sm_ret_off stack0)
             out = mapFromList [ (l', cont_stack)
                               | l' <- successors last ]
         return ( assigs
                , spOffsetForCall sp0 cont_stack (wORD_SIZE dflags)
                , last
                , []
                , out)

        | otherwise = do
          pps <- mapM handleBranch (successors last)
          let lbl_map :: LabelMap Label
              lbl_map = mapFromList [ (l,tmp) | (l,tmp,_,_) <- pps ]
              fix_lbl l = mapFindWithDefault l l lbl_map
          return ( []
                 , 0
                 , mapSuccessors fix_lbl last
                 , concat [ blk | (_,_,_,blk) <- pps ]
                 , mapFromList [ (l, sm) | (l,_,sm,_) <- pps ] )

     -- For each successor of this block
     handleBranch :: BlockId -> UniqSM (BlockId, BlockId, StackMap, [CmmBlock])
     handleBranch l
        --   (a) if the successor already has a stackmap, we need to
        --       shuffle the current stack to make it look the same.
        --       We have to insert a new block to make this happen.
        | Just stack2 <- mapLookup l stackmaps
        = do
             let assigs = fixupStack stack0 stack2
             (tmp_lbl, block) <- makeFixupBlock dflags sp0 l stack2 tscp assigs
             return (l, tmp_lbl, stack2, block)

        --   (b) if the successor is a proc point, save everything
        --       on the stack.
        | l `setMember` procpoints
        = do
             let cont_args = mapFindWithDefault 0 l cont_info
                 (stack2, assigs) =
                      setupStackFrame dflags l liveness (sm_ret_off stack0)
                                                        cont_args stack0
             (tmp_lbl, block) <- makeFixupBlock dflags sp0 l stack2 tscp assigs
             return (l, tmp_lbl, stack2, block)

        --   (c) otherwise, the current StackMap is the StackMap for
        --       the continuation.  But we must remember to remove any
        --       variables from the StackMap that are *not* live at
        --       the destination, because this StackMap might be used
        --       by fixupStack if this is a join point.
        | otherwise = return (l, l, stack1, [])
        where live = mapFindWithDefault (panic "handleBranch") l liveness
              stack1 = stack0 { sm_regs = filterUFM is_live (sm_regs stack0) }
              is_live (r,_) = r `elemRegSet` live

makeFixupBlock :: DynFlags -> ByteOff -> Label -> StackMap
               -> CmmTickScope -> [CmmNode O O]
               -> UniqSM (Label, [CmmBlock])
makeFixupBlock dflags sp0 l stack tscope assigs
  | null assigs && sp0 == sm_sp stack = return (l, [])
  | otherwise = do
    tmp_lbl <- newBlockId
    let sp_off = sp0 - sm_sp stack
        block = blockJoin (CmmEntry tmp_lbl tscope)
                          ( maybeAddSpAdj dflags sp0 sp_off
                           $ blockFromList assigs )
                          (CmmBranch l)
    return (tmp_lbl, [block])

-- Sp is currently pointing to current_sp,
-- we want it to point to
--    (sm_sp cont_stack - sm_args cont_stack + args)
-- so the difference is
--    sp0 - (sm_sp cont_stack - sm_args cont_stack + args)
spOffsetForCall :: ByteOff -> StackMap -> ByteOff -> ByteOff
spOffsetForCall current_sp cont_stack args
  = current_sp - (sm_sp cont_stack - sm_args cont_stack + args)

-- | create a sequence of assignments to establish the new StackMap,
-- given the old StackMap.
fixupStack :: StackMap -> StackMap -> [CmmNode O O]
fixupStack old_stack new_stack = concatMap move new_locs
 where
     old_map  = sm_regs old_stack
     new_locs = stackSlotRegs new_stack

     move (r,n)
       | Just (_,m) <- lookupUFM old_map r, n == m = []
       | otherwise = [CmmStore (CmmStackSlot Old n)
                               (CmmReg (CmmLocal r))]

setupStackFrame
             :: DynFlags
             -> BlockId                 -- label of continuation
             -> LabelMap CmmLocalLive   -- liveness
             -> ByteOff      -- updfr
             -> ByteOff      -- bytes of return values on stack
             -> StackMap     -- current StackMap
             -> (StackMap, [CmmNode O O])

setupStackFrame dflags lbl liveness updfr_off ret_args stack0
  = (cont_stack, assignments)
  where
      -- get the set of LocalRegs live in the continuation
      live = mapFindWithDefault Set.empty lbl liveness

      -- the stack from the base to updfr_off is off-limits.
      -- our new stack frame contains:
      --   * saved live variables
      --   * the return address [young(C) + 8]
      --   * the args for the call,
      --     which are replaced by the return values at the return
      --     point.

      -- everything up to updfr_off is off-limits
      -- stack1 contains updfr_off, plus everything we need to save
      (stack1, assignments) = allocate dflags updfr_off live stack0

      -- And the Sp at the continuation is:
      --   sm_sp stack1 + ret_args
      cont_stack = stack1{ sm_sp = sm_sp stack1 + ret_args
                         , sm_args = ret_args
                         , sm_ret_off = updfr_off
                         }

-- -----------------------------------------------------------------------------
-- Note [diamond proc point]
--
-- This special case looks for the pattern we get from a typical
-- tagged case expression:
--
--    Sp[young(L1)] = L1
--    if (R1 & 7) != 0 goto L1 else goto L2
--  L2:
--    call [R1] returns to L1
--  L1: live: {y}
--    x = R1
--
-- If we let the generic case handle this, we get
--
--    Sp[-16] = L1
--    if (R1 & 7) != 0 goto L1a else goto L2
--  L2:
--    Sp[-8] = y
--    Sp = Sp - 16
--    call [R1] returns to L1
--  L1a:
--    Sp[-8] = y
--    Sp = Sp - 16
--    goto L1
--  L1:
--    x = R1
--
-- The code for saving the live vars is duplicated in each branch, and
-- furthermore there is an extra jump in the fast path (assuming L1 is
-- a proc point, which it probably is if there is a heap check).
--
-- So to fix this we want to set up the stack frame before the
-- conditional jump.  How do we know when to do this, and when it is
-- safe?  The basic idea is, when we see the assignment
--
--   Sp[young(L)] = L
--
-- we know that
--   * we are definitely heading for L
--   * there can be no more reads from another stack area, because young(L)
--     overlaps with it.
--
-- We don't necessarily know that everything live at L is live now
-- (some might be assigned between here and the jump to L).  So we
-- simplify and only do the optimisation when we see
--
--   (1) a block containing an assignment of a return address L
--   (2) ending in a branch where one (and only) continuation goes to L,
--       and no other continuations go to proc points.
--
-- then we allocate the stack frame for L at the end of the block,
-- before the branch.
--
-- We could generalise (2), but that would make it a bit more
-- complicated to handle, and this currently catches the common case.

futureContinuation :: Block CmmNode O O -> Maybe BlockId
futureContinuation middle = foldBlockNodesB f middle Nothing
   where f :: CmmNode a b -> Maybe BlockId -> Maybe BlockId
         f (CmmStore (CmmStackSlot (Young l) _) (CmmLit (CmmBlock _))) _
               = Just l
         f _ r = r

-- -----------------------------------------------------------------------------
-- Saving live registers

-- | Given a set of live registers and a StackMap, save all the registers
-- on the stack and return the new StackMap and the assignments to do
-- the saving.
--
allocate :: DynFlags -> ByteOff -> LocalRegSet -> StackMap
         -> (StackMap, [CmmNode O O])
allocate dflags ret_off live stackmap@StackMap{ sm_sp = sp0
                                              , sm_regs = regs0 }
 =
   -- we only have to save regs that are not already in a slot
   let to_save = filter (not . (`elemUFM` regs0)) (Set.elems live)
       regs1   = filterUFM (\(r,_) -> elemRegSet r live) regs0
   in

   -- make a map of the stack
   let stack = reverse $ Array.elems $
               accumArray (\_ x -> x) Empty (1, toWords dflags (max sp0 ret_off)) $
                 ret_words ++ live_words
            where ret_words =
                   [ (x, Occupied)
                   | x <- [ 1 .. toWords dflags ret_off] ]
                  live_words =
                   [ (toWords dflags x, Occupied)
                   | (r,off) <- nonDetEltsUFM regs1,
                   -- See Note [Unique Determinism and code generation]
                     let w = localRegBytes dflags r,
                     x <- [ off, off - wORD_SIZE dflags .. off - w + 1] ]
   in

   -- Pass over the stack: find slots to save all the new live variables,
   -- choosing the oldest slots first (hence a foldr).
   let
       save slot ([], stack, n, assigs, regs) -- no more regs to save
          = ([], slot:stack, plusW dflags n 1, assigs, regs)
       save slot (to_save, stack, n, assigs, regs)
          = case slot of
               Occupied ->  (to_save, Occupied:stack, plusW dflags n 1, assigs, regs)
               Empty
                 | Just (stack', r, to_save') <-
                       select_save to_save (slot:stack)
                 -> let assig = CmmStore (CmmStackSlot Old n')
                                         (CmmReg (CmmLocal r))
                        n' = plusW dflags n 1
                   in
                        (to_save', stack', n', assig : assigs, (r,(r,n')):regs)

                 | otherwise
                 -> (to_save, slot:stack, plusW dflags n 1, assigs, regs)

       -- we should do better here: right now we'll fit the smallest first,
       -- but it would make more sense to fit the biggest first.
       select_save :: [LocalReg] -> [StackSlot]
                   -> Maybe ([StackSlot], LocalReg, [LocalReg])
       select_save regs stack = go regs []
         where go []     _no_fit = Nothing
               go (r:rs) no_fit
                 | Just rest <- dropEmpty words stack
                 = Just (replicate words Occupied ++ rest, r, rs++no_fit)
                 | otherwise
                 = go rs (r:no_fit)
                 where words = localRegWords dflags r

       -- fill in empty slots as much as possible
       (still_to_save, save_stack, n, save_assigs, save_regs)
          = foldr save (to_save, [], 0, [], []) stack

       -- push any remaining live vars on the stack
       (push_sp, push_assigs, push_regs)
          = foldr push (n, [], []) still_to_save
          where
              push r (n, assigs, regs)
                = (n', assig : assigs, (r,(r,n')) : regs)
                where
                  n' = n + localRegBytes dflags r
                  assig = CmmStore (CmmStackSlot Old n')
                                   (CmmReg (CmmLocal r))

       trim_sp
          | not (null push_regs) = push_sp
          | otherwise
          = plusW dflags n (- length (takeWhile isEmpty save_stack))

       final_regs = regs1 `addListToUFM` push_regs
                          `addListToUFM` save_regs

   in
  -- XXX should be an assert
   if ( n /= max sp0 ret_off ) then pprPanic "allocate" (ppr n <+> ppr sp0 <+> ppr ret_off) else

   if (trim_sp .&. (wORD_SIZE dflags - 1)) /= 0  then pprPanic "allocate2" (ppr trim_sp <+> ppr final_regs <+> ppr push_sp) else

   ( stackmap { sm_regs = final_regs , sm_sp = trim_sp }
   , push_assigs ++ save_assigs )

-- -----------------------------------------------------------------------------
-- Manifesting Sp

-- | Manifest Sp: turn all the CmmStackSlots into CmmLoads from Sp.  The
-- block looks like this:
--
--    middle_pre       -- the middle nodes
--    Sp = Sp + sp_off -- Sp adjustment goes here
--    last             -- the last node
--
-- And we have some extra blocks too (that don't contain Sp adjustments)
--
-- The adjustment for middle_pre will be different from that for
-- middle_post, because the Sp adjustment intervenes.
--
manifestSp
   :: DynFlags
   -> LabelMap StackMap  -- StackMaps for other blocks
   -> StackMap           -- StackMap for this block
   -> ByteOff            -- Sp on entry to the block
   -> ByteOff            -- SpHigh
   -> CmmNode C O        -- first node
   -> [CmmNode O O]      -- middle
   -> ByteOff            -- sp_off
   -> CmmNode O C        -- last node
   -> [CmmBlock]         -- new blocks
   -> [CmmBlock]         -- final blocks with Sp manifest

manifestSp dflags stackmaps stack0 sp0 sp_high
           first middle_pre sp_off last fixup_blocks
  = final_block : fixup_blocks'
  where
    area_off = getAreaOff stackmaps

    adj_pre_sp, adj_post_sp :: CmmNode e x -> CmmNode e x
    adj_pre_sp  = mapExpDeep (areaToSp dflags sp0            sp_high area_off)
    adj_post_sp = mapExpDeep (areaToSp dflags (sp0 - sp_off) sp_high area_off)

    final_middle = maybeAddSpAdj dflags sp0 sp_off
                 . blockFromList
                 . map adj_pre_sp
                 . elimStackStores stack0 stackmaps area_off
                 $ middle_pre
    final_last    = optStackCheck (adj_post_sp last)

    final_block   = blockJoin first final_middle final_last

    fixup_blocks' = map (mapBlock3' (id, adj_post_sp, id)) fixup_blocks

getAreaOff :: LabelMap StackMap -> (Area -> StackLoc)
getAreaOff _ Old = 0
getAreaOff stackmaps (Young l) =
  case mapLookup l stackmaps of
    Just sm -> sm_sp sm - sm_args sm
    Nothing -> pprPanic "getAreaOff" (ppr l)

maybeAddSpAdj
  :: DynFlags -> ByteOff -> ByteOff -> Block CmmNode O O -> Block CmmNode O O
maybeAddSpAdj dflags sp0 sp_off block =
  add_initial_unwind $ add_adj_unwind $ adj block
  where
    adj block
      | sp_off /= 0
      = block `blockSnoc` CmmAssign spReg (cmmOffset dflags spExpr sp_off)
      | otherwise = block
    -- Add unwind pseudo-instruction at the beginning of each block to
    -- document Sp level for debugging
    add_initial_unwind block
      | debugLevel dflags > 0
      = CmmUnwind [(Sp, Just sp_unwind)] `blockCons` block
      | otherwise
      = block
      where sp_unwind = CmmRegOff spReg (sp0 - wORD_SIZE dflags)

    -- Add unwind pseudo-instruction right after the Sp adjustment
    -- if there is one.
    add_adj_unwind block
      | debugLevel dflags > 0
      , sp_off /= 0
      = block `blockSnoc` CmmUnwind [(Sp, Just sp_unwind)]
      | otherwise
      = block
      where sp_unwind = CmmRegOff spReg (sp0 - wORD_SIZE dflags - sp_off)

{- Note [SP old/young offsets]

Sp(L) is the Sp offset on entry to block L relative to the base of the
OLD area.

SpArgs(L) is the size of the young area for L, i.e. the number of
arguments.

 - in block L, each reference to [old + N] turns into
   [Sp + Sp(L) - N]

 - in block L, each reference to [young(L') + N] turns into
   [Sp + Sp(L) - Sp(L') + SpArgs(L') - N]

 - be careful with the last node of each block: Sp has already been adjusted
   to be Sp + Sp(L) - Sp(L')
-}

areaToSp :: DynFlags -> ByteOff -> ByteOff -> (Area -> StackLoc) -> CmmExpr -> CmmExpr

areaToSp dflags sp_old _sp_hwm area_off (CmmStackSlot area n)
  = cmmOffset dflags spExpr (sp_old - area_off area - n)
    -- Replace (CmmStackSlot area n) with an offset from Sp

areaToSp dflags _ sp_hwm _ (CmmLit CmmHighStackMark)
  = mkIntExpr dflags sp_hwm
    -- Replace CmmHighStackMark with the number of bytes of stack used,
    -- the sp_hwm.   See Note [Stack usage] in StgCmmHeap

areaToSp dflags _ _ _ (CmmMachOp (MO_U_Lt _) args)
  | falseStackCheck args
  = zeroExpr dflags
areaToSp dflags _ _ _ (CmmMachOp (MO_U_Ge _) args)
  | falseStackCheck args
  = mkIntExpr dflags 1
    -- Replace a stack-overflow test that cannot fail with a no-op
    -- See Note [Always false stack check]

areaToSp _ _ _ _ other = other

-- | Determine whether a stack check cannot fail.
falseStackCheck :: [CmmExpr] -> Bool
falseStackCheck [ CmmMachOp (MO_Sub _)
                      [ CmmRegOff (CmmGlobal Sp) x_off
                      , CmmLit (CmmInt y_lit _)]
                , CmmReg (CmmGlobal SpLim)]
  = fromIntegral x_off >= y_lit
falseStackCheck _ = False

-- Note [Always false stack check]
-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-- We can optimise stack checks of the form
--
--   if ((Sp + x) - y < SpLim) then .. else ..
--
-- where are non-negative integer byte offsets.  Since we know that
-- SpLim <= Sp (remember the stack grows downwards), this test must
-- yield False if (x >= y), so we can rewrite the comparison to False.
-- A subsequent sinking pass will later drop the dead code.
-- Optimising this away depends on knowing that SpLim <= Sp, so it is
-- really the job of the stack layout algorithm, hence we do it now.
--
-- The control flow optimiser may negate a conditional to increase
-- the likelihood of a fallthrough if the branch is not taken.  But
-- not every conditional is inverted as the control flow optimiser
-- places some requirements on the predecessors of both branch targets.
-- So we better look for the inverted comparison too.

optStackCheck :: CmmNode O C -> CmmNode O C
optStackCheck n = -- Note [Always false stack check]
 case n of
   CmmCondBranch (CmmLit (CmmInt 0 _)) _true false _ -> CmmBranch false
   CmmCondBranch (CmmLit (CmmInt _ _)) true _false _ -> CmmBranch true
   other -> other

-- -----------------------------------------------------------------------------

-- | Eliminate stores of the form
--
--    Sp[area+n] = r
--
-- when we know that r is already in the same slot as Sp[area+n].  We
-- could do this in a later optimisation pass, but that would involve
-- a separate analysis and we already have the information to hand
-- here.  It helps clean up some extra stack stores in common cases.
--
-- Note that we may have to modify the StackMap as we walk through the
-- code using procMiddle, since an assignment to a variable in the
-- StackMap will invalidate its mapping there.
--
elimStackStores :: StackMap
                -> LabelMap StackMap
                -> (Area -> ByteOff)
                -> [CmmNode O O]
                -> [CmmNode O O]
elimStackStores stackmap stackmaps area_off nodes
  = go stackmap nodes
  where
    go _stackmap [] = []
    go stackmap (n:ns)
     = case n of
         CmmStore (CmmStackSlot area m) (CmmReg (CmmLocal r))
            | Just (_,off) <- lookupUFM (sm_regs stackmap) r
            , area_off area + m == off
            -> go stackmap ns
         _otherwise
            -> n : go (procMiddle stackmaps n stackmap) ns

-- -----------------------------------------------------------------------------
-- Update info tables to include stack liveness

setInfoTableStackMap :: DynFlags -> LabelMap StackMap -> CmmDecl -> CmmDecl
setInfoTableStackMap dflags stackmaps (CmmProc top_info@TopInfo{..} l v g)
  = CmmProc top_info{ info_tbls = mapMapWithKey fix_info info_tbls } l v g
  where
    fix_info lbl info_tbl@CmmInfoTable{ cit_rep = StackRep _ } =
       info_tbl { cit_rep = StackRep (get_liveness lbl) }
    fix_info _ other = other

    get_liveness :: BlockId -> Liveness
    get_liveness lbl
      = case mapLookup lbl stackmaps of
          Nothing -> pprPanic "setInfoTableStackMap" (ppr lbl <+> ppr info_tbls)
          Just sm -> stackMapToLiveness dflags sm

setInfoTableStackMap _ _ d = d

stackMapToLiveness :: DynFlags -> StackMap -> Liveness
stackMapToLiveness dflags StackMap{..} =
   reverse $ Array.elems $
        accumArray (\_ x -> x) True (toWords dflags sm_ret_off + 1,
                                     toWords dflags (sm_sp - sm_args)) live_words
   where
     live_words =  [ (toWords dflags off, False)
                   | (r,off) <- nonDetEltsUFM sm_regs
                   , isGcPtrType (localRegType r) ]
                   -- See Note [Unique Determinism and code generation]

-- -----------------------------------------------------------------------------
-- Pass 2
-- -----------------------------------------------------------------------------

insertReloadsAsNeeded
    :: DynFlags
    -> ProcPointSet
    -> LabelMap StackMap
    -> BlockId
    -> [CmmBlock]
    -> UniqSM [CmmBlock]
insertReloadsAsNeeded dflags procpoints final_stackmaps entry blocks = do
    toBlockList . fst <$>
        rewriteCmmBwd liveLattice rewriteCC (ofBlockList entry blocks) mapEmpty
  where
    rewriteCC :: RewriteFun CmmLocalLive
    rewriteCC (BlockCC e_node middle0 x_node) fact_base0 = do
        let entry_label = entryLabel e_node
            stackmap = case mapLookup entry_label final_stackmaps of
                Just sm -> sm
                Nothing -> panic "insertReloadsAsNeeded: rewriteCC: stackmap"

            -- Merge the liveness from successor blocks and analyse the last
            -- node.
            joined = gen_kill dflags x_node $!
                         joinOutFacts liveLattice x_node fact_base0
            -- What is live at the start of middle0.
            live_at_middle0 = foldNodesBwdOO (gen_kill dflags) middle0 joined

            -- If this is a procpoint we need to add the reloads, but only if
            -- they're actually live. Furthermore, nothing is live at the entry
            -- to a proc point.
            (middle1, live_with_reloads)
                | entry_label `setMember` procpoints
                = let reloads = insertReloads dflags stackmap live_at_middle0
                  in (foldr blockCons middle0 reloads, emptyRegSet)
                | otherwise
                = (middle0, live_at_middle0)

            -- Final liveness for this block.
            !fact_base2 = mapSingleton entry_label live_with_reloads

        return (BlockCC e_node middle1 x_node, fact_base2)

insertReloads :: DynFlags -> StackMap -> CmmLocalLive -> [CmmNode O O]
insertReloads dflags stackmap live =
     [ CmmAssign (CmmLocal reg)
                 -- This cmmOffset basically corresponds to manifesting
                 -- @CmmStackSlot Old sp_off@, see Note [SP old/young offsets]
                 (CmmLoad (cmmOffset dflags spExpr (sp_off - reg_off))
                          (localRegType reg))
     | (reg, reg_off) <- stackSlotRegs stackmap
     , reg `elemRegSet` live
     ]
   where
     sp_off = sm_sp stackmap

-- -----------------------------------------------------------------------------
-- Lowering safe foreign calls

{-
Note [Lower safe foreign calls]

We start with

   Sp[young(L1)] = L1
 ,-----------------------
 | r1 = foo(x,y,z) returns to L1
 '-----------------------
 L1:
   R1 = r1 -- copyIn, inserted by mkSafeCall
   ...

the stack layout algorithm will arrange to save and reload everything
live across the call.  Our job now is to expand the call so we get

   Sp[young(L1)] = L1
 ,-----------------------
 | SAVE_THREAD_STATE()
 | token = suspendThread(BaseReg, interruptible)
 | r = foo(x,y,z)
 | BaseReg = resumeThread(token)
 | LOAD_THREAD_STATE()
 | R1 = r  -- copyOut
 | jump Sp[0]
 '-----------------------
 L1:
   r = R1 -- copyIn, inserted by mkSafeCall
   ...

Note the copyOut, which saves the results in the places that L1 is
expecting them (see Note [safe foreign call convention]). Note also
that safe foreign call is replace by an unsafe one in the Cmm graph.
-}

lowerSafeForeignCall :: DynFlags -> CmmBlock -> UniqSM CmmBlock
lowerSafeForeignCall dflags block
  | (entry@(CmmEntry _ tscp), middle, CmmForeignCall { .. }) <- blockSplit block
  = do
    -- Both 'id' and 'new_base' are KindNonPtr because they're
    -- RTS-only objects and are not subject to garbage collection
    id <- newTemp (bWord dflags)
    new_base <- newTemp (cmmRegType dflags baseReg)
    let (caller_save, caller_load) = callerSaveVolatileRegs dflags
    save_state_code <- saveThreadState dflags
    load_state_code <- loadThreadState dflags
    let suspend = save_state_code  <*>
                  caller_save <*>
                  mkMiddle (callSuspendThread dflags id intrbl)
        midCall = mkUnsafeCall tgt res args
        resume  = mkMiddle (callResumeThread new_base id) <*>
                  -- Assign the result to BaseReg: we
                  -- might now have a different Capability!
                  mkAssign baseReg (CmmReg (CmmLocal new_base)) <*>
                  caller_load <*>
                  load_state_code

        (_, regs, copyout) =
             copyOutOflow dflags NativeReturn Jump (Young succ)
                            (map (CmmReg . CmmLocal) res)
                            ret_off []

        -- NB. after resumeThread returns, the top-of-stack probably contains
        -- the stack frame for succ, but it might not: if the current thread
        -- received an exception during the call, then the stack might be
        -- different.  Hence we continue by jumping to the top stack frame,
        -- not by jumping to succ.
        jump = CmmCall { cml_target    = entryCode dflags $
                                         CmmLoad spExpr (bWord dflags)
                       , cml_cont      = Just succ
                       , cml_args_regs = regs
                       , cml_args      = widthInBytes (wordWidth dflags)
                       , cml_ret_args  = ret_args
                       , cml_ret_off   = ret_off }

    graph' <- lgraphOfAGraph ( suspend <*>
                               midCall <*>
                               resume  <*>
                               copyout <*>
                               mkLast jump, tscp)

    case toBlockList graph' of
      [one] -> let (_, middle', last) = blockSplit one
               in return (blockJoin entry (middle `blockAppend` middle') last)
      _ -> panic "lowerSafeForeignCall0"

  -- Block doesn't end in a safe foreign call:
  | otherwise = return block

foreignLbl :: FastString -> CmmExpr
foreignLbl name = CmmLit (CmmLabel (mkForeignLabel name Nothing ForeignLabelInExternalPackage IsFunction))

callSuspendThread :: DynFlags -> LocalReg -> Bool -> CmmNode O O
callSuspendThread dflags id intrbl =
  CmmUnsafeForeignCall
       (ForeignTarget (foreignLbl (fsLit "suspendThread"))
        (ForeignConvention CCallConv [AddrHint, NoHint] [AddrHint] CmmMayReturn))
       [id] [baseExpr, mkIntExpr dflags (fromEnum intrbl)]

callResumeThread :: LocalReg -> LocalReg -> CmmNode O O
callResumeThread new_base id =
  CmmUnsafeForeignCall
       (ForeignTarget (foreignLbl (fsLit "resumeThread"))
            (ForeignConvention CCallConv [AddrHint] [AddrHint] CmmMayReturn))
       [new_base] [CmmReg (CmmLocal id)]

-- -----------------------------------------------------------------------------

plusW :: DynFlags -> ByteOff -> WordOff -> ByteOff
plusW dflags b w = b + w * wORD_SIZE dflags

data StackSlot = Occupied | Empty
     -- Occupied: a return address or part of an update frame

instance Outputable StackSlot where
  ppr Occupied = text "XXX"
  ppr Empty    = text "---"

dropEmpty :: WordOff -> [StackSlot] -> Maybe [StackSlot]
dropEmpty 0 ss           = Just ss
dropEmpty n (Empty : ss) = dropEmpty (n-1) ss
dropEmpty _ _            = Nothing

isEmpty :: StackSlot -> Bool
isEmpty Empty = True
isEmpty _ = False

localRegBytes :: DynFlags -> LocalReg -> ByteOff
localRegBytes dflags r
    = roundUpToWords dflags (widthInBytes (typeWidth (localRegType r)))

localRegWords :: DynFlags -> LocalReg -> WordOff
localRegWords dflags = toWords dflags . localRegBytes dflags

toWords :: DynFlags -> ByteOff -> WordOff
toWords dflags x = x `quot` wORD_SIZE dflags

stackSlotRegs :: StackMap -> [(LocalReg, StackLoc)]
stackSlotRegs sm = nonDetEltsUFM (sm_regs sm)
  -- See Note [Unique Determinism and code generation]


-----------------------------------------------------------------------------
--
-- (c) The University of Glasgow 2011
--
-- CmmLint: checking the correctness of Cmm statements and expressions
--
-----------------------------------------------------------------------------
{-# LANGUAGE GADTs #-}
module CmmLint (
    cmmLint, cmmLintGraph
  ) where

import GhcPrelude

import Hoopl.Block
import Hoopl.Collections
import Hoopl.Graph
import Hoopl.Label
import Cmm
import CmmUtils
import CmmLive
import CmmSwitch (switchTargetsToList)
import PprCmm ()
import Outputable
import DynFlags

import Control.Monad (liftM, ap)

-- Things to check:
--     - invariant on CmmBlock in CmmExpr (see comment there)
--     - check for branches to blocks that don't exist
--     - check types

-- -----------------------------------------------------------------------------
-- Exported entry points:

cmmLint :: (Outputable d, Outputable h)
        => DynFlags -> GenCmmGroup d h CmmGraph -> Maybe SDoc
cmmLint dflags tops = runCmmLint dflags (mapM_ (lintCmmDecl dflags)) tops

cmmLintGraph :: DynFlags -> CmmGraph -> Maybe SDoc
cmmLintGraph dflags g = runCmmLint dflags (lintCmmGraph dflags) g

runCmmLint :: Outputable a => DynFlags -> (a -> CmmLint b) -> a -> Maybe SDoc
runCmmLint dflags l p =
   case unCL (l p) dflags of
     Left err -> Just (vcat [text "Cmm lint error:",
                             nest 2 err,
                             text "Program was:",
                             nest 2 (ppr p)])
     Right _  -> Nothing

lintCmmDecl :: DynFlags -> GenCmmDecl h i CmmGraph -> CmmLint ()
lintCmmDecl dflags (CmmProc _ lbl _ g)
  = addLintInfo (text "in proc " <> ppr lbl) $ lintCmmGraph dflags g
lintCmmDecl _ (CmmData {})
  = return ()

lintCmmGraph :: DynFlags -> CmmGraph -> CmmLint ()
lintCmmGraph dflags g =
    cmmLocalLiveness dflags g `seq` mapM_ (lintCmmBlock labels) blocks
    -- cmmLiveness throws an error if there are registers
    -- live on entry to the graph (i.e. undefined
    -- variables)
  where
       blocks = toBlockList g
       labels = setFromList (map entryLabel blocks)

lintCmmBlock :: LabelSet -> CmmBlock -> CmmLint ()
lintCmmBlock labels block
  = addLintInfo (text "in basic block " <> ppr (entryLabel block)) $ do
        let (_, middle, last) = blockSplit block
        mapM_ lintCmmMiddle (blockToList middle)
        lintCmmLast labels last

-- -----------------------------------------------------------------------------
-- lintCmmExpr

-- Checks whether a CmmExpr is "type-correct", and check for obvious-looking
-- byte/word mismatches.

lintCmmExpr :: CmmExpr -> CmmLint CmmType
lintCmmExpr (CmmLoad expr rep) = do
  _ <- lintCmmExpr expr
  -- Disabled, if we have the inlining phase before the lint phase,
  -- we can have funny offsets due to pointer tagging. -- EZY
  -- when (widthInBytes (typeWidth rep) >= wORD_SIZE) $
  --   cmmCheckWordAddress expr
  return rep
lintCmmExpr expr@(CmmMachOp op args) = do
  dflags <- getDynFlags
  tys <- mapM lintCmmExpr args
  if map (typeWidth . cmmExprType dflags) args == machOpArgReps dflags op
        then cmmCheckMachOp op args tys
        else cmmLintMachOpErr expr (map (cmmExprType dflags) args) (machOpArgReps dflags op)
lintCmmExpr (CmmRegOff reg offset)
  = do dflags <- getDynFlags
       let rep = typeWidth (cmmRegType dflags reg)
       lintCmmExpr (CmmMachOp (MO_Add rep)
                [CmmReg reg, CmmLit (CmmInt (fromIntegral offset) rep)])
lintCmmExpr expr =
  do dflags <- getDynFlags
     return (cmmExprType dflags expr)

-- Check for some common byte/word mismatches (eg. Sp + 1)
cmmCheckMachOp   :: MachOp -> [CmmExpr] -> [CmmType] -> CmmLint CmmType
cmmCheckMachOp op [lit@(CmmLit (CmmInt { })), reg@(CmmReg _)] tys
  = cmmCheckMachOp op [reg, lit] tys
cmmCheckMachOp op _ tys
  = do dflags <- getDynFlags
       return (machOpResultType dflags op tys)

{-
isOffsetOp :: MachOp -> Bool
isOffsetOp (MO_Add _) = True
isOffsetOp (MO_Sub _) = True
isOffsetOp _ = False

-- This expression should be an address from which a word can be loaded:
-- check for funny-looking sub-word offsets.
_cmmCheckWordAddress :: CmmExpr -> CmmLint ()
_cmmCheckWordAddress e@(CmmMachOp op [arg, CmmLit (CmmInt i _)])
  | isOffsetOp op && notNodeReg arg && i `rem` fromIntegral (wORD_SIZE dflags) /= 0
  = cmmLintDubiousWordOffset e
_cmmCheckWordAddress e@(CmmMachOp op [CmmLit (CmmInt i _), arg])
  | isOffsetOp op && notNodeReg arg && i `rem` fromIntegral (wORD_SIZE dflags) /= 0
  = cmmLintDubiousWordOffset e
_cmmCheckWordAddress _
  = return ()

-- No warnings for unaligned arithmetic with the node register,
-- which is used to extract fields from tagged constructor closures.
notNodeReg :: CmmExpr -> Bool
notNodeReg (CmmReg reg) | reg == nodeReg = False
notNodeReg _                             = True
-}

lintCmmMiddle :: CmmNode O O -> CmmLint ()
lintCmmMiddle node = case node of
  CmmComment _ -> return ()
  CmmTick _    -> return ()
  CmmUnwind{}  -> return ()

  CmmAssign reg expr -> do
            dflags <- getDynFlags
            erep <- lintCmmExpr expr
            let reg_ty = cmmRegType dflags reg
            if (erep `cmmEqType_ignoring_ptrhood` reg_ty)
                then return ()
                else cmmLintAssignErr (CmmAssign reg expr) erep reg_ty

  CmmStore l r -> do
            _ <- lintCmmExpr l
            _ <- lintCmmExpr r
            return ()

  CmmUnsafeForeignCall target _formals actuals -> do
            lintTarget target
            mapM_ lintCmmExpr actuals

lintCmmLast :: LabelSet -> CmmNode O C -> CmmLint ()
lintCmmLast labels node = case node of
  CmmBranch id -> checkTarget id

  CmmCondBranch e t f _ -> do
            dflags <- getDynFlags
            mapM_ checkTarget [t,f]
            _ <- lintCmmExpr e
            checkCond dflags e

  CmmSwitch e ids -> do
            dflags <- getDynFlags
            mapM_ checkTarget $ switchTargetsToList ids
            erep <- lintCmmExpr e
            if (erep `cmmEqType_ignoring_ptrhood` bWord dflags)
              then return ()
              else cmmLintErr (text "switch scrutinee is not a word: " <>
                               ppr e <> text " :: " <> ppr erep)

  CmmCall { cml_target = target, cml_cont = cont } -> do
          _ <- lintCmmExpr target
          maybe (return ()) checkTarget cont

  CmmForeignCall tgt _ args succ _ _ _ -> do
          lintTarget tgt
          mapM_ lintCmmExpr args
          checkTarget succ
 where
  checkTarget id
     | setMember id labels = return ()
     | otherwise = cmmLintErr (text "Branch to nonexistent id" <+> ppr id)

lintTarget :: ForeignTarget -> CmmLint ()
lintTarget (ForeignTarget e _) = lintCmmExpr e >> return ()
lintTarget (PrimTarget {})     = return ()

checkCond :: DynFlags -> CmmExpr -> CmmLint ()
checkCond _ (CmmMachOp mop _) | isComparisonMachOp mop = return ()
checkCond dflags (CmmLit (CmmInt x t)) | x == 0 || x == 1, t == wordWidth dflags = return () -- constant values
checkCond _ expr
    = cmmLintErr (hang (text "expression is not a conditional:") 2
                         (ppr expr))

-- -----------------------------------------------------------------------------
-- CmmLint monad

-- just a basic error monad:

newtype CmmLint a = CmmLint { unCL :: DynFlags -> Either SDoc a }

instance Functor CmmLint where
      fmap = liftM

instance Applicative CmmLint where
      pure a = CmmLint (\_ -> Right a)
      (<*>) = ap

instance Monad CmmLint where
  CmmLint m >>= k = CmmLint $ \dflags ->
                                case m dflags of
                                Left e -> Left e
                                Right a -> unCL (k a) dflags

instance HasDynFlags CmmLint where
    getDynFlags = CmmLint (\dflags -> Right dflags)

cmmLintErr :: SDoc -> CmmLint a
cmmLintErr msg = CmmLint (\_ -> Left msg)

addLintInfo :: SDoc -> CmmLint a -> CmmLint a
addLintInfo info thing = CmmLint $ \dflags ->
   case unCL thing dflags of
        Left err -> Left (hang info 2 err)
        Right a  -> Right a

cmmLintMachOpErr :: CmmExpr -> [CmmType] -> [Width] -> CmmLint a
cmmLintMachOpErr expr argsRep opExpectsRep
     = cmmLintErr (text "in MachOp application: " $$
                   nest 2 (ppr  expr) $$
                      (text "op is expecting: " <+> ppr opExpectsRep) $$
                      (text "arguments provide: " <+> ppr argsRep))

cmmLintAssignErr :: CmmNode e x -> CmmType -> CmmType -> CmmLint a
cmmLintAssignErr stmt e_ty r_ty
  = cmmLintErr (text "in assignment: " $$
                nest 2 (vcat [ppr stmt,
                              text "Reg ty:" <+> ppr r_ty,
                              text "Rhs ty:" <+> ppr e_ty]))

{-
cmmLintDubiousWordOffset :: CmmExpr -> CmmLint a
cmmLintDubiousWordOffset expr
   = cmmLintErr (text "offset is not a multiple of words: " $$
                 nest 2 (ppr expr))
-}


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}

module CmmLive
    ( CmmLocalLive
    , cmmLocalLiveness
    , cmmGlobalLiveness
    , liveLattice
    , gen_kill
    )
where

import GhcPrelude

import DynFlags
import BlockId
import Cmm
import PprCmmExpr ()
import Hoopl.Block
import Hoopl.Collections
import Hoopl.Dataflow
import Hoopl.Label

import Maybes
import Outputable

-----------------------------------------------------------------------------
-- Calculating what variables are live on entry to a basic block
-----------------------------------------------------------------------------

-- | The variables live on entry to a block
type CmmLive r = RegSet r
type CmmLocalLive = CmmLive LocalReg

-- | The dataflow lattice
liveLattice :: Ord r => DataflowLattice (CmmLive r)
{-# SPECIALIZE liveLattice :: DataflowLattice (CmmLive LocalReg) #-}
{-# SPECIALIZE liveLattice :: DataflowLattice (CmmLive GlobalReg) #-}
liveLattice = DataflowLattice emptyRegSet add
  where
    add (OldFact old) (NewFact new) =
        let !join = plusRegSet old new
        in changedIf (sizeRegSet join > sizeRegSet old) join

-- | A mapping from block labels to the variables live on entry
type BlockEntryLiveness r = LabelMap (CmmLive r)

-----------------------------------------------------------------------------
-- | Calculated liveness info for a CmmGraph
-----------------------------------------------------------------------------

cmmLocalLiveness :: DynFlags -> CmmGraph -> BlockEntryLiveness LocalReg
cmmLocalLiveness dflags graph =
    check $ analyzeCmmBwd liveLattice (xferLive dflags) graph mapEmpty
  where
    entry = g_entry graph
    check facts =
        noLiveOnEntry entry (expectJust "check" $ mapLookup entry facts) facts

cmmGlobalLiveness :: DynFlags -> CmmGraph -> BlockEntryLiveness GlobalReg
cmmGlobalLiveness dflags graph =
    analyzeCmmBwd liveLattice (xferLive dflags) graph mapEmpty

-- | On entry to the procedure, there had better not be any LocalReg's live-in.
noLiveOnEntry :: BlockId -> CmmLive LocalReg -> a -> a
noLiveOnEntry bid in_fact x =
  if nullRegSet in_fact then x
  else pprPanic "LocalReg's live-in to graph" (ppr bid <+> ppr in_fact)

gen_kill
    :: (DefinerOfRegs r n, UserOfRegs r n)
    => DynFlags -> n -> CmmLive r -> CmmLive r
gen_kill dflags node set =
    let !afterKill = foldRegsDefd dflags deleteFromRegSet set node
    in foldRegsUsed dflags extendRegSet afterKill node
{-# INLINE gen_kill #-}

xferLive
    :: forall r.
       ( UserOfRegs r (CmmNode O O)
       , DefinerOfRegs r (CmmNode O O)
       , UserOfRegs r (CmmNode O C)
       , DefinerOfRegs r (CmmNode O C)
       )
    => DynFlags -> TransferFun (CmmLive r)
xferLive dflags (BlockCC eNode middle xNode) fBase =
    let joined = gen_kill dflags xNode $! joinOutFacts liveLattice xNode fBase
        !result = foldNodesBwdOO (gen_kill dflags) middle joined
    in mapSingleton (entryLabel eNode) result
{-# SPECIALIZE xferLive :: DynFlags -> TransferFun (CmmLive LocalReg) #-}
{-# SPECIALIZE xferLive :: DynFlags -> TransferFun (CmmLive GlobalReg) #-}


module CmmMachOp
    ( MachOp(..)
    , pprMachOp, isCommutableMachOp, isAssociativeMachOp
    , isComparisonMachOp, maybeIntComparison, machOpResultType
    , machOpArgReps, maybeInvertComparison, isFloatComparison

    -- MachOp builders
    , mo_wordAdd, mo_wordSub, mo_wordEq, mo_wordNe,mo_wordMul, mo_wordSQuot
    , mo_wordSRem, mo_wordSNeg, mo_wordUQuot, mo_wordURem
    , mo_wordSGe, mo_wordSLe, mo_wordSGt, mo_wordSLt, mo_wordUGe
    , mo_wordULe, mo_wordUGt, mo_wordULt
    , mo_wordAnd, mo_wordOr, mo_wordXor, mo_wordNot
    , mo_wordShl, mo_wordSShr, mo_wordUShr
    , mo_u_8To32, mo_s_8To32, mo_u_16To32, mo_s_16To32
    , mo_u_8ToWord, mo_s_8ToWord, mo_u_16ToWord, mo_s_16ToWord
    , mo_u_32ToWord, mo_s_32ToWord
    , mo_32To8, mo_32To16, mo_WordTo8, mo_WordTo16, mo_WordTo32, mo_WordTo64

    -- CallishMachOp
    , CallishMachOp(..), callishMachOpHints
    , pprCallishMachOp
    , machOpMemcpyishAlign

    -- Atomic read-modify-write
    , AtomicMachOp(..)
   )
where

import GhcPrelude

import CmmType
import Outputable
import DynFlags

-----------------------------------------------------------------------------
--              MachOp
-----------------------------------------------------------------------------

{- |
Machine-level primops; ones which we can reasonably delegate to the
native code generators to handle.

Most operations are parameterised by the 'Width' that they operate on.
Some operations have separate signed and unsigned versions, and float
and integer versions.
-}

data MachOp
  -- Integer operations (insensitive to signed/unsigned)
  = MO_Add Width
  | MO_Sub Width
  | MO_Eq  Width
  | MO_Ne  Width
  | MO_Mul Width                -- low word of multiply

  -- Signed multiply/divide
  | MO_S_MulMayOflo Width       -- nonzero if signed multiply overflows
  | MO_S_Quot Width             -- signed / (same semantics as IntQuotOp)
  | MO_S_Rem  Width             -- signed % (same semantics as IntRemOp)
  | MO_S_Neg  Width             -- unary -

  -- Unsigned multiply/divide
  | MO_U_MulMayOflo Width       -- nonzero if unsigned multiply overflows
  | MO_U_Quot Width             -- unsigned / (same semantics as WordQuotOp)
  | MO_U_Rem  Width             -- unsigned % (same semantics as WordRemOp)

  -- Signed comparisons
  | MO_S_Ge Width
  | MO_S_Le Width
  | MO_S_Gt Width
  | MO_S_Lt Width

  -- Unsigned comparisons
  | MO_U_Ge Width
  | MO_U_Le Width
  | MO_U_Gt Width
  | MO_U_Lt Width

  -- Floating point arithmetic
  | MO_F_Add  Width
  | MO_F_Sub  Width
  | MO_F_Neg  Width             -- unary -
  | MO_F_Mul  Width
  | MO_F_Quot Width

  -- Floating point comparison
  | MO_F_Eq Width
  | MO_F_Ne Width
  | MO_F_Ge Width
  | MO_F_Le Width
  | MO_F_Gt Width
  | MO_F_Lt Width

  -- Bitwise operations.  Not all of these may be supported
  -- at all sizes, and only integral Widths are valid.
  | MO_And   Width
  | MO_Or    Width
  | MO_Xor   Width
  | MO_Not   Width
  | MO_Shl   Width
  | MO_U_Shr Width      -- unsigned shift right
  | MO_S_Shr Width      -- signed shift right

  -- Conversions.  Some of these will be NOPs.
  -- Floating-point conversions use the signed variant.
  | MO_SF_Conv Width Width      -- Signed int -> Float
  | MO_FS_Conv Width Width      -- Float -> Signed int
  | MO_SS_Conv Width Width      -- Signed int -> Signed int
  | MO_UU_Conv Width Width      -- unsigned int -> unsigned int
  | MO_XX_Conv Width Width      -- int -> int; puts no requirements on the
                                -- contents of upper bits when extending;
                                -- narrowing is simply truncation; the only
                                -- expectation is that we can recover the
                                -- original value by applying the opposite
                                -- MO_XX_Conv, e.g.,
                                --   MO_XX_CONV W64 W8 (MO_XX_CONV W8 W64 x)
                                -- is equivalent to just x.
  | MO_FF_Conv Width Width      -- Float -> Float

  -- Vector element insertion and extraction operations
  | MO_V_Insert  Length Width   -- Insert scalar into vector
  | MO_V_Extract Length Width   -- Extract scalar from vector

  -- Integer vector operations
  | MO_V_Add Length Width
  | MO_V_Sub Length Width
  | MO_V_Mul Length Width

  -- Signed vector multiply/divide
  | MO_VS_Quot Length Width
  | MO_VS_Rem  Length Width
  | MO_VS_Neg  Length Width

  -- Unsigned vector multiply/divide
  | MO_VU_Quot Length Width
  | MO_VU_Rem  Length Width

  -- Floting point vector element insertion and extraction operations
  | MO_VF_Insert  Length Width   -- Insert scalar into vector
  | MO_VF_Extract Length Width   -- Extract scalar from vector

  -- Floating point vector operations
  | MO_VF_Add  Length Width
  | MO_VF_Sub  Length Width
  | MO_VF_Neg  Length Width      -- unary negation
  | MO_VF_Mul  Length Width
  | MO_VF_Quot Length Width

  -- Alignment check (for -falignment-sanitisation)
  | MO_AlignmentCheck Int Width
  deriving (Eq, Show)

pprMachOp :: MachOp -> SDoc
pprMachOp mo = text (show mo)

-- -----------------------------------------------------------------------------
-- Some common MachReps

-- A 'wordRep' is a machine word on the target architecture
-- Specifically, it is the size of an Int#, Word#, Addr#
-- and the unit of allocation on the stack and the heap
-- Any pointer is also guaranteed to be a wordRep.

mo_wordAdd, mo_wordSub, mo_wordEq, mo_wordNe,mo_wordMul, mo_wordSQuot
    , mo_wordSRem, mo_wordSNeg, mo_wordUQuot, mo_wordURem
    , mo_wordSGe, mo_wordSLe, mo_wordSGt, mo_wordSLt, mo_wordUGe
    , mo_wordULe, mo_wordUGt, mo_wordULt
    , mo_wordAnd, mo_wordOr, mo_wordXor, mo_wordNot, mo_wordShl, mo_wordSShr, mo_wordUShr
    , mo_u_8ToWord, mo_s_8ToWord, mo_u_16ToWord, mo_s_16ToWord, mo_u_32ToWord, mo_s_32ToWord
    , mo_WordTo8, mo_WordTo16, mo_WordTo32, mo_WordTo64
    :: DynFlags -> MachOp

mo_u_8To32, mo_s_8To32, mo_u_16To32, mo_s_16To32
    , mo_32To8, mo_32To16
    :: MachOp

mo_wordAdd      dflags = MO_Add (wordWidth dflags)
mo_wordSub      dflags = MO_Sub (wordWidth dflags)
mo_wordEq       dflags = MO_Eq  (wordWidth dflags)
mo_wordNe       dflags = MO_Ne  (wordWidth dflags)
mo_wordMul      dflags = MO_Mul (wordWidth dflags)
mo_wordSQuot    dflags = MO_S_Quot (wordWidth dflags)
mo_wordSRem     dflags = MO_S_Rem (wordWidth dflags)
mo_wordSNeg     dflags = MO_S_Neg (wordWidth dflags)
mo_wordUQuot    dflags = MO_U_Quot (wordWidth dflags)
mo_wordURem     dflags = MO_U_Rem (wordWidth dflags)

mo_wordSGe      dflags = MO_S_Ge  (wordWidth dflags)
mo_wordSLe      dflags = MO_S_Le  (wordWidth dflags)
mo_wordSGt      dflags = MO_S_Gt  (wordWidth dflags)
mo_wordSLt      dflags = MO_S_Lt  (wordWidth dflags)

mo_wordUGe      dflags = MO_U_Ge  (wordWidth dflags)
mo_wordULe      dflags = MO_U_Le  (wordWidth dflags)
mo_wordUGt      dflags = MO_U_Gt  (wordWidth dflags)
mo_wordULt      dflags = MO_U_Lt  (wordWidth dflags)

mo_wordAnd      dflags = MO_And (wordWidth dflags)
mo_wordOr       dflags = MO_Or  (wordWidth dflags)
mo_wordXor      dflags = MO_Xor (wordWidth dflags)
mo_wordNot      dflags = MO_Not (wordWidth dflags)
mo_wordShl      dflags = MO_Shl (wordWidth dflags)
mo_wordSShr     dflags = MO_S_Shr (wordWidth dflags)
mo_wordUShr     dflags = MO_U_Shr (wordWidth dflags)

mo_u_8To32             = MO_UU_Conv W8 W32
mo_s_8To32             = MO_SS_Conv W8 W32
mo_u_16To32            = MO_UU_Conv W16 W32
mo_s_16To32            = MO_SS_Conv W16 W32

mo_u_8ToWord    dflags = MO_UU_Conv W8  (wordWidth dflags)
mo_s_8ToWord    dflags = MO_SS_Conv W8  (wordWidth dflags)
mo_u_16ToWord   dflags = MO_UU_Conv W16 (wordWidth dflags)
mo_s_16ToWord   dflags = MO_SS_Conv W16 (wordWidth dflags)
mo_s_32ToWord   dflags = MO_SS_Conv W32 (wordWidth dflags)
mo_u_32ToWord   dflags = MO_UU_Conv W32 (wordWidth dflags)

mo_WordTo8      dflags = MO_UU_Conv (wordWidth dflags) W8
mo_WordTo16     dflags = MO_UU_Conv (wordWidth dflags) W16
mo_WordTo32     dflags = MO_UU_Conv (wordWidth dflags) W32
mo_WordTo64     dflags = MO_UU_Conv (wordWidth dflags) W64

mo_32To8               = MO_UU_Conv W32 W8
mo_32To16              = MO_UU_Conv W32 W16

-- ----------------------------------------------------------------------------
-- isCommutableMachOp

{- |
Returns 'True' if the MachOp has commutable arguments.  This is used
in the platform-independent Cmm optimisations.

If in doubt, return 'False'.  This generates worse code on the
native routes, but is otherwise harmless.
-}
isCommutableMachOp :: MachOp -> Bool
isCommutableMachOp mop =
  case mop of
        MO_Add _                -> True
        MO_Eq _                 -> True
        MO_Ne _                 -> True
        MO_Mul _                -> True
        MO_S_MulMayOflo _       -> True
        MO_U_MulMayOflo _       -> True
        MO_And _                -> True
        MO_Or _                 -> True
        MO_Xor _                -> True
        MO_F_Add _              -> True
        MO_F_Mul _              -> True
        _other                  -> False

-- ----------------------------------------------------------------------------
-- isAssociativeMachOp

{- |
Returns 'True' if the MachOp is associative (i.e. @(x+y)+z == x+(y+z)@)
This is used in the platform-independent Cmm optimisations.

If in doubt, return 'False'.  This generates worse code on the
native routes, but is otherwise harmless.
-}
isAssociativeMachOp :: MachOp -> Bool
isAssociativeMachOp mop =
  case mop of
        MO_Add {} -> True       -- NB: does not include
        MO_Mul {} -> True --     floatint point!
        MO_And {} -> True
        MO_Or  {} -> True
        MO_Xor {} -> True
        _other    -> False

-- ----------------------------------------------------------------------------
-- isComparisonMachOp

{- |
Returns 'True' if the MachOp is a comparison.

If in doubt, return False.  This generates worse code on the
native routes, but is otherwise harmless.
-}
isComparisonMachOp :: MachOp -> Bool
isComparisonMachOp mop =
  case mop of
    MO_Eq   _  -> True
    MO_Ne   _  -> True
    MO_S_Ge _  -> True
    MO_S_Le _  -> True
    MO_S_Gt _  -> True
    MO_S_Lt _  -> True
    MO_U_Ge _  -> True
    MO_U_Le _  -> True
    MO_U_Gt _  -> True
    MO_U_Lt _  -> True
    MO_F_Eq {} -> True
    MO_F_Ne {} -> True
    MO_F_Ge {} -> True
    MO_F_Le {} -> True
    MO_F_Gt {} -> True
    MO_F_Lt {} -> True
    _other     -> False

{- |
Returns @Just w@ if the operation is an integer comparison with width
@w@, or @Nothing@ otherwise.
-}
maybeIntComparison :: MachOp -> Maybe Width
maybeIntComparison mop =
  case mop of
    MO_Eq   w  -> Just w
    MO_Ne   w  -> Just w
    MO_S_Ge w  -> Just w
    MO_S_Le w  -> Just w
    MO_S_Gt w  -> Just w
    MO_S_Lt w  -> Just w
    MO_U_Ge w  -> Just w
    MO_U_Le w  -> Just w
    MO_U_Gt w  -> Just w
    MO_U_Lt w  -> Just w
    _ -> Nothing

isFloatComparison :: MachOp -> Bool
isFloatComparison mop =
  case mop of
    MO_F_Eq {} -> True
    MO_F_Ne {} -> True
    MO_F_Ge {} -> True
    MO_F_Le {} -> True
    MO_F_Gt {} -> True
    MO_F_Lt {} -> True
    _other     -> False

-- -----------------------------------------------------------------------------
-- Inverting conditions

-- Sometimes it's useful to be able to invert the sense of a
-- condition.  Not all conditional tests are invertible: in
-- particular, floating point conditionals cannot be inverted, because
-- there exist floating-point values which return False for both senses
-- of a condition (eg. !(NaN > NaN) && !(NaN /<= NaN)).

maybeInvertComparison :: MachOp -> Maybe MachOp
maybeInvertComparison op
  = case op of  -- None of these Just cases include floating point
        MO_Eq r   -> Just (MO_Ne r)
        MO_Ne r   -> Just (MO_Eq r)
        MO_U_Lt r -> Just (MO_U_Ge r)
        MO_U_Gt r -> Just (MO_U_Le r)
        MO_U_Le r -> Just (MO_U_Gt r)
        MO_U_Ge r -> Just (MO_U_Lt r)
        MO_S_Lt r -> Just (MO_S_Ge r)
        MO_S_Gt r -> Just (MO_S_Le r)
        MO_S_Le r -> Just (MO_S_Gt r)
        MO_S_Ge r -> Just (MO_S_Lt r)
        _other    -> Nothing

-- ----------------------------------------------------------------------------
-- machOpResultType

{- |
Returns the MachRep of the result of a MachOp.
-}
machOpResultType :: DynFlags -> MachOp -> [CmmType] -> CmmType
machOpResultType dflags mop tys =
  case mop of
    MO_Add {}           -> ty1  -- Preserve GC-ptr-hood
    MO_Sub {}           -> ty1  -- of first arg
    MO_Mul    r         -> cmmBits r
    MO_S_MulMayOflo r   -> cmmBits r
    MO_S_Quot r         -> cmmBits r
    MO_S_Rem  r         -> cmmBits r
    MO_S_Neg  r         -> cmmBits r
    MO_U_MulMayOflo r   -> cmmBits r
    MO_U_Quot r         -> cmmBits r
    MO_U_Rem  r         -> cmmBits r

    MO_Eq {}            -> comparisonResultRep dflags
    MO_Ne {}            -> comparisonResultRep dflags
    MO_S_Ge {}          -> comparisonResultRep dflags
    MO_S_Le {}          -> comparisonResultRep dflags
    MO_S_Gt {}          -> comparisonResultRep dflags
    MO_S_Lt {}          -> comparisonResultRep dflags

    MO_U_Ge {}          -> comparisonResultRep dflags
    MO_U_Le {}          -> comparisonResultRep dflags
    MO_U_Gt {}          -> comparisonResultRep dflags
    MO_U_Lt {}          -> comparisonResultRep dflags

    MO_F_Add r          -> cmmFloat r
    MO_F_Sub r          -> cmmFloat r
    MO_F_Mul r          -> cmmFloat r
    MO_F_Quot r         -> cmmFloat r
    MO_F_Neg r          -> cmmFloat r
    MO_F_Eq  {}         -> comparisonResultRep dflags
    MO_F_Ne  {}         -> comparisonResultRep dflags
    MO_F_Ge  {}         -> comparisonResultRep dflags
    MO_F_Le  {}         -> comparisonResultRep dflags
    MO_F_Gt  {}         -> comparisonResultRep dflags
    MO_F_Lt  {}         -> comparisonResultRep dflags

    MO_And {}           -> ty1  -- Used for pointer masking
    MO_Or {}            -> ty1
    MO_Xor {}           -> ty1
    MO_Not   r          -> cmmBits r
    MO_Shl   r          -> cmmBits r
    MO_U_Shr r          -> cmmBits r
    MO_S_Shr r          -> cmmBits r

    MO_SS_Conv _ to     -> cmmBits to
    MO_UU_Conv _ to     -> cmmBits to
    MO_XX_Conv _ to     -> cmmBits to
    MO_FS_Conv _ to     -> cmmBits to
    MO_SF_Conv _ to     -> cmmFloat to
    MO_FF_Conv _ to     -> cmmFloat to

    MO_V_Insert  l w    -> cmmVec l (cmmBits w)
    MO_V_Extract _ w    -> cmmBits w

    MO_V_Add l w        -> cmmVec l (cmmBits w)
    MO_V_Sub l w        -> cmmVec l (cmmBits w)
    MO_V_Mul l w        -> cmmVec l (cmmBits w)

    MO_VS_Quot l w      -> cmmVec l (cmmBits w)
    MO_VS_Rem  l w      -> cmmVec l (cmmBits w)
    MO_VS_Neg  l w      -> cmmVec l (cmmBits w)

    MO_VU_Quot l w      -> cmmVec l (cmmBits w)
    MO_VU_Rem  l w      -> cmmVec l (cmmBits w)

    MO_VF_Insert  l w   -> cmmVec l (cmmFloat w)
    MO_VF_Extract _ w   -> cmmFloat w

    MO_VF_Add  l w      -> cmmVec l (cmmFloat w)
    MO_VF_Sub  l w      -> cmmVec l (cmmFloat w)
    MO_VF_Mul  l w      -> cmmVec l (cmmFloat w)
    MO_VF_Quot l w      -> cmmVec l (cmmFloat w)
    MO_VF_Neg  l w      -> cmmVec l (cmmFloat w)

    MO_AlignmentCheck _ _ -> ty1
  where
    (ty1:_) = tys

comparisonResultRep :: DynFlags -> CmmType
comparisonResultRep = bWord  -- is it?

-- -----------------------------------------------------------------------------
-- machOpArgReps

-- | This function is used for debugging only: we can check whether an
-- application of a MachOp is "type-correct" by checking that the MachReps of
-- its arguments are the same as the MachOp expects.  This is used when
-- linting a CmmExpr.

machOpArgReps :: DynFlags -> MachOp -> [Width]
machOpArgReps dflags op =
  case op of
    MO_Add    r         -> [r,r]
    MO_Sub    r         -> [r,r]
    MO_Eq     r         -> [r,r]
    MO_Ne     r         -> [r,r]
    MO_Mul    r         -> [r,r]
    MO_S_MulMayOflo r   -> [r,r]
    MO_S_Quot r         -> [r,r]
    MO_S_Rem  r         -> [r,r]
    MO_S_Neg  r         -> [r]
    MO_U_MulMayOflo r   -> [r,r]
    MO_U_Quot r         -> [r,r]
    MO_U_Rem  r         -> [r,r]

    MO_S_Ge r           -> [r,r]
    MO_S_Le r           -> [r,r]
    MO_S_Gt r           -> [r,r]
    MO_S_Lt r           -> [r,r]

    MO_U_Ge r           -> [r,r]
    MO_U_Le r           -> [r,r]
    MO_U_Gt r           -> [r,r]
    MO_U_Lt r           -> [r,r]

    MO_F_Add r          -> [r,r]
    MO_F_Sub r          -> [r,r]
    MO_F_Mul r          -> [r,r]
    MO_F_Quot r         -> [r,r]
    MO_F_Neg r          -> [r]
    MO_F_Eq  r          -> [r,r]
    MO_F_Ne  r          -> [r,r]
    MO_F_Ge  r          -> [r,r]
    MO_F_Le  r          -> [r,r]
    MO_F_Gt  r          -> [r,r]
    MO_F_Lt  r          -> [r,r]

    MO_And   r          -> [r,r]
    MO_Or    r          -> [r,r]
    MO_Xor   r          -> [r,r]
    MO_Not   r          -> [r]
    MO_Shl   r          -> [r, wordWidth dflags]
    MO_U_Shr r          -> [r, wordWidth dflags]
    MO_S_Shr r          -> [r, wordWidth dflags]

    MO_SS_Conv from _   -> [from]
    MO_UU_Conv from _   -> [from]
    MO_XX_Conv from _   -> [from]
    MO_SF_Conv from _   -> [from]
    MO_FS_Conv from _   -> [from]
    MO_FF_Conv from _   -> [from]

    MO_V_Insert  l r    -> [typeWidth (vec l (cmmBits r)),r,wordWidth dflags]
    MO_V_Extract l r    -> [typeWidth (vec l (cmmBits r)),wordWidth dflags]

    MO_V_Add _ r        -> [r,r]
    MO_V_Sub _ r        -> [r,r]
    MO_V_Mul _ r        -> [r,r]

    MO_VS_Quot _ r      -> [r,r]
    MO_VS_Rem  _ r      -> [r,r]
    MO_VS_Neg  _ r      -> [r]

    MO_VU_Quot _ r      -> [r,r]
    MO_VU_Rem  _ r      -> [r,r]

    MO_VF_Insert  l r   -> [typeWidth (vec l (cmmFloat r)),r,wordWidth dflags]
    MO_VF_Extract l r   -> [typeWidth (vec l (cmmFloat r)),wordWidth dflags]

    MO_VF_Add  _ r      -> [r,r]
    MO_VF_Sub  _ r      -> [r,r]
    MO_VF_Mul  _ r      -> [r,r]
    MO_VF_Quot _ r      -> [r,r]
    MO_VF_Neg  _ r      -> [r]

    MO_AlignmentCheck _ r -> [r]

-----------------------------------------------------------------------------
-- CallishMachOp
-----------------------------------------------------------------------------

-- CallishMachOps tend to be implemented by foreign calls in some backends,
-- so we separate them out.  In Cmm, these can only occur in a
-- statement position, in contrast to an ordinary MachOp which can occur
-- anywhere in an expression.
data CallishMachOp
  = MO_F64_Pwr
  | MO_F64_Sin
  | MO_F64_Cos
  | MO_F64_Tan
  | MO_F64_Sinh
  | MO_F64_Cosh
  | MO_F64_Tanh
  | MO_F64_Asin
  | MO_F64_Acos
  | MO_F64_Atan
  | MO_F64_Asinh
  | MO_F64_Acosh
  | MO_F64_Atanh
  | MO_F64_Log
  | MO_F64_Exp
  | MO_F64_Fabs
  | MO_F64_Sqrt
  | MO_F32_Pwr
  | MO_F32_Sin
  | MO_F32_Cos
  | MO_F32_Tan
  | MO_F32_Sinh
  | MO_F32_Cosh
  | MO_F32_Tanh
  | MO_F32_Asin
  | MO_F32_Acos
  | MO_F32_Atan
  | MO_F32_Asinh
  | MO_F32_Acosh
  | MO_F32_Atanh
  | MO_F32_Log
  | MO_F32_Exp
  | MO_F32_Fabs
  | MO_F32_Sqrt

  | MO_UF_Conv Width

  | MO_S_QuotRem Width
  | MO_U_QuotRem Width
  | MO_U_QuotRem2 Width
  | MO_Add2      Width
  | MO_AddWordC  Width
  | MO_SubWordC  Width
  | MO_AddIntC   Width
  | MO_SubIntC   Width
  | MO_U_Mul2    Width

  | MO_WriteBarrier
  | MO_Touch         -- Keep variables live (when using interior pointers)

  -- Prefetch
  | MO_Prefetch_Data Int -- Prefetch hint. May change program performance but not
                     -- program behavior.
                     -- the Int can be 0-3. Needs to be known at compile time
                     -- to interact with code generation correctly.
                     --  TODO: add support for prefetch WRITES,
                     --  currently only exposes prefetch reads, which
                     -- would the majority of use cases in ghc anyways

  -- These three MachOps are parameterised by the known alignment
  -- of the destination and source (for memcpy/memmove) pointers.
  -- This information may be used for optimisation in backends.
  | MO_Memcpy Int
  | MO_Memset Int
  | MO_Memmove Int
  | MO_Memcmp Int

  | MO_PopCnt Width
  | MO_Pdep Width
  | MO_Pext Width
  | MO_Clz Width
  | MO_Ctz Width

  | MO_BSwap Width

  -- Atomic read-modify-write.
  | MO_AtomicRMW Width AtomicMachOp
  | MO_AtomicRead Width
  | MO_AtomicWrite Width
  | MO_Cmpxchg Width
  deriving (Eq, Show)

-- | The operation to perform atomically.
data AtomicMachOp =
      AMO_Add
    | AMO_Sub
    | AMO_And
    | AMO_Nand
    | AMO_Or
    | AMO_Xor
      deriving (Eq, Show)

pprCallishMachOp :: CallishMachOp -> SDoc
pprCallishMachOp mo = text (show mo)

callishMachOpHints :: CallishMachOp -> ([ForeignHint], [ForeignHint])
callishMachOpHints op = case op of
  MO_Memcpy _  -> ([], [AddrHint,AddrHint,NoHint])
  MO_Memset _  -> ([], [AddrHint,NoHint,NoHint])
  MO_Memmove _ -> ([], [AddrHint,AddrHint,NoHint])
  MO_Memcmp _  -> ([], [AddrHint, AddrHint, NoHint])
  _            -> ([],[])
  -- empty lists indicate NoHint

-- | The alignment of a 'memcpy'-ish operation.
machOpMemcpyishAlign :: CallishMachOp -> Maybe Int
machOpMemcpyishAlign op = case op of
  MO_Memcpy  align -> Just align
  MO_Memset  align -> Just align
  MO_Memmove align -> Just align
  MO_Memcmp  align -> Just align
  _                -> Nothing


{-# LANGUAGE CPP #-}

-----------------------------------------------------------------------------
-- A Parser monad with access to the 'DynFlags'.
--
-- The 'P' monad  only has access to the subset of of 'DynFlags'
-- required for parsing Haskell.

-- The parser for C-- requires access to a lot more of the 'DynFlags',
-- so 'PD' provides access to 'DynFlags' via a 'HasDynFlags' instance.
-----------------------------------------------------------------------------
module CmmMonad (
    PD(..)
  , liftP
  ) where

import GhcPrelude

import Control.Monad
import qualified Control.Monad.Fail as MonadFail

import DynFlags
import Lexer

newtype PD a = PD { unPD :: DynFlags -> PState -> ParseResult a }

instance Functor PD where
  fmap = liftM

instance Applicative PD where
  pure = returnPD
  (<*>) = ap

instance Monad PD where
  (>>=) = thenPD
#if !MIN_VERSION_base(4,13,0)
  fail = MonadFail.fail
#endif

instance MonadFail.MonadFail PD where
  fail = failPD

liftP :: P a -> PD a
liftP (P f) = PD $ \_ s -> f s

returnPD :: a -> PD a
returnPD = liftP . return

thenPD :: PD a -> (a -> PD b) -> PD b
(PD m) `thenPD` k = PD $ \d s ->
        case m d s of
                POk s1 a         -> unPD (k a) d s1
                PFailed s1 -> PFailed s1

failPD :: String -> PD a
failPD = liftP . fail

instance HasDynFlags PD where
   getDynFlags = PD $ \d s -> POk s d