ncg/aarch64: Fix sub-word sign extension yet again

In adc7f108141a973b6dcb02a7836eed65d61230e8 we fixed a number of issues
to do with sign extension in the AArch64 NCG found by ghc/test-primops>.
However, this patch made a critical error, assuming that getSomeReg
would allocate a fresh register for the result of its evaluation.
However, this is not the case as `getSomeReg (CmmReg r) == r`.
Consequently, any mutation of the register returned by `getSomeReg` may
have unwanted side-effects on other expressions also mentioning `r`. In
the fix listed above, this manifested as the registers containing the
operands of binary arithmetic operations being incorrectly
sign-extended. This resulted in #22282.

Sadly, the rather simple structure of the tests generated
by `test-primops` meant that this particular case was not exercised.
Even more surprisingly, none of our testsuite caught this case.

Here we fix this by ensuring that intermediate sign extension is
performed in a fresh register.

Fixes #22282.

1 files changed, 20 insertions, 12 deletions

diff --git a/compiler/GHC/CmmToAsm/AArch64/CodeGen.hs b/compiler/GHC/CmmToAsm/AArch64/CodeGen.hs
index af2d63cc8b..c1161832d1 100644
--- a/compiler/GHC/CmmToAsm/AArch64/CodeGen.hs
+++ b/compiler/GHC/CmmToAsm/AArch64/CodeGen.hs
@@ -664,10 +664,11 @@ getRegister' config plat expr
         -- See Note [Signed arithmetic on AArch64].
         negate code w reg = do
             let w' = opRegWidth w
+            (reg', code_sx) <- signExtendReg w w' reg
             return $ Any (intFormat w) $ \dst ->
                 code `appOL`
-                signExtendReg w w' reg `snocOL`
-                NEG (OpReg w' dst) (OpReg w' reg) `appOL`
+                code_sx `snocOL`
+                NEG (OpReg w' dst) (OpReg w' reg') `appOL`
                 truncateReg w' w dst
 
         ss_conv from to reg code =
@@ -817,15 +818,17 @@ getRegister' config plat expr
               -- should be performed.
               let w' = opRegWidth w
                   signExt r
-                    | not is_signed  = nilOL
+                    | not is_signed  = return (r, nilOL)
                     | otherwise      = signExtendReg w w' r
+              (reg_x_sx, code_x_sx) <- signExt reg_x
+              (reg_y_sx, code_y_sx) <- signExt reg_y
               return $ Any (intFormat w) $ \dst ->
                   code_x `appOL`
                   code_y `appOL`
                   -- sign-extend both operands
-                  signExt reg_x `appOL`
-                  signExt reg_y `appOL`
-                  op (OpReg w' dst) (OpReg w' reg_x) (OpReg w' reg_y) `appOL`
+                  code_x_sx `appOL`
+                  code_y_sx `appOL`
+                  op (OpReg w' dst) (OpReg w' reg_x_sx) (OpReg w' reg_y_sx) `appOL`
                   truncateReg w' w dst -- truncate back to the operand's original width
 
           floatOp w op = do
@@ -1021,16 +1024,21 @@ getRegister' config plat expr
 
 -- | Instructions to sign-extend the value in the given register from width @w@
 -- up to width @w'@.
-signExtendReg :: Width -> Width -> Reg -> OrdList Instr
+signExtendReg :: Width -> Width -> Reg -> NatM (Reg, OrdList Instr)
 signExtendReg w w' r =
     case w of
-      W64 -> nilOL
+      W64 -> noop
       W32
-        | w' == W32 -> nilOL
-        | otherwise -> unitOL $ SXTH (OpReg w' r) (OpReg w' r)
-      W16           -> unitOL $ SXTH (OpReg w' r) (OpReg w' r)
-      W8            -> unitOL $ SXTB (OpReg w' r) (OpReg w' r)
+        | w' == W32 -> noop
+        | otherwise -> extend SXTH
+      W16           -> extend SXTH
+      W8            -> extend SXTB
       _             -> panic "intOp"
+  where
+    noop = return (r, nilOL)
+    extend instr = do
+        r' <- getNewRegNat II64
+        return (r', unitOL $ instr (OpReg w' r') (OpReg w' r))
 
 -- | Instructions to truncate the value in the given register from width @w@
 -- down to width @w'@.