Fix a long-standing bug in the demand analyser

This patch fixes Trac #10148, an outright and egregious
bug in the demand analyser.

It is explained in Note [Demand on case-alternative binders]
in Demand.hs.

I did some other minor refactoring.

To my astonishment I got some big compiler perf changes

* perf/compiler/T5837: bytes allocated -76%
* perf/compiler/T5030: bytes allocated -10%
* perf/compiler/T3294: max bytes used  -25%

Happy days

(cherry picked from commit 9f0f99fd41ff82cc223d3b682703e508efb564d2)

2 files changed, 165 insertions, 87 deletions

diff --git a/compiler/basicTypes/Demand.hs b/compiler/basicTypes/Demand.hs
index ecf22bc51f..25dcd130a2 100644
--- a/compiler/basicTypes/Demand.hs
+++ b/compiler/basicTypes/Demand.hs
@@ -18,6 +18,7 @@ module Demand (
         lubDmd, bothDmd, apply1Dmd, apply2Dmd,
         isTopDmd, isBotDmd, isAbsDmd, isSeqDmd,
         peelUseCall, cleanUseDmd_maybe, strictenDmd, bothCleanDmd,
+        addCaseBndrDmd,
 
         DmdType(..), dmdTypeDepth, lubDmdType, bothDmdType,
         nopDmdType, botDmdType, mkDmdType,
@@ -25,7 +26,7 @@ module Demand (
         BothDmdArg, mkBothDmdArg, toBothDmdArg,
 
         DmdEnv, emptyDmdEnv,
-        peelFV,
+        peelFV, findIdDemand,
 
         DmdResult, CPRResult,
         isBotRes, isTopRes,
@@ -200,6 +201,10 @@ seqMaybeStr Lazy    = ()
 seqMaybeStr (Str s) = seqStrDmd s
 
 -- Splitting polymorphic demands
+splitMaybeStrProdDmd :: Int -> MaybeStr -> Maybe [MaybeStr]
+splitMaybeStrProdDmd n Lazy    = Just (replicate n Lazy)
+splitMaybeStrProdDmd n (Str s) = splitStrProdDmd n s
+
 splitStrProdDmd :: Int -> StrDmd -> Maybe [MaybeStr]
 splitStrProdDmd n HyperStr   = Just (replicate n strBot)
 splitStrProdDmd n HeadStr    = Just (replicate n strTop)
@@ -352,7 +357,49 @@ peelUseCall :: UseDmd -> Maybe (Count, UseDmd)
 peelUseCall (UCall c u)   = Just (c,u)
 peelUseCall _             = Nothing
 
-{-
+addCaseBndrDmd :: Demand    -- On the case binder
+               -> [Demand]  -- On the components of the constructor
+               -> [Demand]  -- Final demands for the components of the constructor
+-- See Note [Demand on case-alternative binders]
+addCaseBndrDmd (JD { strd = ms, absd = mu }) alt_dmds
+  = case mu of
+     Abs     -> alt_dmds
+     Use _ u -> zipWith bothDmd alt_dmds (mkJointDmds ss us)
+             where
+                Just ss = splitMaybeStrProdDmd arity ms  -- Guaranteed not to be a call
+                Just us = splitUseProdDmd      arity u   -- Ditto
+  where
+    arity = length alt_dmds
+
+{- Note [Demand on case-alternative binders]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+The demand on a binder in a case alternative comes
+  (a) From the demand on the binder itself
+  (b) From the demand on the case binder
+Forgetting (b) led directly to Trac #10148.
+
+Example. Source code:
+  f x@(p,_) = if p then foo x else True
+
+  foo (p,True) = True
+  foo (p,q)    = foo (q,p)
+
+After strictness analysis:
+  f = \ (x_an1 [Dmd=<S(SL),1*U(U,1*U)>] :: (Bool, Bool)) ->
+      case x_an1
+      of wild_X7 [Dmd=<L,1*U(1*U,1*U)>]
+      { (p_an2 [Dmd=<S,1*U>], ds_dnz [Dmd=<L,A>]) ->
+      case p_an2 of _ {
+        False -> GHC.Types.True;
+        True -> foo wild_X7 }
+
+It's true that ds_dnz is *itself* absent, b ut the use of wild_X7 means
+that it is very much alive and demanded.  See Trac #10148 for how the
+consequences play out.
+
+This is needed even for non-product types, in case the case-binder
+is used but the components of the case alternative are not.
+
 Note [Don't optimise UProd(Used) to Used]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 These two UseDmds:
@@ -586,7 +633,8 @@ f g = (snd (g 3), True)
 should be: <L,C(U(AU))>m
 -}
 
-data CleanDemand = CD { sd :: StrDmd, ud :: UseDmd }
+data CleanDemand   -- A demand that is at least head-strict
+  = CD { sd :: StrDmd, ud :: UseDmd }
   deriving ( Eq, Show )
 
 instance Outputable CleanDemand where
@@ -1339,6 +1387,10 @@ peelFV (DmdType fv ds res) id = -- pprTrace "rfv" (ppr id <+> ppr dmd $$ ppr fv)
 addDemand :: Demand -> DmdType -> DmdType
 addDemand dmd (DmdType fv ds res) = DmdType fv (dmd:ds) res
 
+findIdDemand :: DmdType -> Var -> Demand
+findIdDemand (DmdType fv _ res) id
+  = lookupVarEnv fv id `orElse` defaultDmd res
+
 {-
 Note [Default demand on free variables]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/compiler/stranal/DmdAnal.hs b/compiler/stranal/DmdAnal.hs
index 9d9af64a7e..65e2b0c944 100644
--- a/compiler/stranal/DmdAnal.hs
+++ b/compiler/stranal/DmdAnal.hs
@@ -208,19 +208,16 @@ dmdAnal' env dmd (Lam var body)
     in
     (postProcessUnsat defer_and_use lam_ty, Lam var' body')
 
-dmdAnal' env dmd (Case scrut case_bndr ty [alt@(DataAlt dc, _, _)])
+dmdAnal' env dmd (Case scrut case_bndr ty [(DataAlt dc, bndrs, rhs)])
   -- Only one alternative with a product constructor
   | let tycon = dataConTyCon dc
   , isProductTyCon tycon
   , Just rec_tc' <- checkRecTc (ae_rec_tc env) tycon
   = let
-        env_w_tc              = env { ae_rec_tc = rec_tc' }
-        env_alt               = extendAnalEnv NotTopLevel env_w_tc case_bndr case_bndr_sig
-        (alt_ty, alt')        = dmdAnalAlt env_alt dmd alt
-        (alt_ty1, case_bndr') = annotateBndr env alt_ty case_bndr
-        (_, bndrs', _)        = alt'
-        case_bndr_sig         = cprProdSig (dataConRepArity dc)
-                -- Inside the alternative, the case binder has the CPR property.
+        env_w_tc      = env { ae_rec_tc = rec_tc' }
+        env_alt       = extendAnalEnv NotTopLevel env_w_tc case_bndr case_bndr_sig
+        case_bndr_sig = cprProdSig (dataConRepArity dc)
+                -- cprProdSig: inside the alternative, the case binder has the CPR property.
                 -- Meaning that a case on it will successfully cancel.
                 -- Example:
                 --      f True  x = case x of y { I# x' -> if x' ==# 3 then y else I# 8 }
@@ -231,44 +228,33 @@ dmdAnal' env dmd (Case scrut case_bndr ty [alt@(DataAlt dc, _, _)])
                 --      fw True  x = case x of y { I# x' -> if x' ==# 3 then x' else 8 }
                 --      fw False x = 3
 
-        -- Figure out whether the demand on the case binder is used, and use
-        -- that to set the scrut_dmd.  This is utterly essential.
-        -- Consider     f x = case x of y { (a,b) -> k y a }
-        -- If we just take scrut_demand = U(L,A), then we won't pass x to the
-        -- worker, so the worker will rebuild
-        --      x = (a, absent-error)
-        -- and that'll crash.
-        -- So at one stage I had:
-        --      dead_case_bndr           = isAbsDmd (idDemandInfo case_bndr')
-        --      keepity | dead_case_bndr = Drop
-        --              | otherwise      = Keep
-        --
-        -- But then consider
-        --      case x of y { (a,b) -> h y + a }
-        -- where h : U(LL) -> T
-        -- The above code would compute a Keep for x, since y is not Abs, which is silly
-        -- The insight is, of course, that a demand on y is a demand on the
-        -- scrutinee, so we need to `both` it with the scrut demand
-
-        scrut_dmd1 = mkProdDmd [idDemandInfo b | b <- bndrs', isId b]
-        scrut_dmd2 = strictenDmd (idDemandInfo case_bndr')
-        scrut_dmd  = scrut_dmd1 `bothCleanDmd` scrut_dmd2
+        (rhs_ty, rhs')                  = dmdAnal env_alt dmd 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
 
+        -- Compute demand on the scrutinee
+        -- See Note [Demand on scrutinee of a product case]
+        scrut_dmd          = mkProdDmd (addDataConStrictness dc id_dmds)
         (scrut_ty, scrut') = dmdAnal env scrut_dmd scrut
-        res_ty             = alt_ty1 `bothDmdType` toBothDmdArg scrut_ty
+        res_ty             = alt_ty3 `bothDmdType` toBothDmdArg scrut_ty
+        case_bndr'         = setIdDemandInfo case_bndr case_bndr_dmd
+        bndrs'             = setBndrsDemandInfo bndrs id_dmds
     in
 --    pprTrace "dmdAnal:Case1" (vcat [ text "scrut" <+> ppr scrut
 --                                   , text "dmd" <+> ppr dmd
 --                                   , text "case_bndr_dmd" <+> ppr (idDemandInfo case_bndr')
 --                                   , text "scrut_dmd" <+> ppr scrut_dmd
 --                                   , text "scrut_ty" <+> ppr scrut_ty
---                                   , text "alt_ty" <+> ppr alt_ty1
+--                                   , text "alt_ty" <+> ppr alt_ty2
 --                                   , text "res_ty" <+> ppr res_ty ]) $
-    (res_ty, Case scrut' case_bndr' ty [alt'])
+    (res_ty, Case scrut' case_bndr' ty [(DataAlt dc, bndrs', rhs')])
 
 dmdAnal' env dmd (Case scrut case_bndr ty alts)
   = let      -- Case expression with multiple alternatives
-        (alt_tys, alts')     = mapAndUnzip (dmdAnalAlt env dmd) alts
+        (alt_tys, alts')     = mapAndUnzip (dmdAnalAlt env dmd case_bndr) alts
         (scrut_ty, scrut')   = dmdAnal env cleanEvalDmd scrut
         (alt_ty, case_bndr') = annotateBndr env (foldr lubDmdType botDmdType alt_tys) case_bndr
         res_ty               = alt_ty `bothDmdType` toBothDmdArg scrut_ty
@@ -315,6 +301,32 @@ 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
+  | (bndr:_) <- bndrs
+  = con == unboxedPairDataCon &&
+    idType bndr `eqType` realWorldStatePrimTy
+  | otherwise
+  = False
+
 annLamWithShotness :: Demand -> CoreExpr -> CoreExpr
 annLamWithShotness d e
   | Just u <- cleanUseDmd_maybe d
@@ -334,40 +346,32 @@ setOneShotness :: Count -> Id -> Id
 setOneShotness One  bndr = setOneShotLambda bndr
 setOneShotness Many bndr = bndr
 
-dmdAnalAlt :: AnalEnv -> CleanDemand -> Alt Var -> (DmdType, Alt Var)
-dmdAnalAlt env dmd (con,bndrs,rhs)
-  = let
-        (rhs_ty, rhs')   = dmdAnal env dmd rhs
-        rhs_ty'          = addDataConPatDmds con bndrs rhs_ty
-        (alt_ty, bndrs') = annotateBndrs env rhs_ty' bndrs
-        final_alt_ty | io_hack_reqd = deferAfterIO alt_ty
-                     | otherwise    = alt_ty
+dmdAnalAlt :: AnalEnv -> CleanDemand -> Id -> Alt Var -> (DmdType, Alt Var)
+dmdAnalAlt env dmd case_bndr (con,bndrs,rhs)
+  | null bndrs    -- Literals, DEFAULT, and nullary constructors
+  , (rhs_ty, rhs') <- dmdAnal env dmd rhs
+  = (rhs_ty, (con, [], rhs'))
+
+  | otherwise     -- Non-nullary data constructors
+  , (rhs_ty, rhs') <- dmdAnal env dmd rhs
+  , (alt_ty, dmds) <- findBndrsDmds env rhs_ty bndrs
+  , let case_bndr_dmd = findIdDemand alt_ty case_bndr
+        id_dmds       = addCaseBndrDmd case_bndr_dmd dmds
+  = (alt_ty, (con, setBndrsDemandInfo bndrs id_dmds, rhs'))
+
+{- 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;
+see Note [Demand on case-alternative binders] in Demand.hs.
+This is crucial. Example:
+   f x = case x of y { (a,b) -> k y a }
+If we just take scrut_demand = U(L,A), then we won't pass x to the
+worker, so the worker will rebuild
+     x = (a, absent-error)
+and that'll crash.
 
-        -- 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 == DataAlt unboxedPairDataCon &&
-                       idType (head bndrs) `eqType` realWorldStatePrimTy
-    in
-    (final_alt_ty, (con, bndrs', rhs'))
-
-{-
 Note [Aggregated demand for cardinality]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 We use different strategies for strictness and usage/cardinality to
@@ -426,18 +430,6 @@ In other words, for locally-bound lambdas we can infer
 one-shotness.
 -}
 
-addDataConPatDmds :: AltCon -> [Var] -> DmdType -> DmdType
--- See Note [Add demands for strict constructors]
-addDataConPatDmds DEFAULT    _ dmd_ty = dmd_ty
-addDataConPatDmds (LitAlt _) _ dmd_ty = dmd_ty
-addDataConPatDmds (DataAlt con) bndrs dmd_ty
-  = foldr add dmd_ty str_bndrs
-  where
-    add bndr dmd_ty = addVarDmd dmd_ty bndr seqDmd
-    str_bndrs = [ b | (b,s) <- zipEqual "addDataConPatBndrs"
-                                   (filter isId bndrs)
-                                   (dataConRepStrictness con)
-                    , isMarkedStrict s ]
 
 {-
 Note [Add demands for strict constructors]
@@ -457,8 +449,8 @@ We want the worker for 'foo' too look like this:
     $wfoo :: Int# -> Int# -> Int#
 
 with the first argument unboxed, so that it is not eval'd each time
-around the loop (which would otherwise happen, since 'foo' is not
-strict in 'a'.  It is sound for the wrapper to pass an unboxed arg
+around the 'go' loop (which would otherwise happen, since 'foo' is not
+strict in 'a').  It is sound for the wrapper to pass an unboxed arg
 because X is strict, so its argument must be evaluated.  And if we
 *don't* pass an unboxed argument, we can't even repair it by adding a
 `seq` thus:
@@ -472,6 +464,13 @@ if X is monomorphic, and has an UNPACK pragma, then this optimisation
 is even more important.  We don't want the wrapper to rebox an unboxed
 argument, and pass an Int to $wfoo!
 
+We add these extra strict demands to the demand on the *scrutinee* of
+the case expression; hence the use of addDataConStrictness when
+forming scrut_dmd.  The case alternatives aren't strict in their
+sub-components, but simply evaluating the scrutinee to HNF does force
+those sub-components.
+
+
 ************************************************************************
 *                                                                      *
                     Demand transformer
@@ -746,6 +745,13 @@ conservative thing and refrain from strictifying a dfun's argument
 dictionaries.
 -}
 
+setBndrsDemandInfo :: [Var] -> [Demand] -> [Var]
+setBndrsDemandInfo (b:bs) (d:ds)
+  | isTyVar b = b : setBndrsDemandInfo bs (d:ds)
+  | otherwise = setIdDemandInfo b d : setBndrsDemandInfo bs ds
+setBndrsDemandInfo [] ds = ASSERT( null ds ) []
+setBndrsDemandInfo bs _  = pprPanic "setBndrsDemandInfo" (ppr bs)
+
 annotateBndr :: AnalEnv -> DmdType -> Var -> (DmdType, Var)
 -- The returned env has the var deleted
 -- The returned var is annotated with demand info
@@ -757,9 +763,6 @@ annotateBndr env dmd_ty var
   where
     (dmd_ty', dmd) = findBndrDmd env False dmd_ty var
 
-annotateBndrs :: AnalEnv -> DmdType -> [Var] -> (DmdType, [Var])
-annotateBndrs env = mapAccumR (annotateBndr env)
-
 annotateLamBndrs :: AnalEnv -> DFunFlag -> DmdType -> [Var] -> (DmdType, [Var])
 annotateLamBndrs env args_of_dfun ty bndrs = mapAccumR annotate ty bndrs
   where
@@ -1085,6 +1088,29 @@ extendSigsWithLam env id
   | otherwise
   = env
 
+addDataConStrictness :: DataCon -> [Demand] -> [Demand]
+-- See Note [Add demands for strict constructors]
+addDataConStrictness con ds
+  = ASSERT2( equalLength strs ds, ppr con $$ ppr strs $$ ppr ds )
+    zipWith add ds strs
+  where
+    strs = dataConRepStrictness con
+    add dmd str | isMarkedStrict str = dmd `bothDmd` seqDmd
+                | otherwise          = dmd
+    -- Yes, even if 'dmd' is Absent!
+
+findBndrsDmds :: AnalEnv -> DmdType -> [Var] -> (DmdType, [Demand])
+-- Return the demands on the Ids in the [Var]
+findBndrsDmds env dmd_ty bndrs
+  = go dmd_ty bndrs
+  where
+    go dmd_ty []  = (dmd_ty, [])
+    go dmd_ty (b:bs)
+      | isId b    = let (dmd_ty1, dmds) = go dmd_ty bs
+                        (dmd_ty2, dmd)  = findBndrDmd env False dmd_ty1 b
+                    in (dmd_ty2, dmd : dmds)
+      | otherwise = go dmd_ty bs
+
 findBndrDmd :: AnalEnv -> Bool -> DmdType -> Id -> (DmdType, Demand)
 -- See Note [Trimming a demand to a type] in Demand.lhs
 findBndrDmd env arg_of_dfun dmd_ty id