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.

2 files changed, 65 insertions, 34 deletions

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}