Update JMP_TBL targets during shortcutting in X86 NCG.

Without updating the JMP_TBL information the block list in
JMP_TBL contained blocks which were eliminated in some circumstances.

The actual assembly generation doesn't look at these fields so this
didn't cause any bugs yet. However as long as we carry this information
around we should make an effort to keep it correct.

Especially since it's useful for debugging purposes and can be used
for passes near the end of the codegen pipeline.
In particular it's used by jumpDestsOfInstr which without these changes
returns the wrong destinations.

Test Plan: ci

Reviewers: bgamari

Subscribers: thomie, carter

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

1 files changed, 20 insertions, 8 deletions

diff --git a/compiler/nativeGen/X86/Instr.hs b/compiler/nativeGen/X86/Instr.hs
index f4f625b4a5..49beafa1ff 100644
--- a/compiler/nativeGen/X86/Instr.hs
+++ b/compiler/nativeGen/X86/Instr.hs
@@ -1026,14 +1026,26 @@ canShortcut _                    = Nothing
 -- The blockset helps avoid following cycles.
 shortcutJump :: (BlockId -> Maybe JumpDest) -> Instr -> Instr
 shortcutJump fn insn = shortcutJump' fn (setEmpty :: LabelSet) insn
-  where shortcutJump' fn seen insn@(JXX cc id) =
-          if setMember id seen then insn
-          else case fn id of
-                 Nothing                -> insn
-                 Just (DestBlockId id') -> shortcutJump' fn seen' (JXX cc id')
-                 Just (DestImm imm)     -> shortcutJump' fn seen' (JXX_GBL cc imm)
-               where seen' = setInsert id seen
-        shortcutJump' _ _ other = other
+  where
+    shortcutJump' :: (BlockId -> Maybe JumpDest) -> LabelSet -> Instr -> Instr
+    shortcutJump' fn seen insn@(JXX cc id) =
+        if setMember id seen then insn
+        else case fn id of
+            Nothing                -> insn
+            Just (DestBlockId id') -> shortcutJump' fn seen' (JXX cc id')
+            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)  =
+                case fn bid of
+                    Nothing                 -> Just bid
+                    Just (DestBlockId bid') -> Just bid'
+                    Just (DestImm _)        ->
+                        panic "Can't shortcut jump table to immediate"
+            blocks' = map updateBlock blocks
+        in  JMP_TBL addr blocks' section tblId
+    shortcutJump' _ _ other = other
 
 -- Here because it knows about JumpDest
 shortcutStatics :: (BlockId -> Maybe JumpDest) -> (Alignment, CmmStatics) -> (Alignment, CmmStatics)