Make arityType deal with join points

As Note [Eta-expansion and join points] describes,
this patch makes arityType deal correctly with join points.
What was there before was not wrong, but yielded lower
arities than it could.

Fixes #18328

In base GHC this makes no difference to nofib.

        Program           Size    Allocs   Runtime   Elapsed  TotalMem
--------------------------------------------------------------------------------
         n-body          -0.1%     -0.1%     -1.2%     -1.1%      0.0%
--------------------------------------------------------------------------------
            Min          -0.1%     -0.1%    -55.0%    -56.5%      0.0%
            Max          -0.0%      0.0%    +16.1%    +13.4%      0.0%
 Geometric Mean          -0.0%     -0.0%    -30.1%    -31.0%     -0.0%

But it starts to make real difference when we land the change to the
way mkDupableAlts handles StrictArg, in fixing #13253 and friends.
I think this is because we then get more non-inlined join points.

4 files changed, 195 insertions, 4 deletions

diff --git a/compiler/GHC/Core/Opt/Arity.hs b/compiler/GHC/Core/Opt/Arity.hs
index 8d51457ae0..2b2a7c20ea 100644
--- a/compiler/GHC/Core/Opt/Arity.hs
+++ b/compiler/GHC/Core/Opt/Arity.hs
@@ -36,6 +36,7 @@ import GHC.Core.TyCon     ( initRecTc, checkRecTc )
 import GHC.Core.Predicate ( isDictTy )
 import GHC.Core.Coercion as Coercion
 import GHC.Core.Multiplicity
+import GHC.Types.Var.Set
 import GHC.Types.Basic
 import GHC.Types.Unique
 import GHC.Driver.Session ( DynFlags, GeneralFlag(..), gopt )
@@ -156,7 +157,9 @@ exprBotStrictness_maybe e
         Nothing -> Nothing
         Just ar -> Just (ar, sig ar)
   where
-    env    = AE { ae_ped_bot = True, ae_cheap_fn = \ _ _ -> False }
+    env    = AE { ae_ped_bot = True
+                , ae_cheap_fn = \ _ _ -> False
+                , ae_joins = emptyVarSet }
     sig ar = mkClosedStrictSig (replicate ar topDmd) botDiv
 
 {-
@@ -506,7 +509,8 @@ exprEtaExpandArity dflags e
       ABot n   -> n
   where
     env = AE { ae_cheap_fn = mk_cheap_fn dflags isCheapApp
-             , ae_ped_bot  = gopt Opt_PedanticBottoms dflags }
+             , ae_ped_bot  = gopt Opt_PedanticBottoms dflags
+             , ae_joins    = emptyVarSet }
 
 getBotArity :: ArityType -> Maybe Arity
 -- Arity of a divergent function
@@ -578,7 +582,8 @@ findRhsArity dflags bndr rhs old_arity
           ATop _  -> (0,              False)    -- Note [Eta expanding thunks]
        where
          env = AE { ae_cheap_fn = mk_cheap_fn dflags cheap_app
-                  , ae_ped_bot  = gopt Opt_PedanticBottoms dflags }
+                  , ae_ped_bot  = gopt Opt_PedanticBottoms dflags
+                  , ae_joins    = emptyVarSet }
 
 {-
 Note [Arity analysis]
@@ -736,8 +741,15 @@ type CheapFun = CoreExpr -> Maybe Type -> Bool
 data ArityEnv
   = AE { ae_cheap_fn :: CheapFun
        , ae_ped_bot  :: Bool       -- True <=> be pedantic about bottoms
+       , ae_joins    :: IdSet      -- In-scope join points
+                                   -- See Note [Eta-expansion and join points]
   }
 
+extendJoinEnv :: ArityEnv -> [JoinId] -> ArityEnv
+extendJoinEnv env@(AE { ae_joins = joins }) join_ids
+  = env { ae_joins = joins `extendVarSetList` join_ids }
+
+----------------
 arityType :: ArityEnv -> CoreExpr -> ArityType
 
 arityType env (Cast e co)
@@ -755,7 +767,10 @@ arityType env (Cast e co)
     -- However, do make sure that ATop -> ATop and ABot -> ABot!
     --   Casts don't affect that part. Getting this wrong provoked #5475
 
-arityType _ (Var v)
+arityType env (Var v)
+  | v `elemVarSet` ae_joins env
+  = ABot 0  -- See Note [Eta-expansion and join points]
+
   | strict_sig <- idStrictness v
   , not $ isTopSig strict_sig
   , (ds, res) <- splitStrictSig strict_sig
@@ -804,6 +819,28 @@ arityType env (Case scrut _ _ alts)
   where
     alts_type = foldr1 andArityType [arityType env rhs | (_,_,rhs) <- alts]
 
+arityType env (Let (NonRec j rhs) body)
+  | Just join_arity <- isJoinId_maybe j
+  , (_, rhs_body)   <- collectNBinders join_arity rhs
+  = -- See Note [Eta-expansion and join points]
+    andArityType (arityType env rhs_body)
+                 (arityType env' body)
+  where
+     env' = extendJoinEnv env [j]
+
+arityType env (Let (Rec pairs) body)
+  | ((j,_):_) <- pairs
+  , isJoinId j
+  = -- See Note [Eta-expansion and join points]
+    foldr (andArityType . do_one) (arityType env' body) pairs
+  where
+    env' = extendJoinEnv env (map fst pairs)
+    do_one (j,rhs)
+      | Just arity <- isJoinId_maybe j
+      = arityType env' $ snd $ collectNBinders arity rhs
+      | otherwise
+      = pprPanic "arityType:joinrec" (ppr pairs)
+
 arityType env (Let b e)
   = floatIn (cheap_bind b) (arityType env e)
   where
@@ -816,6 +853,58 @@ arityType env (Tick t e)
 
 arityType _ _ = vanillaArityType
 
+{- Note [Eta-expansion and join points]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Consider this (#18328)
+
+  f x = join j y = case y of
+                      True -> \a. blah
+                      False -> \b. blah
+        in case x of
+              A -> j True
+              B -> \c. blah
+              C -> j False
+
+and suppose the join point is too big to inline.  Now, what is the
+arity of f?  If we inlined the join point, we'd definitely say "arity
+2" because we are prepared to push case-scrutinisation inside a
+lambda.  But currently the join point totally messes all that up,
+because (thought of as a vanilla let-binding) the arity pinned on 'j'
+is just 1.
+
+Why don't we eta-expand j?  Because of
+Note [Do not eta-expand join points] in GHC.Core.Opt.Simplify.Utils
+
+Even if we don't eta-expand j, why is its arity only 1?
+See invariant 2b in Note [Invariants on join points] in GHC.Core.
+
+So we do this:
+
+* Treat the RHS of a join-point binding, /after/ stripping off
+  join-arity lambda-binders, as very like the body of the let.
+  More precisely, do andArityType with the arityType from the
+  body of the let.
+
+* Dually, when we come to a /call/ of a join point, just no-op
+  by returning (ABot 0), the neutral element of ArityType.
+
+* This works if the join point is bound in the expression we are
+  taking the arityType of.  But if it's bound further out, it makes
+  no sense to say that (say) the arityType of (j False) is ABot 0.
+  Bad things happen.  So we keep track of the in-scope join-point Ids
+  in ae_join.
+
+This will make f, above, have arity 2. Then, we'll eta-expand it thus:
+
+  f x eta = (join j y = ... in case x of ...) eta
+
+and the Simplify will automatically push that application of eta into
+the join points.
+
+An alternative (roughly equivalent) idea would be to carry an
+environment mapping let-bound Ids to their ArityType.
+-}
+
 {-
 %************************************************************************
 %*                                                                      *
diff --git a/testsuite/tests/simplCore/should_compile/T18328.hs b/testsuite/tests/simplCore/should_compile/T18328.hs
new file mode 100644
index 0000000000..7f26d23c63
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T18328.hs
@@ -0,0 +1,14 @@
+module T18328 where
+
+f :: Int -> [a] -> [a] -> [a]
+f x ys = let {-# NOINLINE j #-}
+             j y = case x of
+                     3  -> ((++) ys) . ((++) ys) . ((++) ys) . ((++) ys)
+                     _  -> ((++) ys) . ((++) ys) . ((++) ys)
+
+         in
+         case x of
+           1 -> j 2
+           2 -> j 3
+           3 -> j 4
+           _ -> ((++) ys)
diff --git a/testsuite/tests/simplCore/should_compile/T18328.stderr b/testsuite/tests/simplCore/should_compile/T18328.stderr
new file mode 100644
index 0000000000..ce366ea367
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T18328.stderr
@@ -0,0 +1,87 @@
+
+==================== Tidy Core ====================
+Result size of Tidy Core
+  = {terms: 69, types: 61, coercions: 0, joins: 1/1}
+
+-- RHS size: {terms: 42, types: 28, coercions: 0, joins: 1/1}
+T18328.$wf [InlPrag=NOUSERINLINE[2]]
+  :: forall {a}. GHC.Prim.Int# -> [a] -> [a] -> [a]
+[GblId,
+ Arity=3,
+ Str=<S,U><S,U><L,1*U>,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [182 0 0] 312 0}]
+T18328.$wf
+  = \ (@a) (ww :: GHC.Prim.Int#) (w :: [a]) (w1 :: [a]) ->
+      join {
+        $wj [InlPrag=NOINLINE, Dmd=<L,1*C1(U)>]
+          :: forall {p}. GHC.Prim.Void# -> [a]
+        [LclId[JoinId(2)], Arity=1, Str=<L,A>, Unf=OtherCon []]
+        $wj (@p) _ [Occ=Dead, OS=OneShot]
+          = case ww of {
+              __DEFAULT -> ++ @a w (++ @a w (++ @a w w1));
+              3# -> ++ @a w (++ @a w (++ @a w (++ @a w w1)))
+            } } in
+      case ww of {
+        __DEFAULT -> ++ @a w w1;
+        1# -> jump $wj @Integer GHC.Prim.void#;
+        2# -> jump $wj @Integer GHC.Prim.void#;
+        3# -> jump $wj @Integer GHC.Prim.void#
+      }
+
+-- RHS size: {terms: 11, types: 10, coercions: 0, joins: 0/0}
+f [InlPrag=NOUSERINLINE[2]] :: forall a. Int -> [a] -> [a] -> [a]
+[GblId,
+ Arity=3,
+ Str=<S(S),1*U(U)><S,U><L,1*U>,
+ Unf=Unf{Src=InlineStable, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True,
+         Guidance=ALWAYS_IF(arity=3,unsat_ok=True,boring_ok=False)
+         Tmpl= \ (@a)
+                 (w [Occ=Once!] :: Int)
+                 (w1 [Occ=Once] :: [a])
+                 (w2 [Occ=Once] :: [a]) ->
+                 case w of { GHC.Types.I# ww1 [Occ=Once] ->
+                 T18328.$wf @a ww1 w1 w2
+                 }}]
+f = \ (@a) (w :: Int) (w1 :: [a]) (w2 :: [a]) ->
+      case w of { GHC.Types.I# ww1 -> T18328.$wf @a ww1 w1 w2 }
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+T18328.$trModule4 :: GHC.Prim.Addr#
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 20 0}]
+T18328.$trModule4 = "main"#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+T18328.$trModule3 :: GHC.Types.TrName
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 10 20}]
+T18328.$trModule3 = GHC.Types.TrNameS T18328.$trModule4
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+T18328.$trModule2 :: GHC.Prim.Addr#
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 30 0}]
+T18328.$trModule2 = "T18328"#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+T18328.$trModule1 :: GHC.Types.TrName
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 10 20}]
+T18328.$trModule1 = GHC.Types.TrNameS T18328.$trModule2
+
+-- RHS size: {terms: 3, types: 0, coercions: 0, joins: 0/0}
+T18328.$trModule :: GHC.Types.Module
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 10 30}]
+T18328.$trModule
+  = GHC.Types.Module T18328.$trModule3 T18328.$trModule1
+
+
+
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index e363a4ec26..93f7fc155a 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -328,5 +328,6 @@ test('T18231', [ only_ways(['optasm']), grep_errmsg(r'^[\w\.]+ ::.*->.*') ], com
 # Cast WW
 test('T17673', [ only_ways(['optasm']), grep_errmsg(r'^\w+\.\$wf') ], compile, ['-ddump-simpl -dsuppress-uniques -dppr-cols=9999'])
 test('T18078', [ only_ways(['optasm']), grep_errmsg(r'^\w+\.\$wf') ], compile, ['-ddump-simpl -dsuppress-uniques -dppr-cols=9999'])
+test('T18328', [ only_ways(['optasm']), grep_errmsg(r'Arity=') ], compile, ['-ddump-simpl -dsuppress-uniques'])
 test('T18347', normal, compile, ['-dcore-lint -O'])
 test('T18399', normal, compile, ['-dcore-lint -O'])