Change jump targets in JMP_TBL from blocks to X86.JumpDest.

Jump tables always point to blocks when we first generate them.  However
there are rare situations where we can shortcut one of these blocks to a
static address during the asm shortcutting pass.

While we already updated the data section accordingly this patch also
extends this to the references stored in JMP_TBL.

Test Plan: ci

Reviewers: bgamari

Reviewed By: bgamari

Subscribers: thomie, carter

GHC Trac Issues: #15104

Differential Revision: https://phabricator.haskell.org/D4595

2 files changed, 19 insertions, 12 deletions

diff --git a/compiler/nativeGen/X86/CodeGen.hs b/compiler/nativeGen/X86/CodeGen.hs
index 579c726923..4551754200 100644
--- a/compiler/nativeGen/X86/CodeGen.hs
+++ b/compiler/nativeGen/X86/CodeGen.hs
@@ -2857,11 +2857,16 @@ genSwitch dflags expr targets
                     JMP_TBL op ids (Section ReadOnlyData lbl) lbl
                  ]
         return code
-  where (offset, ids) = switchTargetsToTable targets
+  where
+    (offset, blockIds) = switchTargetsToTable targets
+    ids = map (fmap DestBlockId) blockIds
 
 generateJumpTableForInstr :: DynFlags -> Instr -> Maybe (NatCmmDecl (Alignment, CmmStatics) Instr)
 generateJumpTableForInstr dflags (JMP_TBL _ ids section lbl)
-    = Just (createJumpTable dflags ids section lbl)
+    = let getBlockId (DestBlockId id) = id
+          getBlockId _ = panic "Non-Label target in Jump Table"
+          blockIds = map (fmap getBlockId) ids
+      in Just (createJumpTable dflags blockIds section lbl)
 generateJumpTableForInstr _ _ = Nothing
 
 createJumpTable :: DynFlags -> [Maybe BlockId] -> Section -> CLabel
diff --git a/compiler/nativeGen/X86/Instr.hs b/compiler/nativeGen/X86/Instr.hs
index d15f2f784e..ee3e64cf24 100644
--- a/compiler/nativeGen/X86/Instr.hs
+++ b/compiler/nativeGen/X86/Instr.hs
@@ -8,7 +8,7 @@
 --
 -----------------------------------------------------------------------------
 
-module X86.Instr (Instr(..), Operand(..), PrefetchVariant(..), JumpDest,
+module X86.Instr (Instr(..), Operand(..), PrefetchVariant(..), JumpDest(..),
                   getJumpDestBlockId, canShortcut, shortcutStatics,
                   shortcutJump, i386_insert_ffrees, allocMoreStack,
                   maxSpillSlots, archWordFormat)
@@ -322,7 +322,7 @@ data Instr
         | JXX_GBL     Cond Imm      -- non-local version of JXX
         -- Table jump
         | JMP_TBL     Operand   -- Address to jump to
-                      [Maybe BlockId] -- Blocks in the jump table
+                      [Maybe JumpDest] -- Targets of the jump table
                       Section   -- Data section jump table should be put in
                       CLabel    -- Label of jump table
         | CALL        (Either Imm Reg) [Reg]
@@ -704,7 +704,7 @@ x86_jumpDestsOfInstr
 x86_jumpDestsOfInstr insn
   = case insn of
         JXX _ id        -> [id]
-        JMP_TBL _ ids _ _ -> [id | Just id <- ids]
+        JMP_TBL _ ids _ _ -> [id | Just (DestBlockId id) <- ids]
         _               -> []
 
 
@@ -715,8 +715,12 @@ x86_patchJumpInstr insn patchF
   = case insn of
         JXX cc id       -> JXX cc (patchF id)
         JMP_TBL op ids section lbl
-          -> JMP_TBL op (map (fmap patchF) ids) section lbl
+          -> JMP_TBL op (map (fmap (patchJumpDest patchF)) ids) section lbl
         _               -> insn
+    where
+        patchJumpDest f (DestBlockId id) = DestBlockId (f id)
+        patchJumpDest _ dest             = dest
+
 
 
 
@@ -1036,13 +1040,11 @@ shortcutJump fn insn = shortcutJump' fn (setEmpty :: LabelSet) insn
             Just (DestImm imm)     -> shortcutJump' fn seen' (JXX_GBL cc imm)
         where seen' = setInsert id seen
     shortcutJump' fn _ (JMP_TBL addr blocks section tblId) =
-        let updateBlock Nothing     = Nothing
-            updateBlock (Just bid)  =
+        let updateBlock (Just (DestBlockId bid))  =
                 case fn bid of
-                    Nothing                 -> Just bid
-                    Just (DestBlockId bid') -> Just bid'
-                    Just (DestImm _)        ->
-                        panic "Can't shortcut jump table to immediate"
+                    Nothing   -> Just (DestBlockId bid )
+                    Just dest -> Just dest
+            updateBlock dest = dest
             blocks' = map updateBlock blocks
         in  JMP_TBL addr blocks' section tblId
     shortcutJump' _ _ other = other