Make the demand on a binder compatible with type (fixes Trac #8569)

Because of GADTs and casts we were getting binders whose
demand annotation was more deeply nested than made sense
for its type.

See Note [Trimming a demand to a type], in Demand.lhs,
which I reproduce here:

   Note [Trimming a demand to a type]
   ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
   Consider this:

     f :: a -> Bool
     f x = case ... of
             A g1 -> case (x |> g1) of (p,q) -> ...
             B    -> error "urk"

   where A,B are the constructors of a GADT.  We'll get a U(U,U) demand
   on x from the A branch, but that's a stupid demand for x itself, which
   has type 'a'. Indeed we get ASSERTs going off (notably in
   splitUseProdDmd, Trac #8569).

   Bottom line: we really don't want to have a binder whose demand is more
   deeply-nested than its type.  There are various ways to tackle this.
   When processing (x |> g1), we could "trim" the incoming demand U(U,U)
   to match x's type.  But I'm currently doing so just at the moment when
   we pin a demand on a binder, in DmdAnal.findBndrDmd.

3 files changed, 127 insertions, 36 deletions

diff --git a/compiler/basicTypes/Demand.lhs b/compiler/basicTypes/Demand.lhs
index e415c6d938..8a082b98ad 100644
--- a/compiler/basicTypes/Demand.lhs
+++ b/compiler/basicTypes/Demand.lhs
@@ -44,6 +44,7 @@ module Demand (
         splitProdDmd, splitProdDmd_maybe, peelCallDmd, mkCallDmd,
         dmdTransformSig, dmdTransformDataConSig, dmdTransformDictSelSig,
         argOneShots, argsOneShots,
+        trimToType, TypeShape(..),
 
         isSingleUsed, reuseEnv, zapDemand, zapStrictSig,
 
@@ -67,6 +68,7 @@ import Maybes           ( orElse )
 import Type            ( Type )
 import TyCon           ( isNewTyCon, isClassTyCon )
 import DataCon         ( splitDataProductType_maybe )
+import FastString
 \end{code}
 
 %************************************************************************
@@ -442,7 +444,7 @@ seqMaybeUsed _          = ()
 splitUseProdDmd :: Int -> UseDmd -> [MaybeUsed]
 splitUseProdDmd n Used          = replicate n useTop
 splitUseProdDmd n UHead         = replicate n Abs
-splitUseProdDmd n (UProd ds)    = ASSERT2( ds `lengthIs` n, ppr n $$ ppr ds ) ds
+splitUseProdDmd n (UProd ds)    = ASSERT2( ds `lengthIs` n, text "splitUseProdDmd" $$ ppr n $$ ppr ds ) ds
 splitUseProdDmd _ d@(UCall _ _) = pprPanic "attempt to prod-split usage call demand" (ppr d)
 \end{code}
   
@@ -638,8 +640,66 @@ isSingleUsed (JD {absd=a}) = is_used_once a
     is_used_once Abs         = True
     is_used_once (Use One _) = True
     is_used_once _           = False
+
+
+data TypeShape = TsFun TypeShape
+               | TsProd [TypeShape]
+               | TsUnk
+
+instance Outputable TypeShape where
+  ppr TsUnk        = ptext (sLit "TsUnk")
+  ppr (TsFun ts)   = ptext (sLit "TsFun") <> parens (ppr ts)
+  ppr (TsProd tss) = parens (hsep $ punctuate comma $ map ppr tss)
+
+trimToType :: JointDmd -> TypeShape -> JointDmd
+-- See Note [Trimming a demand to a type]
+trimToType (JD ms mu) ts
+  = JD (go_ms ms ts) (go_mu mu ts)
+  where
+    go_ms :: MaybeStr -> TypeShape -> MaybeStr
+    go_ms Lazy    _  = Lazy
+    go_ms (Str s) ts = Str (go_s s ts)
+
+    go_s :: StrDmd -> TypeShape -> StrDmd
+    go_s HyperStr    _            = HyperStr
+    go_s (SCall s)   (TsFun ts)   = SCall (go_s s ts)
+    go_s (SProd mss) (TsProd tss)
+      | equalLength mss tss       = SProd (zipWith go_ms mss tss)
+    go_s _           _            = HeadStr
+
+    go_mu :: MaybeUsed -> TypeShape -> MaybeUsed
+    go_mu Abs _ = Abs
+    go_mu (Use c u) ts = Use c (go_u u ts)
+
+    go_u :: UseDmd -> TypeShape -> UseDmd
+    go_u UHead       _          = UHead
+    go_u (UCall c u) (TsFun ts) = UCall c (go_u u ts)
+    go_u (UProd mus) (TsProd tss)
+      | equalLength mus tss      = UProd (zipWith go_mu mus tss)
+    go_u _           _           = Used
 \end{code}
 
+Note [Trimming a demand to a type]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Consider this:
+
+  f :: a -> Bool
+  f x = case ... of
+          A g1 -> case (x |> g1) of (p,q) -> ...
+          B    -> error "urk"
+
+where A,B are the constructors of a GADT.  We'll get a U(U,U) demand
+on x from the A branch, but that's a stupid demand for x itself, which
+has type 'a'. Indeed we get ASSERTs going off (notably in
+splitUseProdDmd, Trac #8569).
+
+Bottom line: we really don't want to have a binder whose demand is more
+deeply-nested than its type.  There are various ways to tackle this.
+When processing (x |> g1), we could "trim" the incoming demand U(U,U)
+to match x's type.  But I'm currently doing so just at the moment when
+we pin a demand on a binder, in DmdAnal.findBndrDmd.
+
+
 Note [Threshold demands]
 ~~~~~~~~~~~~~~~~~~~~~~~~
 Threshold usage demand is generated to figure out if
@@ -1451,7 +1511,7 @@ dmdTransformDataConSig :: Arity -> StrictSig -> CleanDemand -> DmdType
 -- which has a special kind of demand transformer.
 -- If the constructor is saturated, we feed the demand on 
 -- the result into the constructor arguments.
-dmdTransformDataConSig arity (StrictSig (DmdType _ _ con_res)) 
+dmdTransformDataConSig arity (StrictSig (DmdType _ _ con_res))
                              (CD { sd = str, ud = abs })
   | Just str_dmds <- go_str arity str
   , Just abs_dmds <- go_abs arity abs
diff --git a/compiler/stranal/DmdAnal.lhs b/compiler/stranal/DmdAnal.lhs
index 88eea0c03b..3294371964 100644
--- a/compiler/stranal/DmdAnal.lhs
+++ b/compiler/stranal/DmdAnal.lhs
@@ -13,9 +13,8 @@ module DmdAnal ( dmdAnalProgram ) where
 
 #include "HsVersions.h"
 
-import Var		( isTyVar )
 import DynFlags
-import WwLib            ( deepSplitProductType_maybe )
+import WwLib            ( findTypeShape, deepSplitProductType_maybe )
 import Demand	-- All of it
 import CoreSyn
 import Outputable
@@ -26,11 +25,8 @@ import Data.List
 import DataCon
 import Id
 import CoreUtils	( exprIsHNF, exprType, exprIsTrivial )
--- import PprCore	
 import TyCon
-import Type		( eqType )
--- import Pair
--- import Coercion         ( coercionKind )
+import Type
 import FamInstEnv
 import Util
 import Maybes		( isJust )
@@ -492,8 +488,7 @@ dmdTransform :: AnalEnv		-- The strictness environment
 
 dmdTransform env var dmd
   | isDataConWorkId var		                 -- Data constructor
-  = dmdTransformDataConSig 
-       (idArity var) (idStrictness var) dmd
+  = dmdTransformDataConSig (idArity var) (idStrictness var) dmd
 
   | gopt Opt_DmdTxDictSel (ae_dflags env),
     Just _ <- isClassOpId_maybe var -- Dictionary component selector
@@ -728,9 +723,8 @@ addLazyFVs dmd_ty lazy_fvs
 	-- call to f.  So we just get an L demand for x for g.
 \end{code}
 
-Note [do not strictify the argument dictionaries of a dfun]
+Note [Do not strictify the argument dictionaries of a dfun]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
 The typechecker can tie recursive knots involving dfuns, so we do the
 conservative thing and refrain from strictifying a dfun's argument
 dictionaries.
@@ -742,17 +736,10 @@ annotateBndr :: AnalEnv -> DmdType -> Var -> (DmdType, Var)
 -- according to the result demand of the provided demand type
 -- No effect on the argument demands
 annotateBndr env dmd_ty var
-  | isTyVar var = (dmd_ty, var)
-  | otherwise   = (dmd_ty', set_idDemandInfo env var dmd')
+  | isId var  = (dmd_ty', setIdDemandInfo var dmd)
+  | otherwise = (dmd_ty, var)
   where
-    (dmd_ty', dmd) = peelFV dmd_ty var
-
-    dmd' | gopt Opt_DictsStrict (ae_dflags env)
-             -- We never want to strictify a recursive let. At the moment
-             -- annotateBndr is only call for non-recursive lets; if that
-             -- changes, we need a RecFlag parameter and another guard here.
-         = strictifyDictDmd (idType var) dmd
-         | otherwise = dmd
+    (dmd_ty', dmd) = findBndrDmd env False dmd_ty var
 
 annotateBndrs :: AnalEnv -> DmdType -> [Var] -> (DmdType, [Var])
 annotateBndrs env = mapAccumR (annotateBndr env)
@@ -777,7 +764,7 @@ annotateLamIdBndr env arg_of_dfun dmd_ty one_shot id
 -- Only called for Ids
   = ASSERT( isId id )
     -- pprTrace "annLamBndr" (vcat [ppr id, ppr _dmd_ty]) $
-    (final_ty, setOneShotness one_shot (set_idDemandInfo env id dmd'))
+    (final_ty, setOneShotness one_shot (setIdDemandInfo id dmd))
   where
       -- Watch out!  See note [Lambda-bound unfoldings]
     final_ty = case maybeUnfoldingTemplate (idUnfolding id) of
@@ -787,13 +774,7 @@ annotateLamIdBndr env arg_of_dfun dmd_ty one_shot id
                              (unf_ty, _) = dmdAnalStar env dmd unf
 
     main_ty = addDemand dmd dmd_ty'
-    (dmd_ty', dmd) = peelFV dmd_ty id
-
-    dmd' | gopt Opt_DictsStrict (ae_dflags env),
-           -- see Note [do not strictify the argument dictionaries of a dfun]
-           not arg_of_dfun
-         = strictifyDictDmd (idType id) dmd
-         | otherwise = dmd
+    (dmd_ty', dmd) = findBndrDmd env arg_of_dfun dmd_ty id
 
 deleteFVs :: DmdType -> [Var] -> DmdType
 deleteFVs (DmdType fvs dmds res) bndrs
@@ -1079,18 +1060,39 @@ extendSigsWithLam :: AnalEnv -> Id -> AnalEnv
 -- Extend the AnalEnv when we meet a lambda binder
 extendSigsWithLam env id
   | isId id
-  , isStrictDmd (idDemandInfo id) || ae_virgin env  
+  , isStrictDmd (idDemandInfo id) || ae_virgin env
        -- See Note [Optimistic CPR in the "virgin" case]
        -- See Note [Initial CPR for strict binders]
   , Just (dc,_,_,_) <- deepSplitProductType_maybe (ae_fam_envs env) $ idType id
   = extendAnalEnv NotTopLevel env id (cprProdSig (dataConRepArity dc))
 
-  | otherwise 
+  | otherwise
   = env
 
-set_idDemandInfo :: AnalEnv -> Id -> Demand -> Id
-set_idDemandInfo env id dmd 
-  = setIdDemandInfo id (zapDemand (ae_dflags env) dmd)
+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
+  = (dmd_ty', dmd')
+  where
+    dmd' = zapDemand (ae_dflags env) $
+           strictify $
+           trimToType starting_dmd (findTypeShape fam_envs id_ty)
+
+    (dmd_ty', starting_dmd) = peelFV dmd_ty id
+
+    id_ty = idType id
+
+    strictify dmd
+      | gopt Opt_DictsStrict (ae_dflags env)
+             -- We never want to strictify a recursive let. At the moment
+             -- annotateBndr is only call for non-recursive lets; if that
+             -- changes, we need a RecFlag parameter and another guard here.
+      , not arg_of_dfun -- See Note [Do not strictify the argument dictionaries of a dfun]
+      = strictifyDictDmd id_ty dmd
+      | otherwise
+      = dmd
+
+    fam_envs = ae_fam_envs env
 
 set_idStrictness :: AnalEnv -> Id -> StrictSig -> Id
 set_idStrictness env id sig
diff --git a/compiler/stranal/WwLib.lhs b/compiler/stranal/WwLib.lhs
index f88c9ad54f..68292839ed 100644
--- a/compiler/stranal/WwLib.lhs
+++ b/compiler/stranal/WwLib.lhs
@@ -4,7 +4,9 @@
 \section[WwLib]{A library for the ``worker\/wrapper'' back-end to the strictness analyser}
 
 \begin{code}
-module WwLib ( mkWwBodies, mkWWstr, mkWorkerArgs, deepSplitProductType_maybe ) where
+module WwLib ( mkWwBodies, mkWWstr, mkWorkerArgs
+             , deepSplitProductType_maybe, findTypeShape
+ ) where
 
 #include "HsVersions.h"
 
@@ -506,6 +508,12 @@ match the number of constructor arguments; this happened in Trac #8037.
 If so, the worker/wrapper split doesn't work right and we get a Core Lint
 bug.  The fix here is simply to decline to do w/w if that happens.
 
+%************************************************************************
+%*                                                                      *
+         Type scrutiny that is specfic to demand analysis
+%*                                                                      *
+%************************************************************************
+
 \begin{code}
 deepSplitProductType_maybe :: FamInstEnvs -> Type -> Maybe (DataCon, [Type], [Type], Coercion)
 -- If    deepSplitProductType_maybe ty = Just (dc, tys, arg_tys, co)
@@ -534,6 +542,27 @@ deepSplitCprType_maybe fam_envs con_tag ty
   , let con  = cons !! (con_tag - fIRST_TAG)
   = Just (con, tc_args, dataConInstArgTys con tc_args, co)
 deepSplitCprType_maybe _ _ _ = Nothing
+
+findTypeShape :: FamInstEnvs -> Type -> TypeShape
+-- Uncover the arrow and product shape of a type
+-- The data type TypeShape is defined in Demand
+-- See Note [Trimming a demand to a type] in Demand
+findTypeShape fam_envs ty
+  | Just (_, ty') <- splitForAllTy_maybe ty
+  = findTypeShape fam_envs ty'
+
+  | Just (tc, tc_args)  <- splitTyConApp_maybe ty
+  , Just con <- isDataProductTyCon_maybe tc
+  = TsProd (map (findTypeShape fam_envs) $ dataConInstArgTys con tc_args)
+
+  | Just (_, res) <- splitFunTy_maybe ty
+  = TsFun (findTypeShape fam_envs res)
+
+  | Just (_, ty') <- topNormaliseType_maybe fam_envs ty
+  = findTypeShape fam_envs ty'
+
+  | otherwise
+  = TsUnk
 \end{code}