Introduce the new state hack in the demand analyzerwip/T9388

The new state hack, as proposed by SPJ in #9388, applies only to bound
expressions. It is implemented by constructing an artificial incoming demand that claims that it is called at most once.

(Currently, it also applies to non-top-level let-bound expressions. This
needs to be revisited.)

3 files changed, 18 insertions, 5 deletions

diff --git a/compiler/basicTypes/Demand.hs b/compiler/basicTypes/Demand.hs
index ecf22bc51f..7b08a71cf0 100644
--- a/compiler/basicTypes/Demand.hs
+++ b/compiler/basicTypes/Demand.hs
@@ -37,7 +37,7 @@ module Demand (
 
         seqDemand, seqDemandList, seqDmdType, seqStrictSig,
 
-        evalDmd, cleanEvalDmd, cleanEvalProdDmd, isStrictDmd,
+        evalDmd, cleanEvalDmd, cleanEvalProdDmd, cleanEvalStateHackDmd, isStrictDmd,
         splitDmdTy, splitFVs,
         deferAfterIO,
         postProcessUnsat, postProcessDmdTypeM,
@@ -634,6 +634,9 @@ cleanEvalDmd = mkCleanDmd HeadStr Used
 cleanEvalProdDmd :: Arity -> CleanDemand
 cleanEvalProdDmd n = mkCleanDmd HeadStr (UProd (replicate n useTop))
 
+cleanEvalStateHackDmd :: CleanDemand
+cleanEvalStateHackDmd = mkCleanDmd HeadStr (mkUCall One Used)
+
 isSingleUsed :: JointDmd -> Bool
 isSingleUsed (JD {absd=a}) = is_used_once a
   where
diff --git a/compiler/basicTypes/Id.hs b/compiler/basicTypes/Id.hs
index 23d9c30066..7459cc1ba7 100644
--- a/compiler/basicTypes/Id.hs
+++ b/compiler/basicTypes/Id.hs
@@ -71,7 +71,7 @@ module Id (
         isOneShotBndr, isOneShotLambda, isProbablyOneShotLambda,
         setOneShotLambda, clearOneShotLambda,
         updOneShotInfo, setIdOneShotInfo,
-        isStateHackType,
+        isStateHackType, isStateHackFunType,
 
         -- ** Reading 'IdInfo' fields
         idArity,
@@ -674,6 +674,11 @@ isStateHackType ty
         -- Another good example is in fill_in in PrelPack.hs.  We should be able to
         -- spot that fill_in has arity 2 (and when Keith is done, we will) but we can't yet.
 
+isStateHackFunType :: Type -> Bool
+isStateHackFunType ty
+  = case splitFunTy_maybe ty of
+        Just (arg_ty, _) -> isStateHackType arg_ty
+        Nothing -> False
 
 -- | Returns whether the lambda associated with the 'Id' is certainly applied at most once.
 -- You probably want to use 'isOneShotBndr' instead
diff --git a/compiler/stranal/DmdAnal.hs b/compiler/stranal/DmdAnal.hs
index 27fa35fba0..d6155622dd 100644
--- a/compiler/stranal/DmdAnal.hs
+++ b/compiler/stranal/DmdAnal.hs
@@ -619,9 +619,14 @@ dmdAnalRhs top_lvl rec_flag env id rhs
         -- See Note [NOINLINE and strictness]
 
     -- See Note [Product demands for function body]
-    body_dmd = case deepSplitProductType_maybe (ae_fam_envs env) (exprType body) of
-                 Nothing            -> cleanEvalDmd
-                 Just (dc, _, _, _) -> cleanEvalProdDmd (dataConRepArity dc)
+    body_dmd
+        | Just (dc, _, _, _) <- deepSplitProductType_maybe (ae_fam_envs env) (exprType body)
+        = cleanEvalProdDmd (dataConRepArity dc)
+        | isStateHackFunType $ topNormaliseType (ae_fam_envs env) (exprType body)
+        = -- pprTrace "new state hack" (ppr (exprType body)) $
+          cleanEvalStateHackDmd
+        | otherwise
+        = cleanEvalDmd
 
     -- See Note [Lazy and unleashable free variables]
     -- See Note [Aggregated demand for cardinality]