Take account of loop breakers in specLookupRulewip/T22802

The key change is that in GHC.Core.Opt.Specialise.specLookupRule
we were using realIdUnfolding, which ignores the loop-breaker
flag.  When given a loop breaker, rule matching therefore
looped infinitely -- #22802.

In fixing this I refactored a bit.

* Define GHC.Core.InScopeEnv as a data type, and use it.
  (Previously it was a pair: hard to grep for.)

* Put several functions returning an IdUnfoldingFun into
  GHC.Types.Id, namely
     idUnfolding
     alwaysActiveUnfoldingFun,
     whenActiveUnfoldingFun,
     noUnfoldingFun
  and use them.  (The are all loop-breaker aware.)

11 files changed, 90 insertions, 49 deletions

diff --git a/compiler/GHC/Core.hs b/compiler/GHC/Core.hs
index 92b34ffc21..f504a7cbd5 100644
--- a/compiler/GHC/Core.hs
+++ b/compiler/GHC/Core.hs
@@ -83,7 +83,7 @@ module GHC.Core (
 
         -- * Core rule data types
         CoreRule(..),
-        RuleName, RuleFun, IdUnfoldingFun, InScopeEnv, RuleOpts,
+        RuleName, RuleFun, IdUnfoldingFun, InScopeEnv(..), RuleOpts,
 
         -- ** Operations on 'CoreRule's
         ruleArity, ruleName, ruleIdName, ruleActivation,
@@ -1171,10 +1171,11 @@ data CoreRule
     }
                 -- See Note [Extra args in the target] in GHC.Core.Rules
 
+type RuleFun = RuleOpts -> InScopeEnv -> Id -> [CoreExpr] -> Maybe CoreExpr
+
 -- | The 'InScopeSet' in the 'InScopeEnv' is a /superset/ of variables that are
 -- currently in scope. See Note [The InScopeSet invariant].
-type RuleFun = RuleOpts -> InScopeEnv -> Id -> [CoreExpr] -> Maybe CoreExpr
-type InScopeEnv = (InScopeSet, IdUnfoldingFun)
+data InScopeEnv = ISE InScopeSet IdUnfoldingFun
 
 type IdUnfoldingFun = Id -> Unfolding
 -- A function that embodies how to unfold an Id if you need
diff --git a/compiler/GHC/Core/Opt/ConstantFold.hs b/compiler/GHC/Core/Opt/ConstantFold.hs
index 3d36368d5b..7ace3124e9 100644
--- a/compiler/GHC/Core/Opt/ConstantFold.hs
+++ b/compiler/GHC/Core/Opt/ConstantFold.hs
@@ -2402,7 +2402,7 @@ match_cstring_foldr_lit _ _ _ _ _ = Nothing
 -- Also, look into variable's unfolding just in case the expression we look for
 -- is in a top-level thunk.
 stripStrTopTicks :: InScopeEnv -> CoreExpr -> ([CoreTickish], CoreExpr)
-stripStrTopTicks (_,id_unf) e = case e of
+stripStrTopTicks (ISE _ id_unf) e = case e of
   Var v
     | Just rhs <- expandUnfolding_maybe (id_unf v)
     -> stripTicksTop tickishFloatable rhs
diff --git a/compiler/GHC/Core/Opt/Simplify/Utils.hs b/compiler/GHC/Core/Opt/Simplify/Utils.hs
index 28b1ebc221..19daab0075 100644
--- a/compiler/GHC/Core/Opt/Simplify/Utils.hs
+++ b/compiler/GHC/Core/Opt/Simplify/Utils.hs
@@ -1241,14 +1241,13 @@ getUnfoldingInRuleMatch :: SimplEnv -> InScopeEnv
 -- is 'otherwise' which we want exprIsConApp_maybe to be able to
 -- see very early on
 getUnfoldingInRuleMatch env
-  = (in_scope, id_unf)
+  = ISE in_scope id_unf
   where
     in_scope = seInScope env
-    id_unf id | unf_is_active id = idUnfolding id
-              | otherwise        = NoUnfolding
-    unf_is_active id = isActive (sePhase env) (idInlineActivation id)
-       -- When sm_rules was off we used to test for a /stable/ unfolding,
-       -- but that seems wrong (#20941)
+    phase    = sePhase env
+    id_unf   = whenActiveUnfoldingFun (isActive phase)
+     -- When sm_rules was off we used to test for a /stable/ unfolding,
+     -- but that seems wrong (#20941)
 
 ----------------------
 activeRule :: SimplMode -> Activation -> Bool
diff --git a/compiler/GHC/Core/Opt/Specialise.hs b/compiler/GHC/Core/Opt/Specialise.hs
index fa9323ab3b..13bff9f170 100644
--- a/compiler/GHC/Core/Opt/Specialise.hs
+++ b/compiler/GHC/Core/Opt/Specialise.hs
@@ -1626,11 +1626,11 @@ specCalls spec_imp env existing_rules calls_for_me fn rhs
 --      See Note [Inline specialisations] for why we do not
 --      switch off specialisation for inline functions
 
-  = do { -- debugTraceMsg (text "specCalls: some" <+> vcat
-         --   [ text "function" <+> ppr fn
-         --   , text "calls:" <+> ppr calls_for_me
-         --   , text "subst" <+> ppr (se_subst env) ])
-       ; foldlM spec_call ([], [], emptyUDs) calls_for_me }
+  = -- pprTrace "specCalls: some" (vcat
+    --   [ text "function" <+> ppr fn
+    --   , text "calls:" <+> ppr calls_for_me
+    --   , text "subst" <+> ppr (se_subst env) ]) $
+    foldlM spec_call ([], [], emptyUDs) calls_for_me
 
   | otherwise   -- No calls or RHS doesn't fit our preconceptions
   = warnPprTrace (not (exprIsTrivial rhs) && notNull calls_for_me && not (isClassOpId fn))
@@ -1685,7 +1685,7 @@ specCalls spec_imp env existing_rules calls_for_me fn rhs
              , rule_bndrs, rule_lhs_args
              , spec_bndrs1, dx_binds, spec_args) <- specHeader env rhs_bndrs all_call_args
 
---           ; debugTraceMsg (text "spec_call" <+> vcat
+--           ; pprTrace "spec_call" (vcat
 --                [ text "fun:       "  <+> ppr fn
 --                , text "call info: "  <+> ppr _ci
 --                , text "useful:    "  <+> ppr useful
@@ -1698,7 +1698,8 @@ specCalls spec_imp env existing_rules calls_for_me fn rhs
 --                , text "rhs_bndrs"     <+> ppr rhs_bndrs
 --                , text "rhs_body"     <+> ppr rhs_body
 --                , text "rhs_env2:  "  <+> ppr (se_subst rhs_env2)
---                , ppr dx_binds ]
+--                , ppr dx_binds ]) $
+--             return ()
 
            ; if not useful  -- No useful specialisation
                 || already_covered rhs_env2 rules_acc rule_lhs_args
@@ -1795,12 +1796,13 @@ specLookupRule :: SpecEnv -> Id -> [CoreExpr]
                -> CompilerPhase  -- Look up rules as if we were in this phase
                -> [CoreRule] -> Maybe (CoreRule, CoreExpr)
 specLookupRule env fn args phase rules
-  = lookupRule ropts (in_scope, realIdUnfolding) is_active fn args rules
+  = lookupRule ropts in_scope_env is_active fn args rules
   where
-    dflags    = se_dflags env
-    in_scope  = getSubstInScope (se_subst env)
-    ropts     = initRuleOpts dflags
-    is_active = isActive phase
+    dflags       = se_dflags env
+    in_scope     = getSubstInScope (se_subst env)
+    in_scope_env = ISE in_scope (whenActiveUnfoldingFun is_active)
+    ropts        = initRuleOpts dflags
+    is_active    = isActive phase
 
 {- Note [Specialising DFuns]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/compiler/GHC/Core/Rules.hs b/compiler/GHC/Core/Rules.hs
index d9bd0a912c..d635d6aebe 100644
--- a/compiler/GHC/Core/Rules.hs
+++ b/compiler/GHC/Core/Rules.hs
@@ -514,7 +514,7 @@ lookupRule :: RuleOpts -> InScopeEnv
 
 -- See Note [Extra args in the target]
 -- See comments on matchRule
-lookupRule opts rule_env@(in_scope,_) is_active fn args rules
+lookupRule opts rule_env@(ISE in_scope _) is_active fn args rules
   = -- pprTrace "lookupRule" (ppr fn <+> ppr args $$ ppr rules $$ ppr in_scope) $
     case go [] rules of
         []     -> Nothing
@@ -574,11 +574,12 @@ isMoreSpecific _        (Rule {})        (BuiltinRule {}) = True
 isMoreSpecific in_scope (Rule { ru_bndrs = bndrs1, ru_args = args1 })
                         (Rule { ru_bndrs = bndrs2, ru_args = args2
                               , ru_name = rule_name2, ru_rhs = rhs2 })
-  = isJust (matchN (full_in_scope, id_unfolding_fun)
+  = isJust (matchN in_scope_env
                    rule_name2 bndrs2 args2 args1 rhs2)
   where
-   id_unfolding_fun _ = NoUnfolding     -- Don't expand in templates
    full_in_scope = in_scope `extendInScopeSetList` bndrs1
+   in_scope_env  = ISE full_in_scope noUnfoldingFun
+                   -- noUnfoldingFun: don't expand in templates
 
 noBlackList :: Activation -> Bool
 noBlackList _ = False           -- Nothing is black listed
@@ -687,7 +688,7 @@ matchN  :: InScopeEnv
 -- trailing ones, returning the result of applying the rule to a prefix
 -- of the actual arguments.
 
-matchN (in_scope, id_unf) rule_name tmpl_vars tmpl_es target_es rhs
+matchN (ISE in_scope id_unf) rule_name tmpl_vars tmpl_es target_es rhs
   = do  { rule_subst <- match_exprs init_menv emptyRuleSubst tmpl_es target_es
         ; let (_, matched_es) = mapAccumL (lookup_tmpl rule_subst)
                                           (mkEmptySubst in_scope) $
@@ -872,7 +873,7 @@ see `init_menv` in `matchN`.
 -}
 
 rvInScopeEnv :: RuleMatchEnv -> InScopeEnv
-rvInScopeEnv renv = (rnInScopeSet (rv_lcl renv), rv_unf renv)
+rvInScopeEnv renv = ISE (rnInScopeSet (rv_lcl renv)) (rv_unf renv)
 
 -- * The domain of the TvSubstEnv and IdSubstEnv are the template
 --   variables passed into the match.
@@ -1686,7 +1687,7 @@ ruleAppCheck_help env fn args rules
         = text "Rule" <+> doubleQuotes (ftext name)
 
     rule_info opts rule
-        | Just _ <- matchRule opts (emptyInScopeSet, rc_id_unf env)
+        | Just _ <- matchRule opts (ISE emptyInScopeSet (rc_id_unf env))
                               noBlackList fn args rough_args rule
         = text "matches (which is very peculiar!)"
 
diff --git a/compiler/GHC/Core/SimpleOpt.hs b/compiler/GHC/Core/SimpleOpt.hs
index ba95baec64..609d007a5a 100644
--- a/compiler/GHC/Core/SimpleOpt.hs
+++ b/compiler/GHC/Core/SimpleOpt.hs
@@ -242,7 +242,7 @@ simple_opt_expr env expr
     rec_ids      = soe_rec_ids env
     subst        = soe_subst env
     in_scope     = getSubstInScope subst
-    in_scope_env = (in_scope, simpleUnfoldingFun)
+    in_scope_env = ISE in_scope alwaysActiveUnfoldingFun
 
     ---------------
     go (Var v)
@@ -761,11 +761,6 @@ add_info env old_bndr top_level new_rhs new_bndr
                                     False -- may be bottom or not
                                     new_rhs Nothing
 
-simpleUnfoldingFun :: IdUnfoldingFun
-simpleUnfoldingFun id
-  | isAlwaysActive (idInlineActivation id) = idUnfolding id
-  | otherwise                              = noUnfolding
-
 wrapLet :: Maybe (Id,CoreExpr) -> CoreExpr -> CoreExpr
 wrapLet Nothing      body = body
 wrapLet (Just (b,r)) body = Let (NonRec b r) body
@@ -1184,7 +1179,7 @@ data ConCont = CC [CoreExpr] Coercion
 exprIsConApp_maybe :: HasDebugCallStack
                    => InScopeEnv -> CoreExpr
                    -> Maybe (InScopeSet, [FloatBind], DataCon, [Type], [CoreExpr])
-exprIsConApp_maybe (in_scope, id_unf) expr
+exprIsConApp_maybe ise@(ISE in_scope id_unf) expr
   = go (Left in_scope) [] expr (CC [] (mkRepReflCo (exprType expr)))
   where
     go :: Either InScopeSet Subst
@@ -1304,7 +1299,7 @@ exprIsConApp_maybe (in_scope, id_unf) expr
         | (fun `hasKey` unpackCStringIdKey) ||
           (fun `hasKey` unpackCStringUtf8IdKey)
         , [arg]              <- args
-        , Just (LitString str) <- exprIsLiteral_maybe (in_scope, id_unf) arg
+        , Just (LitString str) <- exprIsLiteral_maybe ise arg
         = succeedWith in_scope floats $
           dealWithStringLiteral fun str co
         where
@@ -1400,7 +1395,7 @@ exprIsLiteral_maybe :: InScopeEnv -> CoreExpr -> Maybe Literal
 -- Nevertheless we do need to look through unfoldings for
 -- string literals, which are vigorously hoisted to top level
 -- and not subsequently inlined
-exprIsLiteral_maybe env@(_, id_unf) e
+exprIsLiteral_maybe env@(ISE _ id_unf) e
   = case e of
       Lit l     -> Just l
       Tick _ e' -> exprIsLiteral_maybe env e' -- dubious?
@@ -1430,14 +1425,14 @@ exprIsLambda_maybe _ (Lam x e)
     = Just (x, e, [])
 
 -- Still straightforward: Ticks that we can float out of the way
-exprIsLambda_maybe (in_scope_set, id_unf) (Tick t e)
+exprIsLambda_maybe ise (Tick t e)
     | tickishFloatable t
-    , Just (x, e, ts) <- exprIsLambda_maybe (in_scope_set, id_unf) e
+    , Just (x, e, ts) <- exprIsLambda_maybe ise e
     = Just (x, e, t:ts)
 
 -- Also possible: A casted lambda. Push the coercion inside
-exprIsLambda_maybe (in_scope_set, id_unf) (Cast casted_e co)
-    | Just (x, e,ts) <- exprIsLambda_maybe (in_scope_set, id_unf) casted_e
+exprIsLambda_maybe ise@(ISE in_scope_set _) (Cast casted_e co)
+    | Just (x, e,ts) <- exprIsLambda_maybe ise casted_e
     -- Only do value lambdas.
     -- this implies that x is not in scope in gamma (makes this code simpler)
     , not (isTyVar x) && not (isCoVar x)
@@ -1448,7 +1443,7 @@ exprIsLambda_maybe (in_scope_set, id_unf) (Cast casted_e co)
       res
 
 -- Another attempt: See if we find a partial unfolding
-exprIsLambda_maybe (in_scope_set, id_unf) e
+exprIsLambda_maybe ise@(ISE in_scope_set id_unf) e
     | (Var f, as, ts) <- collectArgsTicks tickishFloatable e
     , idArity f > count isValArg as
     -- Make sure there is hope to get a lambda
@@ -1456,7 +1451,7 @@ exprIsLambda_maybe (in_scope_set, id_unf) e
     -- Optimize, for beta-reduction
     , let e' = simpleOptExprWith defaultSimpleOpts (mkEmptySubst in_scope_set) (rhs `mkApps` as)
     -- Recurse, because of possible casts
-    , Just (x', e'', ts') <- exprIsLambda_maybe (in_scope_set, id_unf) e'
+    , Just (x', e'', ts') <- exprIsLambda_maybe ise e'
     , let res = Just (x', e'', ts++ts')
     = -- pprTrace "exprIsLambda_maybe:Unfold" (vcat [ppr e, ppr (x',e'')])
       res
diff --git a/compiler/GHC/HsToCore/Pmc/Solver.hs b/compiler/GHC/HsToCore/Pmc/Solver.hs
index b70c868c2f..07176b87cc 100644
--- a/compiler/GHC/HsToCore/Pmc/Solver.hs
+++ b/compiler/GHC/HsToCore/Pmc/Solver.hs
@@ -881,7 +881,7 @@ addCoreCt nabla x e = do
       where
         expr_ty       = exprType e
         expr_in_scope = mkInScopeSet (exprFreeVars e)
-        in_scope_env  = (expr_in_scope, const NoUnfolding)
+        in_scope_env  = ISE expr_in_scope noUnfoldingFun
         -- It's inconvenient to get hold of a global in-scope set
         -- here, but it'll only be needed if exprIsConApp_maybe ends
         -- up substituting inside a forall or lambda (i.e. seldom)
diff --git a/compiler/GHC/Types/Id.hs b/compiler/GHC/Types/Id.hs
index 4744147dcf..5131307a00 100644
--- a/compiler/GHC/Types/Id.hs
+++ b/compiler/GHC/Types/Id.hs
@@ -92,12 +92,14 @@ module GHC.Types.Id (
         -- ** Reading 'IdInfo' fields
         idArity,
         idCallArity, idFunRepArity,
-        idUnfolding, realIdUnfolding,
         idSpecialisation, idCoreRules, idHasRules,
         idCafInfo, idLFInfo_maybe,
         idOneShotInfo,
         idOccInfo,
 
+        IdUnfoldingFun, idUnfolding, realIdUnfolding,
+        alwaysActiveUnfoldingFun, whenActiveUnfoldingFun, noUnfoldingFun,
+
         -- ** Writing 'IdInfo' fields
         setIdUnfolding, zapIdUnfolding, setCaseBndrEvald,
         setIdArity,
@@ -126,8 +128,9 @@ module GHC.Types.Id (
 
 import GHC.Prelude
 
-import GHC.Core ( CoreRule, isStableUnfolding, evaldUnfolding,
-                 isCompulsoryUnfolding, Unfolding( NoUnfolding ), isEvaldUnfolding, hasSomeUnfolding, noUnfolding )
+import GHC.Core ( CoreRule, isStableUnfolding, evaldUnfolding
+                , isCompulsoryUnfolding, Unfolding( NoUnfolding )
+                , IdUnfoldingFun, isEvaldUnfolding, hasSomeUnfolding, noUnfolding )
 
 import GHC.Types.Id.Info
 import GHC.Types.Basic
@@ -744,9 +747,28 @@ idTagSig_maybe = tagSig . idInfo
 -- loop breaker. See 'unfoldingInfo'.
 --
 -- If you really want the unfolding of a strong loopbreaker, call 'realIdUnfolding'.
-idUnfolding :: Id -> Unfolding
+idUnfolding :: IdUnfoldingFun
 idUnfolding id = unfoldingInfo (idInfo id)
 
+noUnfoldingFun :: IdUnfoldingFun
+noUnfoldingFun _id = noUnfolding
+
+-- | Returns an unfolding only if
+--   (a) not a strong loop breaker and
+--   (b) always active
+alwaysActiveUnfoldingFun :: IdUnfoldingFun
+alwaysActiveUnfoldingFun id
+  | isAlwaysActive (idInlineActivation id) = idUnfolding id
+  | otherwise                              = noUnfolding
+
+-- | Returns an unfolding only if
+--   (a) not a strong loop breaker and
+--   (b) active in according to is_active
+whenActiveUnfoldingFun :: (Activation -> Bool) -> IdUnfoldingFun
+whenActiveUnfoldingFun is_active id
+  | is_active (idInlineActivation id) = idUnfolding id
+  | otherwise                         = NoUnfolding
+
 realIdUnfolding :: Id -> Unfolding
 -- ^ Expose the unfolding if there is one, including for loop breakers
 realIdUnfolding id = realUnfoldingInfo (idInfo id)
diff --git a/compiler/GHC/Types/Id/Info.hs b/compiler/GHC/Types/Id/Info.hs
index edd1ba0da0..ea7636bea4 100644
--- a/compiler/GHC/Types/Id/Info.hs
+++ b/compiler/GHC/Types/Id/Info.hs
@@ -469,7 +469,7 @@ setOccInfo        info oc = oc `seq` info { occInfo = oc }
 unfoldingInfo :: IdInfo -> Unfolding
 unfoldingInfo info
   | isStrongLoopBreaker (occInfo info) = trimUnfolding $ realUnfoldingInfo info
-  | otherwise                          =                realUnfoldingInfo info
+  | otherwise                          =                 realUnfoldingInfo info
 
 setUnfoldingInfo :: IdInfo -> Unfolding -> IdInfo
 setUnfoldingInfo info uf
diff --git a/testsuite/tests/simplCore/should_compile/T22802.hs b/testsuite/tests/simplCore/should_compile/T22802.hs
new file mode 100644
index 0000000000..2df903598d
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T22802.hs
@@ -0,0 +1,20 @@
+{-# OPTIONS_GHC -O1 #-}
+module T22802 where
+
+class C a where
+  f :: a -> a -> a
+  g :: a -> a -> a
+instance C () where
+  f = g
+  g = f
+
+h :: a -> () -> ()
+h = mapFB f (const ())
+
+mapFB :: (elt -> lst -> lst) -> (a -> elt) -> a -> lst -> lst
+{-# INLINE [0] mapFB #-}
+mapFB c f = \x ys -> c (f x) ys
+
+{-# RULES
+"my-mapFB" forall c a b. mapFB (mapFB c a) b = mapFB c (a.b)
+  #-}
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index a07aba3940..745bb22cd9 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -470,3 +470,4 @@ test('T22725', normal, compile, ['-O'])
 test('T22502', normal, compile, ['-O'])
 test('T22611', [when(wordsize(32), skip), grep_errmsg(r'\$salterF') ], compile, ['-O -ddump-simpl -dsuppress-uniques -dsuppress-all'])
 test('T22715_2', normal, multimod_compile, ['T22715_2', '-v0 -O -fspecialise-aggressively'])
+test('T22802', normal, compile, ['-O'])