Improve boxity in deferAfterPreciseExceptionwip/T18993a

As #20746 showed, the demand analyser behaved badly in a key I/O
library (`GHC.IO.Handle.Text`), by unnessarily boxing and reboxing.

This patch adjusts the subtle function deferAfterPreciseException;
it's quite easy, just a bit subtle.

See the new Note [deferAfterPreciseException]

And this MR deals only with Problem 2 in #20746.
Problem 1 is still open.

6 files changed, 294 insertions, 5 deletions

diff --git a/compiler/GHC/Types/Demand.hs b/compiler/GHC/Types/Demand.hs
index 80c725bbfc..65f3239a9e 100644
--- a/compiler/GHC/Types/Demand.hs
+++ b/compiler/GHC/Types/Demand.hs
@@ -607,13 +607,17 @@ multCard (Card a) (Card b)
 -- isn't any evaluation at all. If you don't care, simply use '(:*)'.
 data Demand
   = BotDmd
-  -- ^ A bottoming demand, produced by a diverging function, hence there is no
+  -- ^ A bottoming demand, produced by a diverging function ('C_10'), hence there is no
   -- 'SubDemand' that describes how it was evaluated.
+
   | AbsDmd
   -- ^ An absent demand: Evaluated exactly 0 times ('C_00'), hence there is no
   -- 'SubDemand' that describes how it was evaluated.
+
   | D !CardNonAbs !SubDemand
   -- ^ Don't use this internal data constructor; use '(:*)' instead.
+  -- Since BotDmd deals with 'C_10' and AbsDmd deals with 'C_00', the
+  -- cardinality component is CardNonAbs
   deriving Eq
 
 -- | Only meant to be used in the pattern synonym below!
@@ -1561,10 +1565,12 @@ isTopDmdType :: DmdType -> Bool
 isTopDmdType (DmdType env args div)
   = div == topDiv && null args && isEmptyVarEnv env
 
+{-   Unused
 -- | The demand type of an unspecified expression that is guaranteed to
 -- throw a (precise or imprecise) exception or diverge.
 exnDmdType :: DmdType
 exnDmdType = DmdType emptyDmdEnv [] exnDiv
+-}
 
 dmdTypeDepth :: DmdType -> Arity
 dmdTypeDepth = length . dt_args
@@ -1626,21 +1632,66 @@ findIdDemand (DmdType fv _ res) id
 -- exception, we act as if there is an additional control flow path that is
 -- taken if e throws a precise exception. The demand type of this control flow
 -- path
---   * is lazy and absent ('topDmd') in all free variables and arguments
+--   * is lazy and absent ('topDmd') and boxed in all free variables and arguments
 --   * has 'exnDiv' 'Divergence' result
+-- See Note [Precise exceptions and strictness analysis]
+--
 -- So we can simply take a variant of 'nopDmdType', 'exnDmdType'.
 -- Why not 'nopDmdType'? Because then the result of 'e' can never be 'exnDiv'!
 -- That means failure to drop dead-ends, see #18086.
--- See Note [Precise exceptions and strictness analysis]
 deferAfterPreciseException :: DmdType -> DmdType
-deferAfterPreciseException = lubDmdType exnDmdType
+-- deferAfterPreciseException = lubDmdType exnDmdType
+deferAfterPreciseException (DmdType fvs ds r)
+  = DmdType (mapVarEnv defer fvs)
+            (map defer ds)
+            (r `lubDivergence` ExnOrDiv)
+  where
+    defer :: Demand -> Demand
+    defer AbsDmd = AbsDmd
+    defer BotDmd = AbsDmd
+    defer (D n sd) = lubCard n C_00 :* lubSubDmd sd (Poly Boxed C_00)
+
+    -- Roughly:  defer d = d `lubDmd` D C_00 (Poly Boxed C_00)
+    -- It is very important that we `lub` with `Boxed`; see
+    --     Note [deferAfterPreciseException]
+    -- But that formulation fails the assert in :*,
+    --     because (D C_00 (Poly Boxed C_00)) is not a legal demand
+    -- So we write defer out more explicitly here
 
 -- | See 'keepAliveDmdEnv'.
 keepAliveDmdType :: DmdType -> VarSet -> DmdType
 keepAliveDmdType (DmdType fvs ds res) vars =
   DmdType (fvs `keepAliveDmdEnv` vars) ds res
 
-{-
+{- Note [deferAfterPreciseException]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+The big picture is in Note [Precise exceptions and strictness analysis]
+The idea is that we want to treat
+   case <I/O operation> of (# s', r #) -> rhs
+
+as if it was
+   case <I/O operation> of
+      Just (# s', r #) -> rhs
+      Nothing          -> error
+
+That is, the I/O operation might throw an exception, so that 'rhs' never
+gets reached.  For example, we don't want to be strict in the strict free
+variables of 'rhs'.
+
+So roughly speaking:
+  deferAfterPreciseException = lubDmdType (DmdType emptyDmdEnv [] exnDiv)
+
+But that doesn't work quite right for boxity becasuse
+   case <I/O operation> of
+      (# s', r) -> f x
+
+uses `x` *boxed*.  If we `lub` it with `(DmdType emptyDmdEnv [] exnDiv)`
+we'll get an *unboxed* demand on `x`, which led to #20746.  There is
+a fuller example in that ticket.
+
+TL;DR: deferAfterPreciseException is very careful to preserve boxity
+in its argument.
+
 Note [Demand type Divergence]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 In contrast to DmdSigs, DmdTypes are elicited under a specific incoming demand.
diff --git a/testsuite/tests/stranal/should_compile/T20746.hs b/testsuite/tests/stranal/should_compile/T20746.hs
new file mode 100644
index 0000000000..93496acd65
--- /dev/null
+++ b/testsuite/tests/stranal/should_compile/T20746.hs
@@ -0,0 +1,11 @@
+module Foo where
+
+f x = (g, g)
+  where
+    g :: Int -> IO Int
+    g y = do { if y>2 then print x else return ()
+             ; foogle y }
+
+foogle :: Int -> IO Int
+{-# NOINLINE foogle #-}
+foogle n = return n
diff --git a/testsuite/tests/stranal/should_compile/T20746.stderr b/testsuite/tests/stranal/should_compile/T20746.stderr
new file mode 100644
index 0000000000..6e7f56f625
--- /dev/null
+++ b/testsuite/tests/stranal/should_compile/T20746.stderr
@@ -0,0 +1,133 @@
+
+==================== Tidy Core ====================
+Result size of Tidy Core
+  = {terms: 70, types: 113, coercions: 18, joins: 0/2}
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+Foo.$trModule4 :: GHC.Prim.Addr#
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 20 0}]
+Foo.$trModule4 = "main"#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+Foo.$trModule3 :: GHC.Types.TrName
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 10 10}]
+Foo.$trModule3 = GHC.Types.TrNameS Foo.$trModule4
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+Foo.$trModule2 :: GHC.Prim.Addr#
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 20 0}]
+Foo.$trModule2 = "Foo"#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+Foo.$trModule1 :: GHC.Types.TrName
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 10 10}]
+Foo.$trModule1 = GHC.Types.TrNameS Foo.$trModule2
+
+-- RHS size: {terms: 3, types: 0, coercions: 0, joins: 0/0}
+Foo.$trModule :: GHC.Types.Module
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 10 10}]
+Foo.$trModule = GHC.Types.Module Foo.$trModule3 Foo.$trModule1
+
+-- RHS size: {terms: 5, types: 8, coercions: 0, joins: 0/0}
+Foo.f1 [InlPrag=NOINLINE]
+  :: Int
+     -> GHC.Prim.State# GHC.Prim.RealWorld
+     -> (# GHC.Prim.State# GHC.Prim.RealWorld, Int #)
+[GblId, Arity=2, Str=<L><L>, Cpr=1, Unf=OtherCon []]
+Foo.f1
+  = \ (n :: Int)
+      (s [OS=OneShot] :: GHC.Prim.State# GHC.Prim.RealWorld) ->
+      (# s, n #)
+
+-- RHS size: {terms: 1, types: 0, coercions: 6, joins: 0/0}
+foogle [InlPrag=[final]] :: Int -> IO Int
+[GblId,
+ Arity=2,
+ Str=<L><L>,
+ Cpr=1,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True,
+         Guidance=ALWAYS_IF(arity=0,unsat_ok=True,boring_ok=True)}]
+foogle
+  = Foo.f1
+    `cast` (<Int>_R %<'Many>_N ->_R Sym (GHC.Types.N:IO[0] <Int>_R)
+            :: (Int
+                -> GHC.Prim.State# GHC.Prim.RealWorld
+                -> (# GHC.Prim.State# GHC.Prim.RealWorld, Int #))
+               ~R# (Int -> IO Int))
+
+-- RHS size: {terms: 35, types: 38, coercions: 12, joins: 0/2}
+Foo.$wf [InlPrag=[2]]
+  :: forall {a}. Show a => a -> (# Int -> IO Int, Int -> IO Int #)
+[GblId,
+ Arity=2,
+ Str=<MP(A,MCM(L),A)><L>,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [30 0] 181 10}]
+Foo.$wf
+  = \ (@a) ($dShow :: Show a) (x :: a) ->
+      let {
+        lvl :: String
+        [LclId]
+        lvl = show @a $dShow x } in
+      let {
+        g :: Int
+             -> GHC.Prim.State# GHC.Prim.RealWorld
+             -> (# GHC.Prim.State# GHC.Prim.RealWorld, Int #)
+        [LclId, Arity=2, Str=<1L><L>, Unf=OtherCon []]
+        g = \ (y :: Int) (s :: GHC.Prim.State# GHC.Prim.RealWorld) ->
+              case y of wild { GHC.Types.I# x1 ->
+              case GHC.Prim.># x1 2# of {
+                __DEFAULT -> Foo.f1 wild s;
+                1# ->
+                  case GHC.IO.Handle.Text.hPutStr2
+                         GHC.IO.Handle.FD.stdout lvl GHC.Types.True s
+                  of
+                  { (# ipv, ipv1 #) ->
+                  Foo.f1 wild ipv
+                  }
+              }
+              } } in
+      (# g
+         `cast` (<Int>_R %<'Many>_N ->_R Sym (GHC.Types.N:IO[0] <Int>_R)
+                 :: (Int
+                     -> GHC.Prim.State# GHC.Prim.RealWorld
+                     -> (# GHC.Prim.State# GHC.Prim.RealWorld, Int #))
+                    ~R# (Int -> IO Int)),
+         g
+         `cast` (<Int>_R %<'Many>_N ->_R Sym (GHC.Types.N:IO[0] <Int>_R)
+                 :: (Int
+                     -> GHC.Prim.State# GHC.Prim.RealWorld
+                     -> (# GHC.Prim.State# GHC.Prim.RealWorld, Int #))
+                    ~R# (Int -> IO Int)) #)
+
+-- RHS size: {terms: 11, types: 26, coercions: 0, joins: 0/0}
+f [InlPrag=[2]]
+  :: forall {a}. Show a => a -> (Int -> IO Int, Int -> IO Int)
+[GblId,
+ Arity=2,
+ Str=<MP(A,MCM(L),A)><L>,
+ Cpr=1,
+ Unf=Unf{Src=InlineStable, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True,
+         Guidance=ALWAYS_IF(arity=2,unsat_ok=True,boring_ok=False)
+         Tmpl= \ (@a) ($dShow [Occ=Once1] :: Show a) (x [Occ=Once1] :: a) ->
+                 case Foo.$wf @a $dShow x of
+                 { (# ww [Occ=Once1], ww1 [Occ=Once1] #) ->
+                 (ww, ww1)
+                 }}]
+f = \ (@a) ($dShow :: Show a) (x :: a) ->
+      case Foo.$wf @a $dShow x of { (# ww, ww1 #) -> (ww, ww1) }
+
+
+
diff --git a/testsuite/tests/stranal/should_compile/T20746b.hs b/testsuite/tests/stranal/should_compile/T20746b.hs
new file mode 100644
index 0000000000..6804fb4449
--- /dev/null
+++ b/testsuite/tests/stranal/should_compile/T20746b.hs
@@ -0,0 +1,14 @@
+module T20746b where
+
+mightThrow :: Int -> IO Int
+{-# NOINLINE mightThrow #-}
+mightThrow n = return n
+
+-- Should not unbox trp
+-- Recursive because if it's too small
+-- we don't do worker/wrapper at all
+f :: Bool -> (Int, Int, Int) -> IO (Int, Int, Int)
+f False trp = f True trp
+f True trp@(a,b,c) = do
+  _ <- mightThrow a -- this potentially throwing IO action should not force unboxing of trp
+  return trp
diff --git a/testsuite/tests/stranal/should_compile/T20746b.stderr b/testsuite/tests/stranal/should_compile/T20746b.stderr
new file mode 100644
index 0000000000..97f8496c4b
--- /dev/null
+++ b/testsuite/tests/stranal/should_compile/T20746b.stderr
@@ -0,0 +1,77 @@
+
+==================== Tidy Core ====================
+Result size of Tidy Core
+  = {terms: 33, types: 78, coercions: 21, joins: 0/0}
+
+-- RHS size: {terms: 5, types: 8, coercions: 0, joins: 0/0}
+T20746b.mightThrow1 [InlPrag=NOINLINE]
+  :: Int
+     -> GHC.Prim.State# GHC.Prim.RealWorld
+     -> (# GHC.Prim.State# GHC.Prim.RealWorld, Int #)
+[GblId, Arity=2, Str=<L><L>, Cpr=1, Unf=OtherCon []]
+T20746b.mightThrow1
+  = \ (n :: Int)
+      (s [OS=OneShot] :: GHC.Prim.State# GHC.Prim.RealWorld) ->
+      (# s, n #)
+
+-- RHS size: {terms: 1, types: 0, coercions: 6, joins: 0/0}
+mightThrow [InlPrag=[final]] :: Int -> IO Int
+[GblId,
+ Arity=2,
+ Str=<L><L>,
+ Cpr=1,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True,
+         Guidance=ALWAYS_IF(arity=0,unsat_ok=True,boring_ok=True)}]
+mightThrow
+  = T20746b.mightThrow1
+    `cast` (<Int>_R %<'Many>_N ->_R Sym (GHC.Types.N:IO[0] <Int>_R)
+            :: (Int
+                -> GHC.Prim.State# GHC.Prim.RealWorld
+                -> (# GHC.Prim.State# GHC.Prim.RealWorld, Int #))
+               ~R# (Int -> IO Int))
+
+Rec {
+-- RHS size: {terms: 22, types: 32, coercions: 0, joins: 0/0}
+T20746b.f1 [Occ=LoopBreaker]
+  :: Bool
+     -> (Int, Int, Int)
+     -> GHC.Prim.State# GHC.Prim.RealWorld
+     -> (# GHC.Prim.State# GHC.Prim.RealWorld, (Int, Int, Int) #)
+[GblId, Arity=3, Str=<1L><1L><L>, Cpr=1, Unf=OtherCon []]
+T20746b.f1
+  = \ (ds :: Bool)
+      (trp :: (Int, Int, Int))
+      (eta [OS=OneShot] :: GHC.Prim.State# GHC.Prim.RealWorld) ->
+      case ds of {
+        False -> T20746b.f1 GHC.Types.True trp eta;
+        True ->
+          case trp of wild1 { (a, b, c) ->
+          case T20746b.mightThrow1 a eta of { (# ipv, ipv1 #) ->
+          (# ipv, wild1 #)
+          }
+          }
+      }
+end Rec }
+
+-- RHS size: {terms: 1, types: 0, coercions: 15, joins: 0/0}
+f :: Bool -> (Int, Int, Int) -> IO (Int, Int, Int)
+[GblId,
+ Arity=3,
+ Str=<1L><1L><L>,
+ Cpr=1,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True,
+         Guidance=ALWAYS_IF(arity=0,unsat_ok=True,boring_ok=True)}]
+f = T20746b.f1
+    `cast` (<Bool>_R
+            %<'Many>_N ->_R <(Int, Int, Int)>_R
+            %<'Many>_N ->_R Sym (GHC.Types.N:IO[0] <(Int, Int, Int)>_R)
+            :: (Bool
+                -> (Int, Int, Int)
+                -> GHC.Prim.State# GHC.Prim.RealWorld
+                -> (# GHC.Prim.State# GHC.Prim.RealWorld, (Int, Int, Int) #))
+               ~R# (Bool -> (Int, Int, Int) -> IO (Int, Int, Int)))
+
+
+
diff --git a/testsuite/tests/stranal/should_compile/all.T b/testsuite/tests/stranal/should_compile/all.T
index 7a8d90c2d8..1723436f8d 100644
--- a/testsuite/tests/stranal/should_compile/all.T
+++ b/testsuite/tests/stranal/should_compile/all.T
@@ -77,3 +77,6 @@ test('T19882b', normal, compile, [''])
 # We want that the 'go' joinrec in the unfolding has been worker/wrappered.
 # So we simply grep for 'jump $wgo' and hope we find more than 2 call sites:
 test('T20510', [ grep_errmsg(r'jump \$wgo') ], compile, ['-dsuppress-uniques -ddump-exitify'])
+
+test('T20746', normal, compile, ['-dsuppress-uniques -ddump-simpl'])
+test('T20746b', normal, compile, ['-dsuppress-uniques -ddump-simpl -dno-typeable-binds'])