[commit: ghc] wip/spj-early-inline2: Make Specialise work with casts (fe912e0)

[git at git.haskell.org]

Repository : ssh://git@git.haskell.org/ghc

On branch  : wip/spj-early-inline2
Link       : http://ghc.haskell.org/trac/ghc/changeset/fe912e0dc9d848162600cc4269b4943331ea28c3/ghc

>---------------------------------------------------------------

commit fe912e0dc9d848162600cc4269b4943331ea28c3
Author: Simon Peyton Jones <simonpj at microsoft.com>
Date:   Tue Feb 14 13:47:13 2017 +0000

    Make Specialise work with casts
    
    With my upcoming early-inlining patch it turned out that Specialise
    was getting stuck on casts.  This patch fixes it; see Specialise
    [Account for casts in binding].


>---------------------------------------------------------------

fe912e0dc9d848162600cc4269b4943331ea28c3
 compiler/coreSyn/CoreSubst.hs     | 49 +++++++++++++++++++++++++++++++++++++--
 compiler/specialise/Specialise.hs | 35 +++++++++++++++++++++-------
 2 files changed, 73 insertions(+), 11 deletions(-)

diff --git a/compiler/coreSyn/CoreSubst.hs b/compiler/coreSyn/CoreSubst.hs
index 89a92f8..6afa3ba 100644
--- a/compiler/coreSyn/CoreSubst.hs
+++ b/compiler/coreSyn/CoreSubst.hs
@@ -34,7 +34,7 @@ module CoreSubst (
         -- ** Simple expression optimiser
         simpleOptPgm, simpleOptExpr, simpleOptExprWith,
         exprIsConApp_maybe, exprIsLiteral_maybe, exprIsLambda_maybe,
-        pushCoArg, pushCoValArg, pushCoTyArg
+        pushCoArg, pushCoValArg, pushCoTyArg, collectBindersPushingCo
     ) where
 
 #include "HsVersions.h"
@@ -1605,7 +1605,7 @@ exprIsLambda_maybe _ _e
 
 Here we implement the "push rules" from FC papers:
 
-* The push-argument ules, where we can move a coercion past an argument.
+* The push-argument rules, where we can move a coercion past an argument.
   We have
       (fun |> co) arg
   and we want to transform it to
@@ -1775,3 +1775,48 @@ pushCoDataCon dc dc_args co
 
   where
     Pair from_ty to_ty = coercionKind co
+
+collectBindersPushingCo :: CoreExpr -> ([Var], CoreExpr)
+-- Collect lambda binders, pushing coercions inside if possible
+-- E..g   (\x.e) |> g         g :: <Int> -> blah
+--        = (\x. e |> Nth 1 g)
+collectBindersPushingCo e
+  = go [] e
+  where
+    go :: [Var] -> CoreExpr -> ([Var], CoreExpr)
+    -- The accumulator is in reverse order
+    go bs (Lam b e)   = go (b:bs) e
+    go bs (Cast e co) = go_c bs e co
+    go bs e           = (reverse bs, e)
+
+    go_c :: [Var] -> CoreExpr -> Coercion -> ([Var], CoreExpr)
+    -- (go_c bs e c) is same as (go bs e (e |> c))
+    go_c bs (Cast e co1) co2 = go_c bs e (co1 `mkTransCo` co2)
+    go_c bs (Lam b e)    co  = go_lam bs b e co
+    go_c bs e            co  = (reverse bs, mkCast e co)
+
+    go_lam :: [Var] -> Var -> CoreExpr -> Coercion -> ([Var], CoreExpr)
+    -- (go_lam bs b e c) is same as (go_c bs (\b.e) c)
+    go_lam bs b e co
+      | isTyVar b
+      , let Pair tyL tyR = coercionKind co
+      , ASSERT( isForAllTy tyL )
+        isForAllTy tyR
+      , isReflCo (mkNthCo 0 co)  -- See Note [collectBindersPushingCo]
+      = go_c (b:bs) e (mkInstCo co (mkNomReflCo (mkTyVarTy b)))
+
+      | isId b
+      , let Pair tyL tyR = coercionKind co
+      , ASSERT( isFunTy tyL) isFunTy tyR
+      , isReflCo (mkNthCo 0 co)  -- See Note [collectBindersPushingCo]
+      = go_c (b:bs) e (mkNthCo 1 co)
+
+      | otherwise = (reverse bs, mkCast (Lam b e) co)
+
+{- Note [collectBindersPushingCo]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+We just look for coercions of form
+   <type> -> blah
+(and similarly for foralls) to keep this function simple.  We could do
+more elaborate stuff, but it'd involve substitution etc.
+-}
diff --git a/compiler/specialise/Specialise.hs b/compiler/specialise/Specialise.hs
index 3f937c2..fb3c7cc 100644
--- a/compiler/specialise/Specialise.hs
+++ b/compiler/specialise/Specialise.hs
@@ -1153,8 +1153,8 @@ specCalls :: Maybe Module      -- Just this_mod  =>  specialising imported fn
 
 specCalls mb_mod env rules_for_me calls_for_me fn rhs
         -- The first case is the interesting one
-  |  rhs_tyvars `lengthIs`     n_tyvars -- Rhs of fn's defn has right number of big lambdas
-  && rhs_ids    `lengthAtLeast` n_dicts -- and enough dict args
+  |  rhs_tyvars `lengthIs`      n_tyvars -- Rhs of fn's defn has right number of big lambdas
+  && rhs_bndrs1 `lengthAtLeast` n_dicts -- and enough dict args
   && notNull calls_for_me               -- And there are some calls to specialise
   && not (isNeverActive (idInlineActivation fn))
         -- Don't specialise NOINLINE things
@@ -1178,7 +1178,7 @@ specCalls mb_mod env rules_for_me calls_for_me fn rhs
     return ([], [], emptyUDs)
   where
     _trace_doc = sep [ ppr rhs_tyvars, ppr n_tyvars
-                     , ppr rhs_ids, ppr n_dicts
+                     , ppr rhs_bndrs, ppr n_dicts
                      , ppr (idInlineActivation fn) ]
 
     fn_type                 = idType fn
@@ -1194,11 +1194,12 @@ specCalls mb_mod env rules_for_me calls_for_me fn rhs
         -- Figure out whether the function has an INLINE pragma
         -- See Note [Inline specialisations]
 
-    (rhs_tyvars, rhs_ids, rhs_body) = collectTyAndValBinders rhs
-
-    rhs_dict_ids = take n_dicts rhs_ids
-    body         = mkLams (drop n_dicts rhs_ids) rhs_body
-                -- Glue back on the non-dict lambdas
+    (rhs_bndrs, rhs_body)      = CoreSubst.collectBindersPushingCo rhs
+                                 -- See Note [Account for casts in binding]
+    (rhs_tyvars, rhs_bndrs1)   = span isTyVar rhs_bndrs
+    (rhs_dict_ids, rhs_bndrs2) = splitAt n_dicts rhs_bndrs1
+    body                       = mkLams rhs_bndrs2 rhs_body
+                                 -- Glue back on the non-dict lambdas
 
     already_covered :: DynFlags -> [CoreExpr] -> Bool
     already_covered dflags args      -- Note [Specialisations already covered]
@@ -1350,7 +1351,23 @@ specCalls mb_mod env rules_for_me calls_for_me fn rhs
 
            ; return (Just ((spec_f_w_arity, spec_rhs), final_uds, spec_env_rule)) } }
 
-{- Note [Evidence foralls]
+{- Note [Account for casts in binding]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Consider
+   f :: Eq a => a -> IO ()
+   {-# INLINABLE f
+       StableUnf = (/\a \(d:Eq a) (x:a). blah) |> g
+     #-}
+   f = ...
+
+In f's stable unfolding we have done some modest simplification which
+has pushed the cast to the outside.  (I wonder if this is the Right
+Thing, but it's what happens now; see SimplUtils Note [Casts and
+lambdas].)  Now that stable unfolding must be specialised, so we want
+to push the cast back inside. It would be terrible if the cast
+defeated specialisation!  Hence the use of collectBindersPushingCo.
+
+Note [Evidence foralls]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~
 Suppose (Trac #12212) that we are specialising
    f :: forall a b. (Num a, F a ~ F b) => blah