Improve out-of-order inferred type variables

  Don't instantiate type variables for :type in
  `GHC.Tc.Gen.App.tcInstFun`, to avoid inconsistently instantianting
  `r1` but not `r2` in the type

    forall {r1} (a :: TYPE r1) {r2} (b :: TYPE r2). ...

  This fixes #21088.

  This patch also changes the primop pretty-printer to ensure
  that we put all the inferred type variables first. For example,
  the type of reallyUnsafePtrEquality# is now

    forall {l :: Levity} {k :: Levity}
           (a :: TYPE (BoxedRep l))
           (b :: TYPE (BoxedRep k)).
      a -> b -> Int#

  This means we avoid running into issue #21088 entirely with
  the types of primops. Users can still write a type signature where
  the inferred type variables don't come first, however.

  This change to primops had a knock-on consequence, revealing that
  we were sometimes performing eta reduction on keepAlive#.
  This patch updates tryEtaReduce to avoid eta reducing functions
  with no binding, bringing it in line with tryEtaReducePrep,
  and thus fixing #21090.

4 files changed, 42 insertions, 30 deletions

diff --git a/compiler/GHC/Core/Utils.hs b/compiler/GHC/Core/Utils.hs
index baefb7712b..03e2ecee55 100644
--- a/compiler/GHC/Core/Utils.hs
+++ b/compiler/GHC/Core/Utils.hs
@@ -38,7 +38,7 @@ module GHC.Core.Utils (
         diffBinds,
 
         -- * Lambdas and eta reduction
-        tryEtaReduce,
+        tryEtaReduce, canEtaReduceToArity,
 
         -- * Manipulating data constructors and types
         exprToType, exprToCoercion_maybe,
@@ -2448,17 +2448,10 @@ tryEtaReduce bndrs body
     ok_fun _fun                = False
 
     ---------------
-    ok_fun_id fun = -- There are arguments to reduce
+    ok_fun_id fun = -- There are arguments to reduce...
                     fun_arity fun >= incoming_arity &&
-                    -- We always want args for join points so
-                    -- we should never eta-reduce to a trivial expression.
-                    -- See Note [Invariants on join points] in GHC.Core, and #20599
-                    not (isJoinId fun) &&
-                    -- And the function doesn't require visible arguments as part of
-                    -- it's calling convention. See Note [Strict Worker Ids]
-                    idCbvMarkArity fun == 0
-
-
+                    -- ... and the function can be eta reduced to arity 0
+                    canEtaReduceToArity fun 0 0
     ---------------
     fun_arity fun             -- See Note [Arity care]
        | isLocalId fun
@@ -2505,6 +2498,28 @@ tryEtaReduce bndrs body
 
     ok_arg _ _ _ _ = Nothing
 
+-- | Can we eta-reduce the given function to the specified arity?
+-- See Note [Eta reduction conditions].
+canEtaReduceToArity :: Id -> JoinArity -> Arity -> Bool
+canEtaReduceToArity fun dest_join_arity dest_arity =
+  not $
+        hasNoBinding fun
+       -- Don't undersaturate functions with no binding.
+
+    ||  ( isJoinId fun && dest_join_arity < idJoinArity fun )
+       -- Don't undersaturate join points.
+       -- See Note [Invariants on join points] in GHC.Core, and #20599
+
+    || ( dest_arity < idCbvMarkArity fun )
+       -- Don't undersaturate StrictWorkerIds.
+       -- See Note [Strict Worker Ids] in GHC.CoreToStg.Prep.
+
+    ||  isLinearType (idType fun)
+       -- Don't perform eta reduction on linear types.
+       -- If `f :: A %1-> B` and `g :: A -> B`,
+       -- then `g x = f x` is OK but `g = f` is not.
+       -- See Note [Eta reduction conditions].
+
 {-
 Note [Eta reduction of an eval'd function]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/compiler/GHC/CoreToStg/Prep.hs b/compiler/GHC/CoreToStg/Prep.hs
index 2e998bf94e..b8593b47a0 100644
--- a/compiler/GHC/CoreToStg/Prep.hs
+++ b/compiler/GHC/CoreToStg/Prep.hs
@@ -1054,9 +1054,9 @@ cpeApp top_env expr
             (Var f)
             args
         | Just KeepAliveOp <- isPrimOpId_maybe f
-        , CpeApp (Type arg_rep)
-          : CpeApp (Type arg_ty)
+        , CpeApp (Type arg_lev)
           : CpeApp (Type _result_rep)
+          : CpeApp (Type arg_ty)
           : CpeApp (Type result_ty)
           : CpeApp arg
           : CpeApp s0
@@ -1070,13 +1070,14 @@ cpeApp top_env expr
                   _          -> cpe_app env k (CpeApp s0 : rest)
              ; let touchId = primOpId TouchOp
                    expr = Case k' y result_ty [Alt DEFAULT [] rhs]
-                   rhs = let scrut = mkApps (Var touchId) [Type arg_rep, Type arg_ty, arg, Var realWorldPrimId]
+                   rhs = let scrut = mkApps (Var touchId) [Type arg_lev, Type arg_ty, arg, Var realWorldPrimId]
                          in Case scrut s2 result_ty [Alt DEFAULT [] (Var y)]
              ; (floats', expr') <- cpeBody env expr
              ; return (floats `appendFloats` floats', expr')
              }
         | Just KeepAliveOp <- isPrimOpId_maybe f
-        = panic "invalid keepAlive# application"
+        = pprPanic "invalid keepAlive# application" $
+            vcat [ text "args:" <+> ppr args ]
 
     -- runRW# magic
     cpe_app env (Var f) (CpeApp _runtimeRep@Type{} : CpeApp _type@Type{} : CpeApp arg : rest)
@@ -1651,16 +1652,7 @@ tryEtaReducePrep bndrs expr@(App _ _)
     ok bndr (Var arg) = bndr == arg
     ok _    _         = False
 
-    -- We can't eta reduce something which must be saturated.
-    ok_to_eta_reduce (Var f) =  not (hasNoBinding f) &&
-                                not (isLinearType (idType f)) && -- Unsure why this is unsafe.
-                                (not (isJoinId f) || idJoinArity f <= n_remaining) &&
-                                -- Don't undersaturate join points.
-                                -- See Note [Invariants on join points] in GHC.Core, and #20599
-                                (idCbvMarkArity f <= n_remaining_vals)
-                                -- Similar for StrictWorkerIds. See Note [Strict Worker Ids]
-
-
+    ok_to_eta_reduce (Var f) = canEtaReduceToArity f n_remaining n_remaining_vals
     ok_to_eta_reduce _       = False -- Safe. ToDo: generalise
 
 
diff --git a/compiler/GHC/Tc/Gen/App.hs b/compiler/GHC/Tc/Gen/App.hs
index a6e505db96..ddf94f1410 100644
--- a/compiler/GHC/Tc/Gen/App.hs
+++ b/compiler/GHC/Tc/Gen/App.hs
@@ -754,18 +754,23 @@ tcInstFun do_ql inst_final (rn_fun, fun_ctxt) fun_sigma rn_args
           HsUnboundVar {} -> True
           _               -> False
 
-    inst_all :: ArgFlag -> Bool
+    inst_all, inst_inferred, inst_none :: ArgFlag -> Bool
     inst_all (Invisible {}) = True
     inst_all Required       = False
 
-    inst_inferred :: ArgFlag -> Bool
     inst_inferred (Invisible InferredSpec)  = True
     inst_inferred (Invisible SpecifiedSpec) = False
     inst_inferred Required                  = False
 
+    inst_none _ = False
+
     inst_fun :: [HsExprArg 'TcpRn] -> ArgFlag -> Bool
     inst_fun [] | inst_final  = inst_all
-                | otherwise   = inst_inferred
+                | otherwise   = inst_none
+                -- Using `inst_none` for `:type` avoids
+                -- `forall {r1} (a :: TYPE r1) {r2} (b :: TYPE r2). a -> b`
+                -- turning into `forall a {r2} (b :: TYPE r2). a -> b`.
+                -- See #21088.
     inst_fun (EValArg {} : _) = inst_all
     inst_fun _                = inst_inferred
 
diff --git a/compiler/GHC/Tc/Module.hs b/compiler/GHC/Tc/Module.hs
index aa43b7e4e0..aa4be8e76e 100644
--- a/compiler/GHC/Tc/Module.hs
+++ b/compiler/GHC/Tc/Module.hs
@@ -2467,8 +2467,8 @@ tcGhciStmts stmts
 
       ; traceTc "GHC.Tc.Module.tcGhciStmts: done" empty
 
-      -- rec_expr is the expression
-      --      returnIO @ [()] [unsafeCoerce# () x, ..,  unsafeCorece# () z]
+      -- ret_expr is the expression
+      --      returnIO @[()] [unsafeCoerce# () x, ..,  unsafeCoerce# () z]
       --
       -- Despite the inconvenience of building the type applications etc,
       -- this *has* to be done in type-annotated post-typecheck form