CoreLint - When checking for levity polymorphism look through more ticks.wip/andreask/core_lint_lev

For expressions like `(scc<cc_name> primOp#) arg1` we should also look
at arg1 to determine if we call primOp# at a fixed runtime rep.

This is what corePrep already does but CoreLint didn't yet. This patch
will bring them in sync in this regard.

It also uses tickishFloatable in CorePrep instead of CorePrep having
it's own slightly differing definition of when a tick is floatable.

4 files changed, 45 insertions, 13 deletions

diff --git a/compiler/GHC/Core.hs b/compiler/GHC/Core.hs
index e82f0b2d8a..36fa6e2673 100644
--- a/compiler/GHC/Core.hs
+++ b/compiler/GHC/Core.hs
@@ -45,6 +45,7 @@ module GHC.Core (
         collectBinders, collectTyBinders, collectTyAndValBinders,
         collectNBinders,
         collectArgs, stripNArgs, collectArgsTicks, flattenBinds,
+        collectFunSimple,
 
         exprToType,
         wrapLamBody,
@@ -2010,6 +2011,19 @@ collectArgs expr
     go (App f a) as = go f (a:as)
     go e         as = (e, as)
 
+-- | Takes a nested application expression and returns the function
+-- being applied. Looking through casts and ticks to find it.
+collectFunSimple :: Expr b -> Expr b
+collectFunSimple expr
+  = go expr
+  where
+    go expr' =
+      case expr' of
+        App f _a    -> go f
+        Tick _t e   -> go e
+        Cast e _co  -> go e
+        e           -> e
+
 -- | fmap on the body of a lambda.
 --   wrapLamBody f (\x -> body) == (\x -> f body)
 wrapLamBody :: (CoreExpr -> CoreExpr) -> CoreExpr -> CoreExpr
diff --git a/compiler/GHC/Core/Lint.hs b/compiler/GHC/Core/Lint.hs
index 037940eac2..df96afff61 100644
--- a/compiler/GHC/Core/Lint.hs
+++ b/compiler/GHC/Core/Lint.hs
@@ -1003,7 +1003,10 @@ lintCoreExpr e@(App _ _)
        ; checkCanEtaExpand fun args app_ty
        ; return app_pair}
   where
-    (fun, args, _source_ticks) = collectArgsTicks tickishFloatable e
+    skipTick t = case collectFunSimple e of
+      (Var v) -> etaExpansionTick v t
+      _ -> tickishFloatable t
+    (fun, args, _source_ticks) = collectArgsTicks skipTick e
       -- We must look through source ticks to avoid #21152, for example:
       --
       -- reallyUnsafePtrEquality
@@ -1014,6 +1017,8 @@ lintCoreExpr e@(App _ _)
       -- To do this, we use `collectArgsTicks tickishFloatable` to match
       -- the eta expansion behaviour, as per Note [Eta expansion and source notes]
       -- in GHC.Core.Opt.Arity.
+      -- Sadly this was not quite enough. So we now also accept things that CorePrep will allow.
+      -- See Note [Ticks and mandatory eta expansion]
 
 lintCoreExpr (Lam var expr)
   = markAllJoinsBad $
@@ -1319,7 +1324,9 @@ The basic version of these functions checks that the argument is a
 subtype of the required type, as one would expect.
 -}
 
-
+-- Takes the functions type and arguments as argument.
+-- Returns the *result* of applying the function to arguments.
+-- e.g. f :: Int -> Bool -> Int would return `Int` as result type.
 lintCoreArgs  :: (LintedType, UsageEnv) -> [CoreArg] -> LintM (LintedType, UsageEnv)
 lintCoreArgs (fun_ty, fun_ue) args = foldM lintCoreArg (fun_ty, fun_ue) args
 
diff --git a/compiler/GHC/Core/Utils.hs b/compiler/GHC/Core/Utils.hs
index b7b1c9334b..9df16053ac 100644
--- a/compiler/GHC/Core/Utils.hs
+++ b/compiler/GHC/Core/Utils.hs
@@ -31,7 +31,7 @@ module GHC.Core.Utils (
         isCheapApp, isExpandableApp, isSaturatedConApp,
         exprIsTickedString, exprIsTickedString_maybe,
         exprIsTopLevelBindable,
-        altsAreExhaustive,
+        altsAreExhaustive, etaExpansionTick,
 
         -- * Equality
         cheapEqExpr, cheapEqExpr', eqExpr,
@@ -332,6 +332,11 @@ mkTick t orig_expr = mkTick' id id orig_expr
   -- non-counting part having laxer placement properties.
   canSplit = tickishCanSplit t && tickishPlace (mkNoCount t) /= tickishPlace t
 
+  -- mkTick' handles floating of ticks *into* the expression.
+  -- In this function, `top` is applied after adding the tick, and `rest` before.
+  -- This will result in applications that look like (top $ Tick t $ rest expr).
+  -- If we want to push the tick deeper, we pre-compose `top` with a function
+  -- adding the tick.
   mkTick' :: (CoreExpr -> CoreExpr) -- apply after adding tick (float through)
           -> (CoreExpr -> CoreExpr) -- apply before adding tick (float with)
           -> CoreExpr               -- current expression
@@ -1698,6 +1703,15 @@ altsAreExhaustive (Alt con1 _ _ : alts)
       -- we behave conservatively here -- I don't think it's important
       -- enough to deserve special treatment
 
+-- | Should we look past this tick when eta-expanding the given function?
+--
+-- See Note [Ticks and mandatory eta expansion]
+-- Takes the function we are applying as argument.
+etaExpansionTick :: Id -> GenTickish pass -> Bool
+etaExpansionTick id t
+  = hasNoBinding id &&
+    ( tickishFloatable t || isProfTick t )
+
 {- Note [exprOkForSpeculation: case expressions]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 exprOkForSpeculation accepts very special case expressions.
diff --git a/compiler/GHC/CoreToStg/Prep.hs b/compiler/GHC/CoreToStg/Prep.hs
index fa9496b4c5..ed7b60fdc3 100644
--- a/compiler/GHC/CoreToStg/Prep.hs
+++ b/compiler/GHC/CoreToStg/Prep.hs
@@ -792,7 +792,7 @@ cpeRhsE env (Let bind body)
 
 cpeRhsE env (Tick tickish expr)
   -- Pull out ticks if they are allowed to be floated.
-  | floatableTick tickish
+  | tickishFloatable tickish
   = do { (floats, body) <- cpeRhsE env expr
          -- See [Floating Ticks in CorePrep]
        ; return (unitFloat (FloatTick tickish) `appendFloats` floats, body) }
@@ -1011,10 +1011,12 @@ cpeApp top_env expr
             -- Profiling ticks are slightly less strict so we expand their scope
             -- if they cover partial applications of things like primOps.
             -- See Note [Ticks and mandatory eta expansion]
-            | floatableTick tickish || isProfTick tickish
-            , Var vh <- head
+            -- Here we look inside `fun` before we make the final decision about
+            -- floating the tick which isn't optimal for perf. But this only makes
+            -- a difference if we have a non-floatable tick which is somewhat rare.
+            | Var vh <- head
             , Var head' <- lookupCorePrepEnv top_env vh
-            , hasNoBinding head'
+            , etaExpansionTick head' tickish
             = (head,as')
             where
               (head,as') = go fun (CpeTick tickish : as)
@@ -1145,7 +1147,7 @@ cpeApp top_env expr
           case info of
             CpeCast {} -> go infos n
             CpeTick tickish
-              | floatableTick tickish                 -> go infos n
+              | tickishFloatable tickish                 -> go infos n
               -- If we can't guarantee a tick will be floated out of the application
               -- we can't guarantee the value args following it will be applied.
               | otherwise                             -> n
@@ -2236,11 +2238,6 @@ wrapTicks (Floats flag floats0) expr =
         wrapBind t (NonRec binder rhs) = NonRec binder (mkTick t rhs)
         wrapBind t (Rec pairs)         = Rec (mapSnd (mkTick t) pairs)
 
-floatableTick :: GenTickish pass -> Bool
-floatableTick tickish =
-    tickishPlace tickish == PlaceNonLam &&
-    tickish `tickishScopesLike` SoftScope
-
 ------------------------------------------------------------------------------
 -- Numeric literals
 -- ---------------------------------------------------------------------------