compiler: properly handle non-word-sized CmmSwitch scrutinees in the wasm NCGwip/T21776

Currently, the wasm NCG has an implicit assumption: all CmmSwitch
scrutinees are 32-bit integers. This is not always true; #22864 is one
counter-example with a 64-bit scrutinee. This patch fixes the logic by
explicitly converting the scrutinee to a word that can be used as a
br_table operand. Fixes #22871. Also includes a regression test.

5 files changed, 49 insertions, 1 deletions

diff --git a/compiler/GHC/Wasm/ControlFlow/FromCmm.hs b/compiler/GHC/Wasm/ControlFlow/FromCmm.hs
index 0667345162..85eb8e534b 100644
--- a/compiler/GHC/Wasm/ControlFlow/FromCmm.hs
+++ b/compiler/GHC/Wasm/ControlFlow/FromCmm.hs
@@ -75,7 +75,7 @@ flowLeaving platform b =
           let (offset, target_labels) = switchTargetsToTable targets
               (lo, hi) = switchTargetsRange targets
               default_label = switchTargetsDefault targets
-              scrutinee = smartPlus platform e offset
+              scrutinee = smartExtend platform $ smartPlus platform e offset
               range = inclusiveInterval (lo+toInteger offset) (hi+toInteger offset)
           in  Switch scrutinee range target_labels default_label
       CmmCall { cml_cont = Nothing, cml_target = e } -> TailCall e
@@ -314,6 +314,14 @@ structuredControl platform txExpr txBlock g =
 nodeBody :: CmmBlock -> CmmActions
 nodeBody (BlockCC _first middle _last) = middle
 
+-- | A CmmSwitch scrutinee may have any width, but a br_table operand
+-- must be exactly word sized, hence the extension here. (#22871)
+smartExtend :: Platform -> CmmExpr -> CmmExpr
+smartExtend p e | w0 == w1 = e
+                | otherwise = CmmMachOp (MO_UU_Conv w0 w1) [e]
+  where
+    w0 = cmmExprWidth p e
+    w1 = wordWidth p
 
 smartPlus :: Platform -> CmmExpr -> Int -> CmmExpr
 smartPlus _ e 0 = e
diff --git a/testsuite/tests/cmm/should_run/T22871.hs b/testsuite/tests/cmm/should_run/T22871.hs
new file mode 100644
index 0000000000..a990e8a793
--- /dev/null
+++ b/testsuite/tests/cmm/should_run/T22871.hs
@@ -0,0 +1,12 @@
+{-# LANGUAGE GHCForeignImportPrim #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE UnliftedFFITypes #-}
+
+import Data.Foldable
+import GHC.Exts
+import GHC.Int
+
+foreign import prim "foo" foo :: Int64# -> Int64#
+
+main :: IO ()
+main = for_ [0, 42, 114514] $ \(I64# x#) -> print $ I64# (foo x#)
diff --git a/testsuite/tests/cmm/should_run/T22871.stdout b/testsuite/tests/cmm/should_run/T22871.stdout
new file mode 100644
index 0000000000..9f832c4f9b
--- /dev/null
+++ b/testsuite/tests/cmm/should_run/T22871.stdout
@@ -0,0 +1,3 @@
+233
+84
+1919810
diff --git a/testsuite/tests/cmm/should_run/T22871_cmm.cmm b/testsuite/tests/cmm/should_run/T22871_cmm.cmm
new file mode 100644
index 0000000000..f8d5483127
--- /dev/null
+++ b/testsuite/tests/cmm/should_run/T22871_cmm.cmm
@@ -0,0 +1,16 @@
+#include "Cmm.h"
+
+foo (I64 x) {
+    switch [0 .. 114514] (x) {
+    case 0: { return (233 :: I64); }
+    case 1: { return (333 :: I64); }
+    case 2: { return (666 :: I64); }
+    case 3: { return (1989 :: I64); }
+    case 4: { return (1997 :: I64); }
+    case 5: { return (2012 :: I64); }
+    case 6: { return (2019 :: I64); }
+    case 7: { return (2022 :: I64); }
+    case 114514: { return (1919810 :: I64); }
+    default: { return (x * (2 :: I64)); }
+  }
+}
diff --git a/testsuite/tests/cmm/should_run/all.T b/testsuite/tests/cmm/should_run/all.T
index 7290fee784..2a50208476 100644
--- a/testsuite/tests/cmm/should_run/all.T
+++ b/testsuite/tests/cmm/should_run/all.T
@@ -25,3 +25,12 @@ test('ByteSwitch',
      ],
      multi_compile_and_run,
      ['ByteSwitch', [('ByteSwitch_cmm.cmm', '')], ''])
+
+test('T22871',
+     [    extra_run_opts('"' + config.libdir + '"')
+     ,    omit_ways(['ghci'])
+     ,    js_skip
+     ,    when(arch('i386'), skip) # x86 NCG panics with "iselExpr64(i386)"
+     ],
+     multi_compile_and_run,
+     ['T22871', [('T22871_cmm.cmm', '')], ''])