Better arity for join points

A join point was getting too large an arity, leading to #17294.
I've tightened up the invariant: see
  CoreSyn, Note [Invariants on join points], invariant 2b

7 files changed, 49 insertions, 24 deletions

diff --git a/compiler/coreSyn/CoreArity.hs b/compiler/coreSyn/CoreArity.hs
index d940d9d69c..04c8557882 100644
--- a/compiler/coreSyn/CoreArity.hs
+++ b/compiler/coreSyn/CoreArity.hs
@@ -881,7 +881,7 @@ inside the RHS of the join as well as into the body.  AND if j
 has an unfolding we have to push it into there too.  AND j might
 be recursive...
 
-So for now I'm abandonig the no-crap rule in this case. I think
+So for now I'm abandoning the no-crap rule in this case. I think
 that for the use in CorePrep it really doesn't matter; and if
 it does, then CoreToStg.myCollectArgs will fall over.
 
diff --git a/compiler/coreSyn/CorePrep.hs b/compiler/coreSyn/CorePrep.hs
index 2b68c2716b..6757f7aac9 100644
--- a/compiler/coreSyn/CorePrep.hs
+++ b/compiler/coreSyn/CorePrep.hs
@@ -560,6 +560,7 @@ it seems good for CorePrep to be robust.
 cpeJoinPair :: CorePrepEnv -> JoinId -> CoreExpr
             -> UniqSM (JoinId, CpeRhs)
 -- Used for all join bindings
+-- No eta-expansion: see Note [Do not eta-expand join points] in SimplUtils
 cpeJoinPair env bndr rhs
   = ASSERT(isJoinId bndr)
     do { let Just join_arity = isJoinId_maybe bndr
diff --git a/compiler/coreSyn/CoreSyn.hs b/compiler/coreSyn/CoreSyn.hs
index d94761b237..e3ad4715f1 100644
--- a/compiler/coreSyn/CoreSyn.hs
+++ b/compiler/coreSyn/CoreSyn.hs
@@ -681,9 +681,21 @@ Join points must follow these invariants:
   2. For join arity n, the right-hand side must begin with at least n lambdas.
      No ticks, no casts, just lambdas!  C.f. CoreUtils.joinRhsArity.
 
-  2a. Moreover, this same constraint applies to any unfolding of the binder.
-     Reason: if we want to push a continuation into the RHS we must push it
-     into the unfolding as well.
+     2a. Moreover, this same constraint applies to any unfolding of
+         the binder.  Reason: if we want to push a continuation into
+         the RHS we must push it into the unfolding as well.
+
+     2b. The Arity (in the IdInfo) of a join point is the number of value
+         binders in the top n lambdas, where n is the join arity.
+
+         So arity <= join arity; the former counts only value binders
+         while the latter counts all binders.
+         e.g. Suppose $j has join arity 1
+               let j = \x y. e in case x of { A -> j 1; B -> j 2 }
+         Then its ordinary arity is also 1, not 2.
+
+         The arity of a join point isn't very important; but short of setting
+         it to zero, it is helpful to have an invariant.  E.g. #17294.
 
   3. If the binding is recursive, then all other bindings in the recursive group
      must also be join points.
diff --git a/compiler/coreSyn/CoreTidy.hs b/compiler/coreSyn/CoreTidy.hs
index dfb031df7f..135d8e9b5b 100644
--- a/compiler/coreSyn/CoreTidy.hs
+++ b/compiler/coreSyn/CoreTidy.hs
@@ -18,7 +18,6 @@ import GhcPrelude
 
 import CoreSyn
 import CoreSeq ( seqUnfolding )
-import CoreArity
 import Id
 import IdInfo
 import Demand ( zapUsageEnvSig )
@@ -45,14 +44,15 @@ tidyBind :: TidyEnv
          ->  (TidyEnv, CoreBind)
 
 tidyBind env (NonRec bndr rhs)
-  = tidyLetBndr env env (bndr,rhs) =: \ (env', bndr') ->
+  = tidyLetBndr env env bndr =: \ (env', bndr') ->
     (env', NonRec bndr' (tidyExpr env' rhs))
 
 tidyBind env (Rec prs)
   = let
-       (env', bndrs') = mapAccumL (tidyLetBndr env') env prs
+       (bndrs, rhss)  = unzip prs
+       (env', bndrs') = mapAccumL (tidyLetBndr env') env bndrs
     in
-    map (tidyExpr env') (map snd prs)   =: \ rhss' ->
+    map (tidyExpr env') rhss =: \ rhss' ->
     (env', Rec (zip bndrs' rhss'))
 
 
@@ -166,10 +166,10 @@ tidyIdBndr env@(tidy_env, var_env) id
 
 tidyLetBndr :: TidyEnv         -- Knot-tied version for unfoldings
             -> TidyEnv         -- The one to extend
-            -> (Id, CoreExpr) -> (TidyEnv, Var)
+            -> Id -> (TidyEnv, Id)
 -- Used for local (non-top-level) let(rec)s
 -- Just like tidyIdBndr above, but with more IdInfo
-tidyLetBndr rec_tidy_env env@(tidy_env, var_env) (id,rhs)
+tidyLetBndr rec_tidy_env env@(tidy_env, var_env) id
   = case tidyOccName tidy_env (getOccName id) of { (tidy_env', occ') ->
     let
         ty'      = tidyType env (idType id)
@@ -193,13 +193,15 @@ tidyLetBndr rec_tidy_env env@(tidy_env, var_env) (id,rhs)
         --      (See Note [Zapping DmdEnv after Demand Analyzer] in WorkWrap)
         --
         -- Similarly arity info for eta expansion in CorePrep
+        -- Don't attempt to recompute arity here; this is just tidying!
+        -- Trying to do so led to #17294
         --
         -- Set inline-prag info so that we preseve it across
         -- separate compilation boundaries
         old_info = idInfo id
         new_info = vanillaIdInfo
                     `setOccInfo`        occInfo old_info
-                    `setArityInfo`      exprArity rhs
+                    `setArityInfo`      arityInfo old_info
                     `setStrictnessInfo` zapUsageEnvSig (strictnessInfo old_info)
                     `setDemandInfo`     demandInfo old_info
                     `setInlinePragInfo` inlinePragInfo old_info
@@ -209,6 +211,7 @@ tidyLetBndr rec_tidy_env env@(tidy_env, var_env) (id,rhs)
         new_unf | isStableUnfolding old_unf = tidyUnfolding rec_tidy_env old_unf old_unf
                 | otherwise                 = zapUnfolding old_unf
                                               -- See Note [Preserve evaluatedness]
+
     in
     ((tidy_env', var_env'), id') }
 
diff --git a/compiler/simplCore/SimplMonad.hs b/compiler/simplCore/SimplMonad.hs
index 732805e9ee..32c277cc55 100644
--- a/compiler/simplCore/SimplMonad.hs
+++ b/compiler/simplCore/SimplMonad.hs
@@ -36,7 +36,8 @@ import Outputable
 import FastString
 import MonadUtils
 import ErrUtils as Err
-import Panic (throwGhcExceptionIO, GhcException (..))
+import Util                ( count )
+import Panic               (throwGhcExceptionIO, GhcException (..))
 import BasicTypes          ( IntWithInf, treatZeroAsInf, mkIntWithInf )
 import Control.Monad       ( ap )
 
@@ -186,8 +187,8 @@ newJoinId bndrs body_ty
   = do { uniq <- getUniqueM
        ; let name       = mkSystemVarName uniq (fsLit "$j")
              join_id_ty = mkLamTypes bndrs body_ty  -- Note [Funky mkLamTypes]
-             -- Note [idArity for join points] in SimplUtils
-             arity      = length (filter isId bndrs)
+             arity      = count isId bndrs
+             -- arity: See Note [Invariants on join points] invariant 2b, in CoreSyn
              join_arity = length bndrs
              details    = JoinId join_arity
              id_info    = vanillaIdInfo `setArityInfo` arity
diff --git a/compiler/simplCore/SimplUtils.hs b/compiler/simplCore/SimplUtils.hs
index 4eeb51ceaa..eb57720d9e 100644
--- a/compiler/simplCore/SimplUtils.hs
+++ b/compiler/simplCore/SimplUtils.hs
@@ -1517,7 +1517,7 @@ tryEtaExpandRhs mode bndr rhs
          -- Note [Do not eta-expand join points]
          -- But do return the correct arity and bottom-ness, because
          -- these are used to set the bndr's IdInfo (#15517)
-         -- Note [idArity for join points]
+         -- Note [Invariants on join points] invariant 2b, in CoreSyn
 
   | otherwise
   = do { (new_arity, is_bot, new_rhs) <- try_expand
@@ -1611,13 +1611,6 @@ CorePrep comes around, the code is very likely to look more like this:
              $j2 = if n > 0 then $j1
                             else (...) eta
 
-Note [idArity for join points]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-Because of Note [Do not eta-expand join points] we have it that the idArity
-of a join point is always (less than or) equal to the join arity.
-Essentially, for join points we set `idArity $j = count isId join_lam_bndrs`.
-It really can be less if there are type-level binders in join_lam_bndrs.
-
 Note [Do not eta-expand PAPs]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 We used to have old_arity = manifestArity rhs, which meant that we
diff --git a/compiler/stranal/DmdAnal.hs b/compiler/stranal/DmdAnal.hs
index 9a4c64bdbb..afde951e60 100644
--- a/compiler/stranal/DmdAnal.hs
+++ b/compiler/stranal/DmdAnal.hs
@@ -603,8 +603,8 @@ dmdAnalRhsLetDown top_lvl rec_flag env let_dmd id rhs
     rhs_arity      = idArity id
     rhs_dmd
       -- See Note [Demand analysis for join points]
-      -- See Note [idArity for join points] in SimplUtils
-      -- rhs_arity matches the join arity of the join point
+      -- See Note [Invariants on join points] invariant 2b, in CoreSyn
+      --     rhs_arity matches the join arity of the join point
       | isJoinId id
       = mkCallDmds rhs_arity let_dmd
       | otherwise
@@ -727,6 +727,21 @@ let_dmd here).
 
 Another win for join points!  #13543.
 
+However, note that the strictness signature for a join point can
+look a little puzzling.  E.g.
+
+    (join j x = \y. error "urk")
+    (in case v of              )
+    (     A -> j 3             )  x
+    (     B -> j 4             )
+    (     C -> \y. blah        )
+
+The entire thing is in a C(S) context, so j's strictness signature
+will be    [A]b
+meaning one absent argument, returns bottom.  That seems odd because
+there's a \y inside.  But it's right because when consumed in a C(1)
+context the RHS of the join point is indeed bottom.
+
 Note [Demand signatures are computed for a threshold demand based on idArity]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 We compute demand signatures assuming idArity incoming arguments to approximate