Fix arity type of coerced types in CoreArity

Previously if we had

    f |> co

where `f` had arity type `ABot N` and `co` had arity M and M < N,
`arityType` would return `ABot M` which is wrong, because `f` is only
known to diverge when applied to `N` args, as described in Note
[ArityType]:

    If at = ABot n, then (f x1..xn) definitely diverges. Partial
    applications to fewer than n args may *or may not* diverge.

This caused incorrect eta expansion in the simplifier, causing #16066.

We now return `ATop M` for the same expression so the simplifier can't
assume partial applications of `f |> co` is divergent.

A regression test T16066 is also added.

4 files changed, 57 insertions, 2 deletions

diff --git a/compiler/coreSyn/CoreArity.hs b/compiler/coreSyn/CoreArity.hs
index 2f2418e38e..bb45773785 100644
--- a/compiler/coreSyn/CoreArity.hs
+++ b/compiler/coreSyn/CoreArity.hs
@@ -11,7 +11,7 @@
 -- | Arity and eta expansion
 module CoreArity (
         manifestArity, joinRhsArity, exprArity, typeArity,
-        exprEtaExpandArity, findRhsArity, CheapFun, etaExpand,
+        exprEtaExpandArity, findRhsArity, etaExpand,
         etaExpandToJoinPoint, etaExpandToJoinPointRule,
         exprBotStrictness_maybe
     ) where
@@ -702,6 +702,28 @@ lambda wasn't one-shot we don't want to do this.
 
 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.
+
+Note [Arity trimming]
+~~~~~~~~~~~~~~~~~~~~~
+Consider ((\x y. blah) |> co), where co :: (Int->Int->Int) ~ (Int -> F a) , and
+F is some type family.
+
+Because of Note [exprArity invariant], item (2), we must return with arity at
+most 1, because typeArity (Int -> F a) = 1.  So we have to trim the result of
+calling arityType on (\x y. blah).  Failing to do so, and hence breaking the
+exprArity invariant, led to #5441.
+
+How to trim?  For ATop, it's easy.  But we must take great care with ABot.
+Suppose the expression was (\x y. error "urk"), we'll get (ABot 2).  We
+absolutely must not trim that to (ABot 1), because that claims that
+((\x y. error "urk") |> co) diverges when given one argument, which it
+absolutely does not. And Bad Things happen if we think something returns bottom
+when it doesn't (#16066).
+
+So, do not reduce the 'n' in (ABot n); rather, switch (conservatively) to ATop.
+
+Historical note: long ago, we unconditionally switched to ATop when we
+encountered a cast, but that is far too conservative: see #5475
 -}
 
 ---------------------------
@@ -720,7 +742,9 @@ arityType :: ArityEnv -> CoreExpr -> ArityType
 arityType env (Cast e co)
   = case arityType env e of
       ATop os -> ATop (take co_arity os)
-      ABot n  -> ABot (n `min` co_arity)
+      -- See Note [Arity trimming]
+      ABot n | co_arity < n -> ATop (replicate co_arity noOneShotInfo)
+             | otherwise    -> ABot n
   where
     co_arity = length (typeArity (pSnd (coercionKind co)))
     -- See Note [exprArity invariant] (2); must be true of
diff --git a/testsuite/tests/simplCore/should_run/T16066.hs b/testsuite/tests/simplCore/should_run/T16066.hs
new file mode 100644
index 0000000000..3a1f44084e
--- /dev/null
+++ b/testsuite/tests/simplCore/should_run/T16066.hs
@@ -0,0 +1,27 @@
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE TypeFamilies #-}
+
+module Main (main) where
+
+import Control.Monad (join)
+import Control.Monad.Reader (ReaderT(..))
+import Control.Concurrent.STM (STM, atomically)
+import Data.Kind (Type)
+
+class Monad (Transaction m) => MonadPersist m where
+  type Transaction m :: Type -> Type
+  atomicTransaction :: Transaction m y -> m y
+
+instance MonadPersist (ReaderT () IO) where
+  type Transaction (ReaderT () IO) = ReaderT () STM
+  atomicTransaction act = ReaderT (atomically . runReaderT act)
+
+main :: IO ()
+main = join (runReaderT doPure2 ()) >>= \x -> seq x (return ())
+
+doPure2 :: MonadPersist m => m (IO ())
+doPure2 = atomicTransaction $ do
+  () <- pure ()
+  () <- pure ()
+  error "exit never happens"
diff --git a/testsuite/tests/simplCore/should_run/T16066.stderr b/testsuite/tests/simplCore/should_run/T16066.stderr
new file mode 100644
index 0000000000..85cfbabec1
--- /dev/null
+++ b/testsuite/tests/simplCore/should_run/T16066.stderr
@@ -0,0 +1,3 @@
+T16066: exit never happens
+CallStack (from HasCallStack):
+  error, called at T16066.hs:31:3 in main:Main
diff --git a/testsuite/tests/simplCore/should_run/all.T b/testsuite/tests/simplCore/should_run/all.T
index 8896ad5ee3..9ec45e05b0 100644
--- a/testsuite/tests/simplCore/should_run/all.T
+++ b/testsuite/tests/simplCore/should_run/all.T
@@ -90,3 +90,4 @@ test('T15114', only_ways('optasm'), compile_and_run, [''])
 test('T15436', normal, compile_and_run, [''])
 test('T15840', normal, compile_and_run, [''])
 test('T15840a', normal, compile_and_run, [''])
+test('T16066', exit_code(1), compile_and_run, ['-O1'])