Improve strictness analysis for exceptions

Two things here:

* For exceptions-catching primops like catch#, we know
  that the main argument function will be called, so
  we can use strictApply1Dmd, rather than lazy

  Changes in primops.txt.pp

* When a 'case' scrutinises a I/O-performing primop,
  the Note [IO hack in the demand analyser] was
  throwing away all strictness from the code that
  followed.

  I found that this was causing quite a bit of unnecessary
  reboxing in the (heavily used) function
  GHC.IO.Handle.Internals.wantReadableHandle

  So this patch prevents the hack applying when the
  case scrutinises a primop.  See the revised
  Note [IO hack in the demand analyser]

Thse two things buy us quite a lot in programs that do a lot of IO.

        Program           Size    Allocs   Runtime   Elapsed  TotalMem
--------------------------------------------------------------------------------
            hpg          -0.4%     -2.9%     -0.9%     -1.0%     +0.0%
reverse-complem          -0.4%    -10.9%    +10.7%    +10.9%     +0.0%
         simple          -0.3%     -0.0%    +26.2%    +26.2%     +3.7%
         sphere          -0.3%     -6.3%      0.09      0.09     +0.0%
--------------------------------------------------------------------------------
            Min          -0.7%    -10.9%     -4.6%     -4.7%     -1.7%
            Max          -0.2%     +0.0%    +26.2%    +26.2%     +6.5%
 Geometric Mean          -0.4%     -0.3%     +2.1%     +2.1%     +0.1%

I think the increase in runtime for 'simple' is measurement error.

3 files changed, 85 insertions, 37 deletions

diff --git a/compiler/basicTypes/Demand.hs b/compiler/basicTypes/Demand.hs
index b942f4ecd5..bfb346efb3 100644
--- a/compiler/basicTypes/Demand.hs
+++ b/compiler/basicTypes/Demand.hs
@@ -15,7 +15,8 @@ module Demand (
         mkProdDmd, mkOnceUsedDmd, mkManyUsedDmd, mkHeadStrict, oneifyDmd,
         getUsage, toCleanDmd,
         absDmd, topDmd, botDmd, seqDmd,
-        lubDmd, bothDmd, apply1Dmd, apply2Dmd,
+        lubDmd, bothDmd,
+        lazyApply1Dmd, lazyApply2Dmd, strictApply1Dmd,
         isTopDmd, isBotDmd, isAbsDmd, isSeqDmd,
         peelUseCall, cleanUseDmd_maybe, strictenDmd, bothCleanDmd,
         addCaseBndrDmd,
@@ -522,10 +523,11 @@ mkJointDmds ss as = zipWithEqual "mkJointDmds" mkJointDmd ss as
 absDmd :: JointDmd
 absDmd = mkJointDmd Lazy Abs
 
-apply1Dmd, apply2Dmd :: Demand
+lazyApply1Dmd, lazyApply2Dmd, strictApply1Dmd :: Demand
 -- C1(U), C1(C1(U)) respectively
-apply1Dmd = JD { strd = Lazy, absd = Use Many (UCall One Used) }
-apply2Dmd = JD { strd = Lazy, absd = Use Many (UCall One (UCall One Used)) }
+strictApply1Dmd = JD { strd = Str (SCall HeadStr), absd = Use Many (UCall One Used) }
+lazyApply1Dmd   = JD { strd = Lazy, absd = Use Many (UCall One Used) }
+lazyApply2Dmd   = JD { strd = Lazy, absd = Use Many (UCall One (UCall One Used)) }
 
 topDmd :: JointDmd
 topDmd = mkJointDmd Lazy useTop
diff --git a/compiler/prelude/primops.txt.pp b/compiler/prelude/primops.txt.pp
index 72110fefc3..c29e9d825a 100644
--- a/compiler/prelude/primops.txt.pp
+++ b/compiler/prelude/primops.txt.pp
@@ -1919,6 +1919,19 @@ primop  CasMutVarOp "casMutVar#" GenPrimOp
 section "Exceptions"
 ------------------------------------------------------------------------
 
+-- Note [Strictness for mask/unmask/catch]
+-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+-- Consider this example, which comes from GHC.IO.Handle.Internals:
+--    wantReadableHandle3 f ma b st
+--      = case ... of
+--          DEFAULT -> case ma of MVar a -> ...
+--          0#      -> maskAsynchExceptions# (\st -> case ma of MVar a -> ...)
+-- The outer case just decides whether to mask exceptions, but we don't want
+-- thereby to hide the strictness in 'ma'!  Hence the use of strictApply1Dmd.
+--
+-- For catch, we know that the first branch will be evaluated, but not
+-- necessarily the second.  Hence strictApply1Dmd and lazyApply1Dmd
+
 primop  CatchOp "catch#" GenPrimOp
           (State# RealWorld -> (# State# RealWorld, a #) )
        -> (b -> State# RealWorld -> (# State# RealWorld, a #) )
@@ -1928,7 +1941,8 @@ primop  CatchOp "catch#" GenPrimOp
         -- Catch is actually strict in its first argument
         -- but we don't want to tell the strictness
         -- analyser about that, so that exceptions stay inside it.
-   strictness  = { \ _arity -> mkClosedStrictSig [apply1Dmd,apply2Dmd,topDmd] topRes }
+   strictness  = { \ _arity -> mkClosedStrictSig [strictApply1Dmd,lazyApply2Dmd,topDmd] topRes }
+                 -- See Note [Strictness for mask/unmask/catch]
    out_of_line = True
    has_side_effects = True
 
@@ -1965,7 +1979,8 @@ primop  MaskAsyncExceptionsOp "maskAsyncExceptions#" GenPrimOp
         (State# RealWorld -> (# State# RealWorld, a #))
      -> (State# RealWorld -> (# State# RealWorld, a #))
    with
-   strictness  = { \ _arity -> mkClosedStrictSig [apply1Dmd,topDmd] topRes }
+   strictness  = { \ _arity -> mkClosedStrictSig [strictApply1Dmd,topDmd] topRes }
+                 -- See Note [Strictness for mask/unmask/catch]
    out_of_line = True
    has_side_effects = True
 
@@ -1973,7 +1988,7 @@ primop  MaskUninterruptibleOp "maskUninterruptible#" GenPrimOp
         (State# RealWorld -> (# State# RealWorld, a #))
      -> (State# RealWorld -> (# State# RealWorld, a #))
    with
-   strictness  = { \ _arity -> mkClosedStrictSig [apply1Dmd,topDmd] topRes }
+   strictness  = { \ _arity -> mkClosedStrictSig [strictApply1Dmd,topDmd] topRes }
    out_of_line = True
    has_side_effects = True
 
@@ -1981,7 +1996,8 @@ primop  UnmaskAsyncExceptionsOp "unmaskAsyncExceptions#" GenPrimOp
         (State# RealWorld -> (# State# RealWorld, a #))
      -> (State# RealWorld -> (# State# RealWorld, a #))
    with
-   strictness  = { \ _arity -> mkClosedStrictSig [apply1Dmd,topDmd] topRes }
+   strictness  = { \ _arity -> mkClosedStrictSig [strictApply1Dmd,topDmd] topRes }
+                 -- See Note [Strictness for mask/unmask/catch]
    out_of_line = True
    has_side_effects = True
 
@@ -2001,7 +2017,8 @@ primop  AtomicallyOp "atomically#" GenPrimOp
       (State# RealWorld -> (# State# RealWorld, a #) )
    -> State# RealWorld -> (# State# RealWorld, a #)
    with
-   strictness  = { \ _arity -> mkClosedStrictSig [apply1Dmd,topDmd] topRes }
+   strictness  = { \ _arity -> mkClosedStrictSig [strictApply1Dmd,topDmd] topRes }
+                 -- See Note [Strictness for mask/unmask/catch]
    out_of_line = True
    has_side_effects = True
 
@@ -2027,7 +2044,8 @@ primop  CatchRetryOp "catchRetry#" GenPrimOp
    -> (State# RealWorld -> (# State# RealWorld, a #) )
    -> (State# RealWorld -> (# State# RealWorld, a #) )
    with
-   strictness  = { \ _arity -> mkClosedStrictSig [apply1Dmd,apply1Dmd,topDmd] topRes }
+   strictness  = { \ _arity -> mkClosedStrictSig [strictApply1Dmd,lazyApply1Dmd,topDmd] topRes }
+                 -- See Note [Strictness for mask/unmask/catch]
    out_of_line = True
    has_side_effects = True
 
@@ -2036,7 +2054,8 @@ primop  CatchSTMOp "catchSTM#" GenPrimOp
    -> (b -> State# RealWorld -> (# State# RealWorld, a #) )
    -> (State# RealWorld -> (# State# RealWorld, a #) )
    with
-   strictness  = { \ _arity -> mkClosedStrictSig [apply1Dmd,apply2Dmd,topDmd] topRes }
+   strictness  = { \ _arity -> mkClosedStrictSig [strictApply1Dmd,lazyApply2Dmd,topDmd] topRes }
+                 -- See Note [Strictness for mask/unmask/catch]
    out_of_line = True
    has_side_effects = True
 
diff --git a/compiler/stranal/DmdAnal.hs b/compiler/stranal/DmdAnal.hs
index 79dd492ce7..41d9abb921 100644
--- a/compiler/stranal/DmdAnal.hs
+++ b/compiler/stranal/DmdAnal.hs
@@ -220,8 +220,8 @@ dmdAnal' env dmd (Case scrut case_bndr ty [(DataAlt dc, bndrs, rhs)])
         (alt_ty1, dmds)          = findBndrsDmds env rhs_ty bndrs
         (alt_ty2, case_bndr_dmd) = findBndrDmd env False alt_ty1 case_bndr
         id_dmds                  = addCaseBndrDmd case_bndr_dmd dmds
-        alt_ty3 | io_hack_reqd dc bndrs = deferAfterIO alt_ty2
-                | otherwise             = alt_ty2
+        alt_ty3 | io_hack_reqd scrut dc bndrs = deferAfterIO alt_ty2
+                | otherwise                   = alt_ty2
 
         -- Compute demand on the scrutinee
         -- See Note [Demand on scrutinee of a product case]
@@ -292,29 +292,16 @@ dmdAnal' env dmd (Let (Rec pairs) body)
     body_ty2 `seq`
     (body_ty2,  Let (Rec pairs') body')
 
-io_hack_reqd :: DataCon -> [Var] -> Bool
--- Note [IO hack in the demand analyser]
---
--- There's a hack here for I/O operations.  Consider
---      case foo x s of { (# s, r #) -> y }
--- Is this strict in 'y'.  Normally yes, but what if 'foo' is an I/O
--- operation that simply terminates the program (not in an erroneous way)?
--- In that case we should not evaluate 'y' before the call to 'foo'.
--- Hackish solution: spot the IO-like situation and add a virtual branch,
--- as if we had
---      case foo x s of
---         (# s, r #) -> y
---         other      -> return ()
--- So the 'y' isn't necessarily going to be evaluated
---
--- A more complete example (Trac #148, #1592) where this shows up is:
---      do { let len = <expensive> ;
---         ; when (...) (exitWith ExitSuccess)
---         ; print len }
-io_hack_reqd con bndrs
+io_hack_reqd :: CoreExpr -> DataCon -> [Var] -> Bool
+-- See Note [IO hack in the demand analyser]
+io_hack_reqd scrut con bndrs
   | (bndr:_) <- bndrs
-  = con == unboxedPairDataCon &&
-    idType bndr `eqType` realWorldStatePrimTy
+  , con == unboxedPairDataCon
+  , idType bndr `eqType` realWorldStatePrimTy
+  , (fun, _) <- collectArgs scrut
+  = case fun of
+      Var f -> not (isPrimOpId f)
+      _     -> True
   | otherwise
   = False
 
@@ -350,8 +337,48 @@ dmdAnalAlt env dmd case_bndr (con,bndrs,rhs)
         id_dmds       = addCaseBndrDmd case_bndr_dmd dmds
   = (alt_ty, (con, setBndrsDemandInfo bndrs id_dmds, rhs'))
 
-{- Note [Demand on the scrutinee of a product case]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+{- Note [IO hack in the demand analyser]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+There's a hack here for I/O operations.  Consider
+     case foo x s of { (# s, r #) -> y }
+Is this strict in 'y'.  Normally yes, but what if 'foo' is an I/O
+operation that simply terminates the program (not in an erroneous way)?
+In that case we should not evaluate 'y' before the call to 'foo'.
+Hackish solution: spot the IO-like situation and add a virtual branch,
+as if we had
+     case foo x s of
+        (# s, r #) -> y
+        other      -> return ()
+So the 'y' isn't necessarily going to be evaluated
+
+A more complete example (Trac #148, #1592) where this shows up is:
+     do { let len = <expensive> ;
+        ; when (...) (exitWith ExitSuccess)
+        ; print len }
+
+However, consider
+  f x s = case getMaskingState# s of
+            (# s, r #) ->
+          case x of I# x2 -> ...
+
+Here it is terribly sad to make 'f' lazy in 's'.  After all,
+getMaskingState# is not going to diverge or throw an exception!  This
+situation actually arises in GHC.IO.Handle.Internals.wantReadableHandle
+(on an MVar not an Int), and make a material difference.
+
+So if the scrutinee is a primop call, we *don't* apply the
+state hack:
+  - If is a simple, terminating one like getMaskingState,
+    applying the hack is over-conservative.
+  - If the primop is raise# then it returns bottom, so
+    the case alternatives are alraedy discarded.
+  - If the primop can raise a non-IO exception, like
+    divide by zero or seg-fault (eg writing an array
+    out of bounds) then we don't mind evaluating 'x' first.
+
+Note [Demand on the scrutinee of a product case]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 When figuring out the demand on the scrutinee of a product case,
 we use the demands of the case alternative, i.e. id_dmds.
 But note that these include the demand on the case binder;