content: {-# LANGUAGE GADTs #-}
  module GHC.Cmm.Switch.Implement
    ( cmmImplementSwitchPlans
    )
  where
  
  import GHC.Prelude
  
  import GHC.Platform
  import GHC.Cmm.Dataflow.Block
  import GHC.Cmm.BlockId
  import GHC.Cmm
  import GHC.Cmm.Utils
  import GHC.Cmm.Switch
  import GHC.Utils.Monad (concatMapM)
  import GHC.Types.Unique.DSM
  
  --
  -- This module replaces Switch statements as generated by the Stg -> Cmm
  -- transformation, which might be huge and sparse and hence unsuitable for
  -- assembly code, by proper constructs (if-then-else trees, dense jump tables).
  --
  -- The actual, abstract strategy is determined by createSwitchPlan in
  -- GHC.Cmm.Switch and returned as a SwitchPlan; here is just the implementation in
  -- terms of Cmm code. See Note [Cmm Switches, the general plan] in GHC.Cmm.Switch.
  --
  -- This division into different modules is both to clearly separate concerns,
  -- but also because createSwitchPlan needs access to the constructors of
  -- SwitchTargets, a data type exported abstractly by GHC.Cmm.Switch.
  --
  
  -- | Traverses the 'CmmGraph', making sure that 'CmmSwitch' are suitable for
  -- code generation.
  cmmImplementSwitchPlans :: Platform -> CmmGraph -> UniqDSM CmmGraph
  cmmImplementSwitchPlans platform g =
      -- Switch generation done by backend (LLVM/C)
      do
      blocks' <- concatMapM (visitSwitches platform) (toBlockList g)
      return $ ofBlockList (g_entry g) blocks'
  
  visitSwitches :: Platform -> CmmBlock -> UniqDSM [CmmBlock]
  visitSwitches platform block
    | (entry@(CmmEntry _ scope), middle, CmmSwitch vanillaExpr ids) <- blockSplit block
    = do
      let plan = createSwitchPlan ids
      -- See Note [Floating switch expressions]
      (assignSimple, simpleExpr) <- floatSwitchExpr platform vanillaExpr
  
      (newTail, newBlocks) <- implementSwitchPlan platform scope simpleExpr plan
  
      let block' = entry `blockJoinHead` middle `blockAppend` assignSimple `blockAppend` newTail
  
      return $ block' : newBlocks
  
    | otherwise
    = return [block]
  
  -- Note [Floating switch expressions]
  -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  -- When we translate a sparse switch into a search tree we would like
  -- to compute the value we compare against only once.
  --
  -- For this purpose we assign the switch expression to a local register
  -- and then use this register when constructing the actual binary tree.
  --
  -- This is important as the expression could contain expensive code like
  -- memory loads or divisions which we REALLY don't want to duplicate.
  --
  -- This happened in parts of the handwritten RTS Cmm code. See also #16933
  
  -- See Note [Floating switch expressions]
  floatSwitchExpr :: Platform -> CmmExpr -> UniqDSM (Block CmmNode O O, CmmExpr)
  floatSwitchExpr _        reg@(CmmReg {})  = return (emptyBlock, reg)
  floatSwitchExpr platform expr             = do
    (assign, expr') <- cmmMkAssign platform expr <$> getUniqueDSM
    return (BMiddle assign, expr')
  
  
  -- Implementing a switch plan (returning a tail block)
  implementSwitchPlan :: Platform -> CmmTickScope -> CmmExpr -> SwitchPlan -> UniqDSM (Block CmmNode O C, [CmmBlock])
  implementSwitchPlan platform scope expr = go
    where
      width = typeWidth $ cmmExprType platform expr
  
      go (Unconditionally l)
        = return (emptyBlock `blockJoinTail` CmmBranch l, [])
      go (JumpTable ids)
        = return (emptyBlock `blockJoinTail` CmmSwitch expr ids, [])
      go (IfLT signed i ids1 ids2)
        = do
          (bid1, newBlocks1) <- go' ids1
          (bid2, newBlocks2) <- go' ids2
  
          let lt | signed    = MO_S_Lt
                 | otherwise = MO_U_Lt
              scrut = CmmMachOp (lt width) [expr, CmmLit $ CmmInt i width]
              lastNode = CmmCondBranch scrut bid1 bid2 Nothing
              lastBlock = emptyBlock `blockJoinTail` lastNode
          return (lastBlock, newBlocks1++newBlocks2)
      go (IfEqual i l ids2)
        = do
          (bid2, newBlocks2) <- go' ids2
  
          let scrut = CmmMachOp (MO_Ne width) [expr, CmmLit $ CmmInt i width]
              lastNode = CmmCondBranch scrut bid2 l Nothing
              lastBlock = emptyBlock `blockJoinTail` lastNode
          return (lastBlock, newBlocks2)
  
      -- Same but returning a label to branch to
      go' (Unconditionally l)
        = return (l, [])
      go' p
        = do
          bid <- mkBlockId `fmap` getUniqueDSM
          (last, newBlocks) <- go p
          let block = CmmEntry bid scope `blockJoinHead` last
          return (bid, block: newBlocks)
  