Fix shadowing in mkWwBodies

This bug, exposed by Trac #12562 was very obscure, and has been
lurking for a long time.  What happened was that, in the
worker/wrapper split

  a tyvar binder for a worker function
  accidentally shadowed an in-scope term variable
  that was mentioned in the body of the function

It's jolly hard to provoke, so I have not even attempted to make
a test case.  There's a Note [Freshen WW arguments] to explain.

Interestingly, fixing the bug (which meant fresher type variables)
revealed a second lurking bug: I'd failed to apply the substitution to
the coercion in the second last case of mkWWArgs, which introduces a
Cast.

2 files changed, 55 insertions, 27 deletions

diff --git a/compiler/stranal/WorkWrap.hs b/compiler/stranal/WorkWrap.hs
index 80d966b392..9acc461c20 100644
--- a/compiler/stranal/WorkWrap.hs
+++ b/compiler/stranal/WorkWrap.hs
@@ -10,6 +10,7 @@ module WorkWrap ( wwTopBinds ) where
 import CoreSyn
 import CoreUnfold       ( certainlyWillInline, mkWwInlineRule, mkWorkerUnfolding )
 import CoreUtils        ( exprType, exprIsHNF )
+import CoreFVs          ( exprFreeVars )
 import Var
 import Id
 import IdInfo
@@ -365,7 +366,7 @@ splitFun :: DynFlags -> FamInstEnvs -> Id -> IdInfo -> [Demand] -> DmdResult ->
 splitFun dflags fam_envs fn_id fn_info wrap_dmds res_info rhs
   = WARN( not (wrap_dmds `lengthIs` arity), ppr fn_id <+> (ppr arity $$ ppr wrap_dmds $$ ppr res_info) ) do
     -- The arity should match the signature
-    stuff <- mkWwBodies dflags fam_envs fun_ty wrap_dmds res_info
+    stuff <- mkWwBodies dflags fam_envs rhs_fvs fun_ty wrap_dmds res_info
     case stuff of
       Just (work_demands, wrap_fn, work_fn) -> do
         work_uniq <- getUniqueM
@@ -432,6 +433,7 @@ splitFun dflags fam_envs fn_id fn_info wrap_dmds res_info rhs
 
       Nothing -> return [(fn_id, rhs)]
   where
+    rhs_fvs         = exprFreeVars rhs
     fun_ty          = idType fn_id
     inl_prag        = inlinePragInfo fn_info
     rule_match_info = inlinePragmaRuleMatchInfo inl_prag
diff --git a/compiler/stranal/WwLib.hs b/compiler/stranal/WwLib.hs
index 64de0e0078..1370bbce06 100644
--- a/compiler/stranal/WwLib.hs
+++ b/compiler/stranal/WwLib.hs
@@ -24,6 +24,7 @@ import MkId             ( voidArgId, voidPrimId )
 import TysPrim          ( voidPrimTy )
 import TysWiredIn       ( tupleDataCon )
 import VarEnv           ( mkInScopeSet )
+import VarSet           ( VarSet )
 import Type
 import RepType          ( isVoidTy )
 import Coercion
@@ -109,14 +110,19 @@ the unusable strictness-info into the interfaces.
 @mkWwBodies@ is called when doing the worker\/wrapper split inside a module.
 -}
 
+type WwResult
+  = ([Demand],              -- Demands for worker (value) args
+     Id -> CoreExpr,        -- Wrapper body, lacking only the worker Id
+     CoreExpr -> CoreExpr)  -- Worker body, lacking the original function rhs
+
 mkWwBodies :: DynFlags
            -> FamInstEnvs
-           -> Type                                  -- Type of original function
-           -> [Demand]                              -- Strictness of original function
-           -> DmdResult                             -- Info about function result
-           -> UniqSM (Maybe ([Demand],              -- Demands for worker (value) args
-                             Id -> CoreExpr,        -- Wrapper body, lacking only the worker Id
-                             CoreExpr -> CoreExpr)) -- Worker body, lacking the original function rhs
+           -> VarSet         -- Free vars of RHS
+                             -- See Note [Freshen WW arguments]
+           -> Type           -- Type of original function
+           -> [Demand]       -- Strictness of original function
+           -> DmdResult      -- Info about function result
+           -> UniqSM (Maybe WwResult)
 
 -- wrap_fn_args E       = \x y -> E
 -- work_fn_args E       = E x y
@@ -129,8 +135,9 @@ mkWwBodies :: DynFlags
 --                        let x = (a,b) in
 --                        E
 
-mkWwBodies dflags fam_envs fun_ty demands res_info
-  = do  { let empty_subst = mkEmptyTCvSubst (mkInScopeSet (tyCoVarsOfType fun_ty))
+mkWwBodies dflags fam_envs rhs_fvs fun_ty demands res_info
+  = do  { let empty_subst = mkEmptyTCvSubst (mkInScopeSet rhs_fvs)
+                -- See Note [Freshen WW arguments]
 
         ; (wrap_args, wrap_fn_args, work_fn_args, res_ty) <- mkWWargs empty_subst fun_ty demands
         ; (useful1, work_args, wrap_fn_str, work_fn_str) <- mkWWstr dflags fam_envs wrap_args
@@ -296,7 +303,7 @@ the \x to get what we want.
 -- and keeps repeating that until it's satisfied the supplied arity
 
 mkWWargs :: TCvSubst            -- Freshening substitution to apply to the type
-                                --   See Note [Freshen type variables]
+                                --   See Note [Freshen WW arguments]
          -> Type                -- The type of the function
          -> [Demand]     -- Demands and one-shot info for value arguments
          -> UniqSM  ([Var],            -- Wrapper args
@@ -321,9 +328,9 @@ mkWWargs subst fun_ty demands
                   res_ty) }
 
   | Just (tv, fun_ty') <- splitForAllTy_maybe fun_ty
-  = do  { let (subst', tv') = substTyVarBndr subst tv
-                -- This substTyVarBndr clones the type variable when necy
-                -- See Note [Freshen type variables]
+  = do  { uniq <- getUniqueM
+        ; let (subst', tv') = cloneTyVarBndr subst tv uniq
+                -- See Note [Freshen WW arguments]
         ; (wrap_args, wrap_fn_args, work_fn_args, res_ty)
              <- mkWWargs subst' fun_ty' demands
         ; return (tv' : wrap_args,
@@ -342,9 +349,10 @@ mkWWargs subst fun_ty demands
 
   = do { (wrap_args, wrap_fn_args, work_fn_args, res_ty)
             <-  mkWWargs subst rep_ty demands
-        ; return (wrap_args,
-                  \e -> Cast (wrap_fn_args e) (mkSymCo co),
-                  \e -> work_fn_args (Cast e co),
+       ; let co' = substCo subst co
+       ; return (wrap_args,
+                  \e -> Cast (wrap_fn_args e) (mkSymCo co'),
+                  \e -> work_fn_args (Cast e co'),
                   res_ty) }
 
   | otherwise
@@ -359,17 +367,35 @@ mk_wrap_arg uniq ty dmd
   = mkSysLocalOrCoVar (fsLit "w") uniq ty
        `setIdDemandInfo` dmd
 
-{-
-Note [Freshen type variables]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-Wen we do a worker/wrapper split, we must not use shadowed names,
-else we'll get
-   f = /\ a /\a. fw a a
-which is obviously wrong.  Type variables can can in principle shadow,
-within a type (e.g. forall a. a -> forall a. a->a).  But type
-variables *are* mentioned in <blah>, so we must substitute.
-
-That's why we carry the TCvSubst through mkWWargs
+{- Note [Freshen WW arguments]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Wen we do a worker/wrapper split, we must not in-scope names as the arguments
+of the worker, else we'll get name capture.  E.g.
+
+   -- y1 is in scope from further out
+   f x = ..y1..
+
+If we accidentally choose y1 as a worker argument disaster results:
+
+   fww y1 y2 = let x = (y1,y2) in ...y1...
+
+To avoid this:
+
+  * We use a fresh unique for both type-variable and term-variable binders
+    Originally we lacked this freshness for type variables, and that led
+    to the very obscure Trac #12562.  (A type varaible in the worker shadowed
+    an outer term-variable binding.)
+
+  * Because of this cloning we have to substitute in the type/kind of the
+    new binders.  That's why we carry the TCvSubst through mkWWargs.
+
+    So we need a decent in-scope set, just in case that type/kind
+    itself has foralls.  We get this from the free vars of the RHS of the
+    function since those are the only variables that might be captured.
+    It's a lazy thunk, which will only be poked if the type/kind has a forall.
+
+    Another tricky case was when f :: forall a. a -> forall a. a->a
+    (i.e. with shadowing), and then the worker used the same 'a' twice.
 
 ************************************************************************
 *                                                                      *