[commit: ghc] wip/dmd-arity: Compute strictness signatures assuming exprArity (00a9b1d)

[git at git.haskell.org]

Repository : ssh://git@git.haskell.org/ghc

On branch  : wip/dmd-arity
Link       : http://ghc.haskell.org/trac/ghc/changeset/00a9b1d39abf6fcf6deecdd219db588f906c88f3/ghc

>---------------------------------------------------------------

commit 00a9b1d39abf6fcf6deecdd219db588f906c88f3
Author: Sebastian Graf <sebastian.graf at kit.edu>
Date:   Thu Feb 7 19:02:28 2019 +0100

    Compute strictness signatures assuming exprArity


>---------------------------------------------------------------

00a9b1d39abf6fcf6deecdd219db588f906c88f3
 compiler/basicTypes/Demand.hs |  6 +-----
 compiler/stranal/DmdAnal.hs   | 28 +++++++++++++++++-----------
 2 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/compiler/basicTypes/Demand.hs b/compiler/basicTypes/Demand.hs
index 52bdf67..a4ba7c9 100644
--- a/compiler/basicTypes/Demand.hs
+++ b/compiler/basicTypes/Demand.hs
@@ -22,7 +22,7 @@ module Demand (
 
         DmdType(..), dmdTypeDepth, lubDmdType, bothDmdType,
         nopDmdType, botDmdType, mkDmdType,
-        addDemand, removeDmdTyArgs,
+        addDemand, ensureArgs,
         BothDmdArg, mkBothDmdArg, toBothDmdArg,
 
         DmdEnv, emptyDmdEnv,
@@ -1207,10 +1207,6 @@ mkDmdType fv ds res = DmdType fv ds res
 dmdTypeDepth :: DmdType -> Arity
 dmdTypeDepth (DmdType _ ds _) = length ds
 
--- Remove any demand on arguments. This is used in dmdAnalRhs on the body
-removeDmdTyArgs :: DmdType -> DmdType
-removeDmdTyArgs = ensureArgs 0
-
 -- This makes sure we can use the demand type with n arguments,
 -- It extends the argument list with the correct resTypeArgDmd
 -- It also adjusts the DmdResult: Divergence survives additional arguments,
diff --git a/compiler/stranal/DmdAnal.hs b/compiler/stranal/DmdAnal.hs
index 10b9da6..684d379 100644
--- a/compiler/stranal/DmdAnal.hs
+++ b/compiler/stranal/DmdAnal.hs
@@ -628,6 +628,7 @@ dmdAnalRhsLetDown top_lvl rec_flag env let_dmd id rhs
   | otherwise
   = (lazy_fv, id', mkLams bndrs' body')
   where
+    rhs_arity        = exprArity rhs
     (bndrs, body, body_dmd)
        = case isJoinId_maybe id of
            Just join_arity  -- See Note [Demand analysis for join points]
@@ -635,15 +636,15 @@ dmdAnalRhsLetDown top_lvl rec_flag env let_dmd id rhs
                    -> (bndrs, body, let_dmd)
 
            Nothing | (bndrs, body) <- collectBinders rhs
-                   -> (bndrs, body, mkBodyDmd env body)
+                   -> (bndrs, body, mkBodyDmd env (exprArity body) body)
 
     env_body         = foldl' extendSigsWithLam env bndrs
     (body_ty, body') = dmdAnal env_body body_dmd body
-    body_ty'         = removeDmdTyArgs body_ty -- zap possible deep CPR info
-    (DmdType rhs_fv rhs_dmds rhs_res, bndrs')
-                     = annotateLamBndrs env (isDFunId id) body_ty' bndrs
+    (rhs_ty, bndrs') = annotateLamBndrs env (isDFunId id) body_ty bndrs
+    DmdType rhs_fv rhs_dmds rhs_res
+                     = ensureArgs rhs_arity rhs_ty -- zap possible deep CPR info
     sig_ty           = mkStrictSig (mkDmdType sig_fv rhs_dmds rhs_res')
-    id'              = set_idStrictness env id sig_ty
+    id'              = set_idStrictness env id rhs_arity sig_ty
         -- See Note [NOINLINE and strictness]
 
 
@@ -667,10 +668,14 @@ dmdAnalRhsLetDown top_lvl rec_flag env let_dmd id rhs
        || not (isStrictDmd (idDemandInfo id) || ae_virgin env)
           -- See Note [Optimistic CPR in the "virgin" case]
 
-mkBodyDmd :: AnalEnv -> CoreExpr -> CleanDemand
+-- | Creates a 'CleanDemand' appropriate for unleashing on the given function
+-- body, by wrapping a head demand into @arity@ many calls.
 -- See Note [Product demands for function body]
-mkBodyDmd env body
-  = case deepSplitProductType_maybe (ae_fam_envs env) (exprType body) of
+mkBodyDmd :: AnalEnv -> Arity -> CoreExpr -> CleanDemand
+mkBodyDmd env arity body
+  = iterate mkCallDmd base !! arity
+  where
+    base = case deepSplitProductType_maybe (ae_fam_envs env) (exprType body) of
        Nothing            -> cleanEvalDmd
        Just (dc, _, _, _) -> cleanEvalProdDmd (dataConRepArity dc)
 
@@ -1247,9 +1252,10 @@ findBndrDmd env arg_of_dfun dmd_ty id
 
     fam_envs = ae_fam_envs env
 
-set_idStrictness :: AnalEnv -> Id -> StrictSig -> Id
-set_idStrictness env id sig
-  = setIdStrictness id (killUsageSig (ae_dflags env) sig)
+set_idStrictness :: AnalEnv -> Id -> Arity -> StrictSig -> Id
+set_idStrictness env id arity sig
+  = id `setIdStrictness` (killUsageSig (ae_dflags env) sig)
+       `setIdArity` arity -- computed by exprArity and must match sig
 
 dumpStrSig :: CoreProgram -> SDoc
 dumpStrSig binds = vcat (map printId ids)