Abstract over the right free varswip/T22459

Fix #22459, in two ways:

(1) Make the Specialiser not create a bogus specialisation if
    it is presented by strangely polymorphic dictionary.
    See Note [Weird special case in SpecDict] in
    GHC.Core.Opt.Specialise

(2) Be more careful in abstractFloats
    See Note [Which type variables to abstract over]
    in GHC.Core.Opt.Simplify.Utils.

So (2) stops creating the excessively polymorphic dictionary in
abstractFloats, while (1) stops crashing if some other pass should
nevertheless create a weirdly polymorphic dictionary.

5 files changed, 142 insertions, 53 deletions

diff --git a/compiler/GHC/Core/Opt/Simplify/Utils.hs b/compiler/GHC/Core/Opt/Simplify/Utils.hs
index 3a8a6b4acc..25a7779274 100644
--- a/compiler/GHC/Core/Opt/Simplify/Utils.hs
+++ b/compiler/GHC/Core/Opt/Simplify/Utils.hs
@@ -2063,34 +2063,51 @@ it is guarded by the doFloatFromRhs call in simplLazyBind.
 Note [Which type variables to abstract over]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 Abstract only over the type variables free in the rhs wrt which the
-new binding is abstracted.  Note that
-
-  * The naive approach of abstracting wrt the
-    tyvars free in the Id's /type/ fails. Consider:
-        /\ a b -> let t :: (a,b) = (e1, e2)
-                      x :: a     = fst t
-                  in ...
-    Here, b isn't free in x's type, but we must nevertheless
-    abstract wrt b as well, because t's type mentions b.
-    Since t is floated too, we'd end up with the bogus:
-         poly_t = /\ a b -> (e1, e2)
-         poly_x = /\ a   -> fst (poly_t a *b*)
-
-  * We must do closeOverKinds.  Example (#10934):
+new binding is abstracted.  Several points worth noting
+
+(AB1) The naive approach of abstracting wrt the
+      tyvars free in the Id's /type/ fails. Consider:
+          /\ a b -> let t :: (a,b) = (e1, e2)
+                        x :: a     = fst t
+                    in ...
+      Here, b isn't free in x's type, but we must nevertheless
+      abstract wrt b as well, because t's type mentions b.
+      Since t is floated too, we'd end up with the bogus:
+           poly_t = /\ a b -> (e1, e2)
+           poly_x = /\ a   -> fst (poly_t a *b*)
+
+(AB2) We must do closeOverKinds.  Example (#10934):
        f = /\k (f:k->*) (a:k). let t = AccFailure @ (f a) in ...
-    Here we want to float 't', but we must remember to abstract over
-    'k' as well, even though it is not explicitly mentioned in the RHS,
-    otherwise we get
-       t = /\ (f:k->*) (a:k). AccFailure @ (f a)
-    which is obviously bogus.
-
-  * We get the variables to abstract over by filtering down the
-    the main_tvs for the original function, picking only ones
-    mentioned in the abstracted body. This means:
-    - they are automatically in dependency order, because main_tvs is
-    - there is no issue about non-determinism
-    - we don't gratuitously change order, which may help (in a tiny
-      way) with CSE and/or the compiler-debugging experience
+      Here we want to float 't', but we must remember to abstract over
+      'k' as well, even though it is not explicitly mentioned in the RHS,
+      otherwise we get
+         t = /\ (f:k->*) (a:k). AccFailure @ (f a)
+      which is obviously bogus.
+
+(AB3) We get the variables to abstract over by filtering down the
+      the main_tvs for the original function, picking only ones
+      mentioned in the abstracted body. This means:
+      - they are automatically in dependency order, because main_tvs is
+      - there is no issue about non-determinism
+      - we don't gratuitously change order, which may help (in a tiny
+        way) with CSE and/or the compiler-debugging experience
+
+(AB4) For a recursive group, it's a bit of a pain to work out the minimal
+      set of tyvars over which to abstract:
+           /\ a b c.  let x = ...a... in
+                      letrec { p = ...x...q...
+                               q = .....p...b... } in
+                      ...
+      Since 'x' is abstracted over 'a', the {p,q} group must be abstracted
+      over 'a' (because x is replaced by (poly_x a)) as well as 'b'.
+      Remember this bizarre case too:
+           x::a = x
+      Here, we must abstract 'x' over 'a'.
+
+      Why is it worth doing this?  Partly tidiness; and partly #22459
+      which showed that it's harder to do polymorphic specialisation well
+      if there are dictionaries abstracted over unnecessary type variables.
+      See Note [Weird special case for SpecDict] in GHC.Core.Opt.Specialise
 -}
 
 abstractFloats :: UnfoldingOpts -> TopLevelFlag -> [OutTyVar] -> SimplFloats
@@ -2115,33 +2132,40 @@ abstractFloats uf_opts top_lvl main_tvs floats body
         rhs' = GHC.Core.Subst.substExpr subst rhs
 
         -- tvs_here: see Note [Which type variables to abstract over]
-        tvs_here = filter (`elemVarSet` free_tvs) main_tvs
-        free_tvs = closeOverKinds $
-                   exprSomeFreeVars isTyVar rhs'
+        tvs_here = choose_tvs (exprSomeFreeVars isTyVar rhs')
 
     abstract subst (Rec prs)
-       = do { (poly_ids, poly_apps) <- mapAndUnzipM (mk_poly1 tvs_here) ids
-            ; let subst' = GHC.Core.Subst.extendSubstList subst (ids `zip` poly_apps)
-                  poly_pairs = [ mk_poly2 poly_id tvs_here rhs'
-                               | (poly_id, rhs) <- poly_ids `zip` rhss
-                               , let rhs' = GHC.Core.Subst.substExpr subst' rhs ]
-            ; return (subst', Rec poly_pairs) }
+      = do { (poly_ids, poly_apps) <- mapAndUnzipM (mk_poly1 tvs_here) ids
+           ; let subst' = GHC.Core.Subst.extendSubstList subst (ids `zip` poly_apps)
+                 poly_pairs = [ mk_poly2 poly_id tvs_here rhs'
+                              | (poly_id, rhs) <- poly_ids `zip` rhss
+                              , let rhs' = GHC.Core.Subst.substExpr subst' rhs ]
+           ; return (subst', Rec poly_pairs) }
+      where
+        (ids,rhss) = unzip prs
+
+
+        -- tvs_here: see Note [Which type variables to abstract over]
+        tvs_here = choose_tvs (mapUnionVarSet get_bind_fvs prs)
+
+        -- See wrinkle (AB4) in Note [Which type variables to abstract over]
+        get_bind_fvs (id,rhs) = tyCoVarsOfType (idType id) `unionVarSet` get_rec_rhs_tvs rhs
+        get_rec_rhs_tvs rhs   = nonDetStrictFoldVarSet get_tvs emptyVarSet (exprFreeVars rhs)
+
+        get_tvs :: Var -> VarSet -> VarSet
+        get_tvs var free_tvs
+           | isTyVar var      -- CoVars have been substituted away
+           = extendVarSet free_tvs var
+           | Just poly_app <- GHC.Core.Subst.lookupIdSubst_maybe subst var
+           = -- 'var' is like 'x' in (AB4)
+             exprSomeFreeVars isTyVar poly_app `unionVarSet` free_tvs
+           | otherwise
+           = free_tvs
+
+    choose_tvs free_tvs
+       = filter (`elemVarSet` all_free_tvs) main_tvs  -- (AB3)
        where
-         (ids,rhss) = unzip prs
-                -- For a recursive group, it's a bit of a pain to work out the minimal
-                -- set of tyvars over which to abstract:
-                --      /\ a b c.  let x = ...a... in
-                --                 letrec { p = ...x...q...
-                --                          q = .....p...b... } in
-                --                 ...
-                -- Since 'x' is abstracted over 'a', the {p,q} group must be abstracted
-                -- over 'a' (because x is replaced by (poly_x a)) as well as 'b'.
-                -- Since it's a pain, we just use the whole set, which is always safe
-                --
-                -- If you ever want to be more selective, remember this bizarre case too:
-                --      x::a = x
-                -- Here, we must abstract 'x' over 'a'.
-         tvs_here = scopedSort main_tvs
+         all_free_tvs = closeOverKinds free_tvs       -- (AB2)
 
     mk_poly1 :: [TyVar] -> Id -> SimplM (Id, CoreExpr)
     mk_poly1 tvs_here var
diff --git a/compiler/GHC/Core/Opt/Specialise.hs b/compiler/GHC/Core/Opt/Specialise.hs
index ac4934edbf..dda10da34e 100644
--- a/compiler/GHC/Core/Opt/Specialise.hs
+++ b/compiler/GHC/Core/Opt/Specialise.hs
@@ -2516,6 +2516,8 @@ specHeader env (bndr : bndrs) (UnspecType : args)
 -- the nitty-gritty), as a LHS rule and unfolding details.
 specHeader env (bndr : bndrs) (SpecDict d : args)
   | not (isDeadBinder bndr)
+  , allVarSet (`elemInScopeSet` in_scope) (exprFreeVars d)
+    -- See Note [Weird special case for SpecDict]
   = do { (env1, bndr') <- newDictBndr env bndr -- See Note [Zap occ info in rule binders]
        ; let (env2, dx_bind, spec_dict) = bindAuxiliaryDict env1 bndr bndr' d
        ; (_, env3, leftover_bndrs, rule_bs, rule_es, bs', dx, spec_args)
@@ -2531,6 +2533,8 @@ specHeader env (bndr : bndrs) (SpecDict d : args)
               , spec_dict : spec_args
               )
        }
+   where
+     in_scope = Core.getSubstInScope (se_subst env)
 
 -- Finally, we don't want to specialise on this argument 'i':
 --   - It's an UnSpecArg, or
@@ -2752,6 +2756,8 @@ monomorpic, and specialised in one go.
 
 Wrinkles.
 
+* See Note [Weird special case for SpecDict]
+
 * With -XOverlappingInstances you might worry about this:
     class C a where ...
     instance C (Maybe Int) where ...   -- $df1 :: C (Maybe Int)
@@ -2777,6 +2783,33 @@ Wrinkles.
   it's a hard test to make.)
 
   But see Note [Specialisation and overlapping instances].
+
+Note [Weird special case for SpecDict]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Suppose we are trying to specialise for this this call:
+   $wsplit @T (mkD @k @(a::k) :: C T)
+where
+   mkD :: forall k (a::k). C T
+is a top-level dictionary-former.  This actually happened in #22459,
+because of (MP1) of Note [Specialising polymorphic dictionaries].
+
+How can we speicalise $wsplit?  We might try
+
+   RULE "SPEC" forall (d :: C T). $wsplit @T d = $s$wsplit
+
+but then in the body of $s$wsplit what will we use for the dictionary
+evidence?  We can't use (mkD @k @(a::k)) because k and a aren't in scope.
+We could zap `k` to (Any @Type) and `a` to (Any @(Any @Type)), but that
+is a lot of hard work for a very strange case.
+
+So we simply refrain from specialising in this case; hence the guard
+   allVarSet (`elemInScopeSet` in_scope) (exprFreeVars d)
+in the SpecDict cased of specHeader.
+
+How did this strange polymorphic mkD arise in the first place?
+From GHC.Core.Opt.Utils.abstractFloats, which was abstracting
+over too many type variables. But that too is now fixed;
+see Note [Which type variables to abstract over] in that module.
 -}
 
 instance Outputable DictBind where
diff --git a/compiler/GHC/Core/Subst.hs b/compiler/GHC/Core/Subst.hs
index 3d52c655f6..f0ad737fb6 100644
--- a/compiler/GHC/Core/Subst.hs
+++ b/compiler/GHC/Core/Subst.hs
@@ -15,7 +15,7 @@ module GHC.Core.Subst (
         deShadowBinds, substRuleInfo, substRulesForImportedIds,
         substTyUnchecked, substCo, substExpr, substExprSC, substBind, substBindSC,
         substUnfolding, substUnfoldingSC,
-        lookupIdSubst, substIdType, substIdOcc,
+        lookupIdSubst, lookupIdSubst_maybe, substIdType, substIdOcc,
         substTickish, substDVarSet, substIdInfo,
 
         -- ** Operations on substitutions
@@ -184,9 +184,11 @@ extendSubstList subst []              = subst
 extendSubstList subst ((var,rhs):prs) = extendSubstList (extendSubst subst var rhs) prs
 
 -- | Find the substitution for an 'Id' in the 'Subst'
+-- The Id should not be a CoVar
 lookupIdSubst :: HasDebugCallStack => Subst -> Id -> CoreExpr
 lookupIdSubst (Subst in_scope ids _ _) v
-  | not (isLocalId v) = Var v
+  | assertPpr (isId v && not (isCoVar v)) (ppr v)
+    not (isLocalId v)                   = Var v
   | Just e  <- lookupVarEnv ids       v = e
   | Just v' <- lookupInScope in_scope v = Var v'
         -- Vital! See Note [Extending the IdSubstEnv]
@@ -194,6 +196,12 @@ lookupIdSubst (Subst in_scope ids _ _) v
         -- it's a bad bug and we really want to know
   | otherwise = pprPanic "lookupIdSubst" (ppr v $$ ppr in_scope)
 
+lookupIdSubst_maybe :: HasDebugCallStack => Subst -> Id -> Maybe CoreExpr
+-- Just look up in the substitution; do not check the in-scope set
+lookupIdSubst_maybe (Subst _ ids _ _) v
+  = assertPpr (isId v && not (isCoVar v)) (ppr v) $
+    lookupVarEnv ids v
+
 delBndr :: Subst -> Var -> Subst
 delBndr (Subst in_scope ids tvs cvs) v
   | isCoVar v = Subst in_scope ids tvs (delVarEnv cvs v)
diff --git a/testsuite/tests/simplCore/should_compile/T22459.hs b/testsuite/tests/simplCore/should_compile/T22459.hs
new file mode 100644
index 0000000000..be28c5e2c0
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T22459.hs
@@ -0,0 +1,23 @@
+{-# LANGUAGE UndecidableInstances #-}
+
+{-# OPTIONS_GHC -O #-}
+
+module Lib (foo) where
+
+import qualified Data.Map as M
+
+newtype Fix f = Fix (f (Fix f))
+
+instance Eq (f (Fix f)) => Eq (Fix f) where
+  Fix a == Fix b = a == b
+
+instance Ord (f (Fix f)) => Ord (Fix f) where
+  Fix a `compare` Fix b = a `compare` b
+
+data Foo i r = Foo i r
+  deriving (Eq, Ord)
+
+newtype Bar a = Bar (M.Map Char (M.Map (Fix (Foo ())) Word))
+
+foo :: Bar a -> Bar a -> Bar a
+foo (Bar a) (Bar b) = Bar (M.unionWith M.union a b)
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index 5209c0c7cb..df60aa1c4f 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -458,3 +458,4 @@ test('T22494', [grep_errmsg(r'case') ], compile, ['-O -ddump-simpl -dsuppress-un
 test('T22491', normal, compile, ['-O2'])
 test('T21476', normal, compile, [''])
 test('T22272', normal, multimod_compile, ['T22272', '-O -fexpose-all-unfoldings -fno-omit-interface-pragmas -fno-ignore-interface-pragmas'])
+test('T22459', normal, compile, [''])