Fix combination of ArityType in andArityType

When combining the ArityType of two case branches we need to make the
conservative decision to

Before this patch `\1. T` when combined with `T` would result in `\1.
T`, the result being that we would then eta-expand the branch of type
`T` (even though we concluded it wasn't necessarily safe to do so).

In particular, this goes wrong when the branch contains a call to a join point, if
we decide to eta-expand it anyway then the join-point gets oversatured.
This is a bit of latent bug which was only triggered quite indirectly by
inserting cost-centres but seems like it could have happened in other
scenarios.

Therefore the correct result of combining `\1. T` and `T` is the
conservative `T`. This patch corrects the logic in `andArityType` to
produce that result.

Fixes #21694

-------------------------
Metric Increase:
    ManyAlternatives
    ManyConstructors
    MultiComponentModules
    MultiComponentModulesRecomp
    MultiLayerModules
    MultiLayerModulesRecomp
    T10421
    T12425
    T12707
    T13035
    T13379
    T13701
    T14683
    T15703
    T16875
    T1969
    T3064
    T3294
    T4801
    T5321FD
    T5321Fun
    T5631
    T783
    T9020
    T9198
    T9233
    T9961
-------------------------

(cherry picked from commit 07e7d0fd84662074ce73ed0d5e19ffe849a7aa36)

3 files changed, 103 insertions, 22 deletions

diff --git a/compiler/GHC/Core/Opt/Arity.hs b/compiler/GHC/Core/Opt/Arity.hs
index cfa0ad93c3..795d256917 100644
--- a/compiler/GHC/Core/Opt/Arity.hs
+++ b/compiler/GHC/Core/Opt/Arity.hs
@@ -776,12 +776,12 @@ andArityType :: ArityType -> ArityType -> ArityType
 andArityType (AT (os1:oss1) div1) (AT (os2:oss2) div2)
   | AT oss' div' <- andArityType (AT oss1 div1) (AT oss2 div2)
   = AT ((os1 `bestOneShot` os2) : oss') div' -- See Note [Combining case branches]
-andArityType (AT []         div1) at2
+andArityType at1@(AT []         div1) at2
   | isDeadEndDiv div1 = at2                  -- Note [ABot branches: max arity wins]
-  | otherwise         = takeWhileOneShot at2 -- See Note [Combining case branches]
-andArityType at1                  (AT []         div2)
+  | otherwise         = at1                  -- See Note [Combining case branches]
+andArityType at1                  at2@(AT []         div2)
   | isDeadEndDiv div2 = at1                  -- Note [ABot branches: max arity wins]
-  | otherwise         = takeWhileOneShot at1 -- See Note [Combining case branches]
+  | otherwise         = at2                  -- See Note [Combining case branches]
 
 {- Note [ABot branches: max arity wins]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -794,24 +794,13 @@ So we need \??.⊥ for the whole thing, the /max/ of both arities.
 
 Note [Combining case branches]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-Consider
-  go = \x. let z = go e0
-               go2 = \x. case x of
-                           True  -> z
-                           False -> \s(one-shot). e1
-           in go2 x
-We *really* want to respect the one-shot annotation provided by the
-user and eta-expand go and go2.
-When combining the branches of the case we have
-     T `andAT` \1.T
-and we want to get \1.T.
-But if the inner lambda wasn't one-shot (\?.T) we don't want to do this.
-(We need a usage analysis to justify that.)
-
-So we combine the best of the two branches, on the (slightly dodgy)
-basis that if we know one branch is one-shot, then they all must be.
-Surprisingly, this means that the one-shot arity type is effectively the top
-element of the lattice.
+
+Unless we can conclude that **all** branches are safe to eta-expand then we
+must pessimisticaly conclude that we can't eta-expand. See #21694 for where this
+went wrong.
+We can do better in the long run, but for the 9.4/9.2 branches we choose to simply
+ignore oneshot annotations for the time being.
+
 
 Note [Arity trimming]
 ~~~~~~~~~~~~~~~~~~~~~
diff --git a/testsuite/tests/simplCore/should_compile/T21694.hs b/testsuite/tests/simplCore/should_compile/T21694.hs
new file mode 100644
index 0000000000..98c5a55c59
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T21694.hs
@@ -0,0 +1,91 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE MagicHash #-}
+{-# OPTIONS_GHC -Wall #-}
+module Bug (go_fast_end) where
+
+import Control.Monad.ST (ST)
+import qualified Data.ByteString.Internal as BS
+import qualified Data.ByteString.Unsafe as BS
+import Data.ByteString (ByteString)
+import Foreign.ForeignPtr (withForeignPtr)
+import Foreign.Ptr (plusPtr)
+import GHC.Exts ( Int(..), Int#, Ptr(..), Word(..)
+                , (<#), (>#), indexWord64OffAddr#, isTrue#, orI#
+                )
+import GHC.Word (Word8(..), Word64(..))
+import System.IO.Unsafe (unsafeDupablePerformIO)
+
+#if MIN_VERSION_ghc_prim(0,8,0)
+import GHC.Exts (word8ToWord#)
+#endif
+
+#if __GLASGOW_HASKELL__ >= 904
+import GHC.Exts (byteSwap64#, int64ToInt#, word64ToInt64#, ltWord64#, wordToWord64#)
+#else
+import GHC.Exts (byteSwap#, ltWord#, word2Int#)
+#endif
+
+go_fast_end :: ByteString -> DecodeAction s a -> ST s (SlowPath s a)
+go_fast_end !bs (ConsumeInt32 k) =
+  case tryConsumeInt (BS.unsafeHead bs) bs of
+    DecodeFailure           -> return $! SlowFail bs "expected int32"
+    DecodedToken sz (I# n#) ->
+      case (n# ># 0x7fffffff#) `orI#` (n# <# -0x80000000#) of
+        0#                  -> k n# >>= go_fast_end (BS.unsafeDrop sz bs)
+        _                   -> return $! SlowFail bs "expected int32"
+
+data SlowPath s a = SlowFail {-# UNPACK #-} !ByteString String
+
+data DecodeAction s a = ConsumeInt32 (Int# -> ST s (DecodeAction s a))
+
+data DecodedToken a = DecodedToken !Int !a | DecodeFailure
+
+tryConsumeInt :: Word8 -> ByteString -> DecodedToken Int
+tryConsumeInt hdr !bs = case word8ToWord hdr of
+  0x17 -> DecodedToken 1 23
+  0x1b -> case word64ToInt (eatTailWord64 bs) of
+            Just n  -> DecodedToken 9 n
+            Nothing -> DecodeFailure
+  _    -> DecodeFailure
+{-# INLINE tryConsumeInt #-}
+
+eatTailWord64 :: ByteString -> Word64
+eatTailWord64 xs = withBsPtr grabWord64 (BS.unsafeTail xs)
+{-# INLINE eatTailWord64 #-}
+
+word64ToInt :: Word64 -> Maybe Int
+#if __GLASGOW_HASKELL__ >= 904
+word64ToInt (W64# w#) =
+  case isTrue# (w# `ltWord64#` wordToWord64# 0x80000000##) of
+    True  -> Just (I# (int64ToInt# (word64ToInt64# w#)))
+    False -> Nothing
+#else
+word64ToInt (W64# w#) =
+  case isTrue# (w# `ltWord#` 0x8000000000000000##) of
+    True  -> Just (I# (word2Int# w#))
+    False -> Nothing
+#endif
+{-# INLINE word64ToInt #-}
+
+withBsPtr :: (Ptr b -> a) -> ByteString -> a
+withBsPtr f (BS.PS x off _) =
+    unsafeDupablePerformIO $ withForeignPtr x $
+        \(Ptr addr#) -> return $! (f (Ptr addr# `plusPtr` off))
+{-# INLINE withBsPtr #-}
+
+grabWord64 :: Ptr () -> Word64
+#if __GLASGOW_HASKELL__ >= 904
+grabWord64 (Ptr ip#) = W64# (byteSwap64# (indexWord64OffAddr# ip# 0#))
+#else
+grabWord64 (Ptr ip#) = W64# (byteSwap# (indexWord64OffAddr# ip# 0#))
+#endif
+{-# INLINE grabWord64 #-}
+
+word8ToWord :: Word8 -> Word
+#if MIN_VERSION_ghc_prim(0,8,0)
+word8ToWord (W8# w#) = W# (word8ToWord# w#)
+#else
+word8ToWord (W8# w#) = W# w#
+#endif
+{-# INLINE word8ToWord #-}
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index a636532f7d..fc0266819e 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -291,6 +291,7 @@ test('T16348', normal, compile, ['-O'])
 test('T16918', normal, compile, ['-O'])
 test('T16918a', normal, compile, ['-O'])
 test('T16978a', normal, compile, ['-O'])
+test('T21694', [ req_profiling ] , compile, ['-O -prof -fprof-auto -funfolding-use-threshold=50 '])
 test('T16978b', normal, compile, ['-O'])
 test('T16979a', normal, compile, ['-O'])
 test('T16979b', normal, compile, ['-O'])