    {-# 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.Types.Unique.Supply
    import GHC.Utils.Monad (concatMapM)

    --
    -- 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 -> UniqSM CmmGraph
    cmmImplementSwitchPlans :: Platform -> CmmGraph -> UniqSM CmmGraph
    cmmImplementSwitchPlans Platform
    platform CmmGraph
    g =
        -- Switch generation done by backend (LLVM/C)
        do
        blocks' <- (CmmBlock -> UniqSM [CmmBlock]) -> [CmmBlock] -> UniqSM [CmmBlock]
    forall (m :: * -> *) (f :: * -> *) a b.
    (Monad m, Traversable f) =>
    (a -> m [b]) -> f a -> m [b]
    concatMapM (Platform -> CmmBlock -> UniqSM [CmmBlock]
    visitSwitches Platform
    platform) (CmmGraph -> [CmmBlock]
    toBlockList CmmGraph
    g)
        return $ ofBlockList (g_entry g) blocks'

    visitSwitches :: Platform -> CmmBlock -> UniqSM [CmmBlock]
    visitSwitches :: Platform -> CmmBlock -> UniqSM [CmmBlock]
    visitSwitches Platform
    platform CmmBlock
    block
      | (entry :: CmmNode C O
    entry@(CmmEntry BlockId
    _ CmmTickScope
    scope), Block CmmNode O O
    middle, CmmSwitch CmmExpr
    vanillaExpr SwitchTargets
    ids) <- CmmBlock -> (CmmNode C O, Block CmmNode O O, CmmNode O C)
    forall (n :: Extensibility -> Extensibility -> *).
    Block n C C -> (n C O, Block n O O, n O C)
    blockSplit CmmBlock
    block
      = do
        let plan :: SwitchPlan
    plan = SwitchTargets -> SwitchPlan
    createSwitchPlan SwitchTargets
    ids
        -- See Note [Floating switch expressions]
        (assignSimple, simpleExpr) <- Platform -> CmmExpr -> UniqSM (Block CmmNode O O, CmmExpr)
    floatSwitchExpr Platform
    platform CmmExpr
    vanillaExpr

        (newTail, newBlocks) <- implementSwitchPlan platform scope simpleExpr plan

        let block' = CmmNode C O
    entry CmmNode C O -> Block CmmNode O O -> Block CmmNode C O
    forall (n :: Extensibility -> Extensibility -> *)
           (x :: Extensibility).
    n C O -> Block n O x -> Block n C x
    `blockJoinHead` Block CmmNode O O
    middle Block CmmNode C O -> Block CmmNode O O -> Block CmmNode C O
    forall (n :: Extensibility -> Extensibility -> *)
           (e :: Extensibility) (x :: Extensibility).
    Block n e O -> Block n O x -> Block n e x
    `blockAppend` Block CmmNode O O
    assignSimple Block CmmNode C O -> Block CmmNode O C -> CmmBlock
    forall (n :: Extensibility -> Extensibility -> *)
           (e :: Extensibility) (x :: Extensibility).
    Block n e O -> Block n O x -> Block n e x
    `blockAppend` Block CmmNode O C
    newTail

        return $ block' : newBlocks

      | Bool
    otherwise
      = [CmmBlock] -> UniqSM [CmmBlock]
    forall a. a -> UniqSM a
    forall (m :: * -> *) a. Monad m => a -> m a
    return [CmmBlock
    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 -> UniqSM (Block CmmNode O O, CmmExpr)
    floatSwitchExpr :: Platform -> CmmExpr -> UniqSM (Block CmmNode O O, CmmExpr)
    floatSwitchExpr Platform
    _        reg :: CmmExpr
    reg@(CmmReg {})  = (Block CmmNode O O, CmmExpr) -> UniqSM (Block CmmNode O O, CmmExpr)
    forall a. a -> UniqSM a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Block CmmNode O O
    forall (n :: Extensibility -> Extensibility -> *). Block n O O
    emptyBlock, CmmExpr
    reg)
    floatSwitchExpr Platform
    platform CmmExpr
    expr             = do
      (assign, expr') <- Platform -> CmmExpr -> Unique -> (CmmNode O O, CmmExpr)
    cmmMkAssign Platform
    platform CmmExpr
    expr (Unique -> (CmmNode O O, CmmExpr))
    -> UniqSM Unique -> UniqSM (CmmNode O O, CmmExpr)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> UniqSM Unique
    forall (m :: * -> *). MonadUnique m => m Unique
    getUniqueM
      return (BMiddle assign, expr')

    -- Implementing a switch plan (returning a tail block)
    implementSwitchPlan :: Platform -> CmmTickScope -> CmmExpr -> SwitchPlan -> UniqSM (Block CmmNode O C, [CmmBlock])
    implementSwitchPlan :: Platform
    -> CmmTickScope
    -> CmmExpr
    -> SwitchPlan
    -> UniqSM (Block CmmNode O C, [CmmBlock])
    implementSwitchPlan Platform
    platform CmmTickScope
    scope CmmExpr
    expr = SwitchPlan -> UniqSM (Block CmmNode O C, [CmmBlock])
    go
      where
        width :: Width
    width = CmmType -> Width
    typeWidth (CmmType -> Width) -> CmmType -> Width
    forall a b. (a -> b) -> a -> b
    $ Platform -> CmmExpr -> CmmType
    cmmExprType Platform
    platform CmmExpr
    expr

        go :: SwitchPlan -> UniqSM (Block CmmNode O C, [CmmBlock])
    go (Unconditionally BlockId
    l)
          = (Block CmmNode O C, [CmmBlock])
    -> UniqSM (Block CmmNode O C, [CmmBlock])
    forall a. a -> UniqSM a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Block CmmNode O O
    forall (n :: Extensibility -> Extensibility -> *). Block n O O
    emptyBlock Block CmmNode O O -> CmmNode O C -> Block CmmNode O C
    forall (n :: Extensibility -> Extensibility -> *)
           (e :: Extensibility).
    Block n e O -> n O C -> Block n e C
    `blockJoinTail` BlockId -> CmmNode O C
    CmmBranch BlockId
    l, [])
        go (JumpTable SwitchTargets
    ids)
          = (Block CmmNode O C, [CmmBlock])
    -> UniqSM (Block CmmNode O C, [CmmBlock])
    forall a. a -> UniqSM a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Block CmmNode O O
    forall (n :: Extensibility -> Extensibility -> *). Block n O O
    emptyBlock Block CmmNode O O -> CmmNode O C -> Block CmmNode O C
    forall (n :: Extensibility -> Extensibility -> *)
           (e :: Extensibility).
    Block n e O -> n O C -> Block n e C
    `blockJoinTail` CmmExpr -> SwitchTargets -> CmmNode O C
    CmmSwitch CmmExpr
    expr SwitchTargets
    ids, [])
        go (IfLT Bool
    signed Integer
    i SwitchPlan
    ids1 SwitchPlan
    ids2)
          = do
            (bid1, newBlocks1) <- SwitchPlan -> UniqSM (BlockId, [CmmBlock])
    go' SwitchPlan
    ids1
            (bid2, newBlocks2) <- go' ids2

            let lt | Bool
    signed    = Width -> MachOp
    MO_S_Lt
                   | Bool
    otherwise = Width -> MachOp
    MO_U_Lt
                scrut = MachOp -> [CmmExpr] -> CmmExpr
    CmmMachOp (Width -> MachOp
    lt Width
    width) [CmmExpr
    expr, CmmLit -> CmmExpr
    CmmLit (CmmLit -> CmmExpr) -> CmmLit -> CmmExpr
    forall a b. (a -> b) -> a -> b
    $ Integer -> Width -> CmmLit
    CmmInt Integer
    i Width
    width]
                lastNode = CmmExpr -> BlockId -> BlockId -> Maybe Bool -> CmmNode O C
    CmmCondBranch CmmExpr
    scrut BlockId
    bid1 BlockId
    bid2 Maybe Bool
    forall a. Maybe a
    Nothing
                lastBlock = Block CmmNode O O
    forall (n :: Extensibility -> Extensibility -> *). Block n O O
    emptyBlock Block CmmNode O O -> CmmNode O C -> Block CmmNode O C
    forall (n :: Extensibility -> Extensibility -> *)
           (e :: Extensibility).
    Block n e O -> n O C -> Block n e C
    `blockJoinTail` CmmNode O C
    lastNode
            return (lastBlock, newBlocks1++newBlocks2)
        go (IfEqual Integer
    i BlockId
    l SwitchPlan
    ids2)
          = do
            (bid2, newBlocks2) <- SwitchPlan -> UniqSM (BlockId, [CmmBlock])
    go' SwitchPlan
    ids2

            let scrut = MachOp -> [CmmExpr] -> CmmExpr
    CmmMachOp (Width -> MachOp
    MO_Ne Width
    width) [CmmExpr
    expr, CmmLit -> CmmExpr
    CmmLit (CmmLit -> CmmExpr) -> CmmLit -> CmmExpr
    forall a b. (a -> b) -> a -> b
    $ Integer -> Width -> CmmLit
    CmmInt Integer
    i Width
    width]
                lastNode = CmmExpr -> BlockId -> BlockId -> Maybe Bool -> CmmNode O C
    CmmCondBranch CmmExpr
    scrut BlockId
    bid2 BlockId
    l Maybe Bool
    forall a. Maybe a
    Nothing
                lastBlock = Block CmmNode O O
    forall (n :: Extensibility -> Extensibility -> *). Block n O O
    emptyBlock Block CmmNode O O -> CmmNode O C -> Block CmmNode O C
    forall (n :: Extensibility -> Extensibility -> *)
           (e :: Extensibility).
    Block n e O -> n O C -> Block n e C
    `blockJoinTail` CmmNode O C
    lastNode
            return (lastBlock, newBlocks2)

        -- Same but returning a label to branch to
        go' :: SwitchPlan -> UniqSM (BlockId, [CmmBlock])
    go' (Unconditionally BlockId
    l)
          = (BlockId, [CmmBlock]) -> UniqSM (BlockId, [CmmBlock])
    forall a. a -> UniqSM a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (BlockId
    l, [])
        go' SwitchPlan
    p
          = do
            bid <- Unique -> BlockId
    mkBlockId (Unique -> BlockId) -> UniqSM Unique -> UniqSM BlockId
    forall a b. (a -> b) -> UniqSM a -> UniqSM b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    `fmap` UniqSM Unique
    forall (m :: * -> *). MonadUnique m => m Unique
    getUniqueM
            (last, newBlocks) <- go p
            let block = BlockId -> CmmTickScope -> CmmNode C O
    CmmEntry BlockId
    bid CmmTickScope
    scope CmmNode C O -> Block CmmNode O C -> CmmBlock
    forall (n :: Extensibility -> Extensibility -> *)
           (x :: Extensibility).
    n C O -> Block n O x -> Block n C x
    `blockJoinHead` Block CmmNode O C
    last
            return (bid, block: newBlocks)