Eta reduction with casted function

We want to be able to eta-reduce
   \x y. ((f x) |> co) y
by pushing 'co' inwards.  A very small change accommodates this
See Note [Eta reduction with casted function]

1 files changed, 22 insertions, 7 deletions

diff --git a/compiler/GHC/Core/Utils.hs b/compiler/GHC/Core/Utils.hs
index 8ade2a981a..33999c5070 100644
--- a/compiler/GHC/Core/Utils.hs
+++ b/compiler/GHC/Core/Utils.hs
@@ -2511,11 +2511,19 @@ variable arguments only) thus:
    f @ (g:t1~t2)   |> co  -->  (f |> (t1~t2 => co)) @ (g:t1~t2)
 These are the equations for ok_arg.
 
-It's true that we could also hope to eta reduce these:
+Note [Eta reduction with casted function]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Since we are pushing a coercion inwards, it is easy to accommodate
     (\xy. (f x |> g) y)
     (\xy. (f x y) |> g)
-But the simplifier pushes those casts outwards, so we don't
-need to address that here.
+
+See the `(Cast e co)` equation for `go` in `tryEtaReduce`.  The
+eta-expander pushes those casts outwards, so you might think we won't
+ever see a cast here, but if we have
+  \xy. (f x y |> g)
+we will call tryEtaReduce [x,y] (f x y |> g), and we'd like that to
+work.  This happens in GHC.Core.Opt.Simplify.Utils.mkLam, where
+eta-expansion may be turned off (by sm_eta_expand).
 
 Note [Eta reduction based on evaluation context]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -2544,7 +2552,6 @@ Then this is how the pieces are put together:
     sub-demands `Cn` and see whether all of the n's (here: `S=C_1N` and
     `1=C_11`) were strict. And strict they are! Thus, it will eta-reduce
     `\x y. e x y` to `e`.
-
 -}
 
 -- | `tryEtaReduce [x,y,z] e sd` returns `Just e'` if `\x y z -> e` is evaluated
@@ -2552,8 +2559,7 @@ Then this is how the pieces are put together:
 -- See Note [Eta reduction soundness]
 -- and Note [Eta reduction makes sense] when that is the case.
 tryEtaReduce :: [Var] -> CoreExpr -> SubDemand -> Maybe CoreExpr
--- When updating this function, make sure to update
--- CorePrep.tryEtaReducePrep as well!
+-- Return an expression equal to (\bndrs. body)
 tryEtaReduce bndrs body eval_sd
   = go (reverse bndrs) body (mkRepReflCo (exprType body))
   where
@@ -2565,6 +2571,14 @@ tryEtaReduce bndrs body eval_sd
        -> Maybe CoreExpr   -- Of type a1 -> a2 -> a3 -> ts
     -- See Note [Eta reduction with casted arguments]
     -- for why we have an accumulating coercion
+    --
+    -- Invariant: (go bs body co) returns an expression
+    --            equivalent to (\(reverse bs). body |> co)
+
+    -- See Note [Eta reduction with casted function]
+    go bs (Cast e co1) co2
+      = go bs e (co1 `mkTransCo` co2)
+
     go bs (Tick t e) co
       | tickishFloatable t
       = fmap (Tick t) $ go bs e co
@@ -2629,7 +2643,8 @@ tryEtaReduce bndrs body eval_sd
     ok_arg :: Var              -- Of type bndr_t
            -> CoreExpr         -- Of type arg_t
            -> Coercion         -- Of kind (t1~t2)
-           -> Type             -- Type of the function to which the argument is applied
+           -> Type             -- Type (arg_t -> t1) of the function
+                               --      to which the argument is supplied
            -> Maybe (Coercion  -- Of type (arg_t -> t1 ~  bndr_t -> t2)
                                --   (and similarly for tyvars, coercion args)
                     , [CoreTickish])