First PoC partially donewip/andreask/spec-transitive

6 files changed, 265 insertions, 12 deletions

diff --git a/compiler/GHC/Core/FVs.hs b/compiler/GHC/Core/FVs.hs
index 65b654356e..74bed60e38 100644
--- a/compiler/GHC/Core/FVs.hs
+++ b/compiler/GHC/Core/FVs.hs
@@ -31,10 +31,11 @@ module GHC.Core.FVs (
         bndrRuleAndUnfoldingIds,
         idFVs,
         idRuleVars, stableUnfoldingVars,
-        ruleFreeVars, rulesFreeVars,
+        ruleFreeVars, rulesFreeVars, rulesSomeFreeVars, rulesSomeFreeVarsList,
         rulesFreeVarsDSet, mkRuleInfo,
         ruleLhsFreeIds, ruleLhsFreeIdsList,
         ruleRhsFreeVars, rulesRhsFreeIds,
+        rulesRhsSomeFVs, ruleRhsSomeFreeVars,
 
         exprFVs,
 
@@ -466,34 +467,47 @@ data RuleFVsFrom
 -- | Those locally-defined variables free in the left and/or right hand sides
 -- of the rule, depending on the first argument. Returns an 'FV' computation.
 ruleFVs :: RuleFVsFrom -> CoreRule -> FV
-ruleFVs !_   (BuiltinRule {}) = emptyFV
-ruleFVs from (Rule { ru_fn = _do_not_include
+ruleFVs = ruleSomeFVs isLocalVar
+
+ruleSomeFVs :: InterestingVarFun -> RuleFVsFrom -> CoreRule -> FV
+ruleSomeFVs _fv_cand !_   (BuiltinRule {}) = emptyFV
+ruleSomeFVs fv_cand from (Rule { ru_fn = _do_not_include
                      -- See Note [Rule free var hack]
                    , ru_bndrs = bndrs
                    , ru_rhs = rhs, ru_args = args })
-  = filterFV isLocalVar $ addBndrs bndrs (exprs_fvs exprs)
+  = filterFV fv_cand $ addBndrs bndrs (exprs_fvs exprs)
   where
     exprs = case from of
       LhsOnly   -> args
       RhsOnly   -> [rhs]
       BothSides -> rhs:args
-
 -- | Those locally-defined variables free in the left and/or right hand sides
 -- from several rules, depending on the first argument.
 -- Returns an 'FV' computation.
 rulesFVs :: RuleFVsFrom -> [CoreRule] -> FV
 rulesFVs from = mapUnionFV (ruleFVs from)
 
+rulesSomeFVs :: InterestingVarFun -> RuleFVsFrom -> [CoreRule] -> FV
+rulesSomeFVs fv_cand from = mapUnionFV (ruleSomeFVs fv_cand from)
+
 -- | Those variables free in the right hand side of a rule returned as a
 -- non-deterministic set
 ruleRhsFreeVars :: CoreRule -> VarSet
 ruleRhsFreeVars = fvVarSet . ruleFVs RhsOnly
 
+-- | Those variables free in the right hand side of a rule returned as a
+-- non-deterministic set
+ruleRhsSomeFreeVars :: InterestingVarFun -> CoreRule -> VarSet
+ruleRhsSomeFreeVars fv_cand = fvVarSet . ruleSomeFVs fv_cand RhsOnly
+
 -- | Those locally-defined free 'Id's in the right hand side of several rules
 -- returned as a non-deterministic set
 rulesRhsFreeIds :: [CoreRule] -> VarSet
 rulesRhsFreeIds = fvVarSet . filterFV isLocalId . rulesFVs RhsOnly
 
+rulesRhsSomeFVs :: InterestingVarFun -> [CoreRule] -> VarSet
+rulesRhsSomeFVs fv_cand = fvVarSet . rulesSomeFVs fv_cand RhsOnly
+
 ruleLhsFreeIds :: CoreRule -> VarSet
 -- ^ This finds all locally-defined free Ids on the left hand side of a rule
 -- and returns them as a non-deterministic set
@@ -518,6 +532,12 @@ rulesFreeVarsDSet rules = fvDVarSet $ rulesFVs BothSides rules
 rulesFreeVars :: [CoreRule] -> VarSet
 rulesFreeVars rules = fvVarSet $ rulesFVs BothSides rules
 
+rulesSomeFreeVars :: InterestingVarFun -> [CoreRule] -> VarSet
+rulesSomeFreeVars fv_cand rules = fvVarSet $ rulesSomeFVs fv_cand BothSides rules
+
+rulesSomeFreeVarsList :: InterestingVarFun -> [CoreRule] -> [Var]
+rulesSomeFreeVarsList fv_cand rules = fvVarList $ rulesSomeFVs fv_cand BothSides rules
+
 -- | Make a 'RuleInfo' containing a number of 'CoreRule's, suitable
 -- for putting into an 'IdInfo'
 mkRuleInfo :: [CoreRule] -> RuleInfo
diff --git a/compiler/GHC/Core/Opt/Simplify/Iteration.hs b/compiler/GHC/Core/Opt/Simplify/Iteration.hs
index 82d84d0012..8dabac1888 100644
--- a/compiler/GHC/Core/Opt/Simplify/Iteration.hs
+++ b/compiler/GHC/Core/Opt/Simplify/Iteration.hs
@@ -677,6 +677,7 @@ mkCastWrapperPragInfo prag_info
                     , inl_rule   = rule_info }  --     in GHC.Core.Opt.WorkWrap
                                     -- RuleMatchInfo is (and must be) unaffected
         (pragHasInlineable prag_info)
+        (pragSpecRec prag_info)
   where
     -- See Note [Wrapper activation] in GHC.Core.Opt.WorkWrap
     -- But simpler, because we don't need to disable during InitialPhase
diff --git a/compiler/GHC/Core/Opt/SpecRec.hs b/compiler/GHC/Core/Opt/SpecRec.hs
new file mode 100644
index 0000000000..aa40e54e9e
--- /dev/null
+++ b/compiler/GHC/Core/Opt/SpecRec.hs
@@ -0,0 +1,211 @@
+{-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE DeriveDataTypeable #-}
+{-# LANGUAGE DeriveAnyClass #-}
+{-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE TupleSections #-}
+
+-- | Transfer specrec pragmas from functions having such a pragma
+-- to functions calling such functions.
+module GHC.Core.Opt.SpecRec
+    ( transferSpecRecs
+    ) where
+
+import GHC.Prelude
+
+
+import GHC.Types.Basic
+import GHC.Driver.Session
+import GHC.Types.Name hiding (varName)
+import GHC.Types.Id
+import GHC.Unit.Module.ModGuts
+import GHC.Types.Var.Set
+import GHC.Types.Name.Env
+import GHC.Unit.Types
+import GHC.Core
+import GHC.Core.Rules
+import GHC.Core.FVs
+import GHC.Utils.Outputable
+
+import Data.Graph
+import GHC.Utils.Monad.State.Strict
+import Control.Monad
+import Data.Maybe
+{-
+-- We need to transfer the pragma in these cases:
+
+{-# SPECREC foo #-}
+foo = ...
+
+We transfer the pragma if foo is mentioned in:
+* The RHS of a function
+* The unfolding. -- TODO: Not needed after desugar?
+-- TODO: Rules
+
+-}
+
+transferSpecRecs :: ModGuts -> ModGuts
+-- transferSpecRecs _dflags guts = guts
+transferSpecRecs guts =
+  let env :: Env
+      env = Env
+        { thisModule = mg_module guts
+        , orphanRules = mkRuleBase (mg_rules guts)
+        }
+  in guts { mg_binds = doCoreProgram env (mg_binds guts)
+          }
+
+-- bind_fvs (NonRec _ rhs) =
+doCoreProgram :: Env -> CoreProgram -> CoreProgram
+doCoreProgram env binds = flip evalState (SS mempty mempty) $ do
+    -- pprTraceM "binds_in" $ vcat $ map ppr binds
+    let sorted_binds = depAnal (map getName . bindersOf) bind_deps binds
+    -- pprTraceM "binds_sorted" $ vcat $ map (ppr . bindersOf) sorted_binds
+    done_binds <- doSccs env sorted_binds
+    -- pprTraceM "binds_out" $ vcat $ map (ppr . bindersOf) done_binds
+    return done_binds
+    where
+        bind_deps bind =
+            let bs = bindersOf bind
+                rhss = rhssOfBind bind
+                rhss_fvs = exprsSomeFreeVarsList (const True) rhss
+                unfs = map realIdUnfolding bs
+                unf_rhs = catMaybes . map maybeUnfoldingTemplate $ unfs
+                unf_fvs = exprsSomeFreeVarsList (const True) unf_rhs
+                id_rules = concatMap idCoreRules bs
+                id_rules_fvs = rulesSomeFreeVarsList (const True) id_rules
+            in map getName $ unf_fvs ++ id_rules_fvs ++ unf_fvs ++ rhss_fvs
+
+mcons :: Monad m => a -> m [a] -> m [a]
+mcons x xs = liftM (x:) xs
+
+doSccs :: Env -> [SCC CoreBind] -> M [CoreBind]
+doSccs env binds = do
+    bindss <- mapM (doScc env) binds
+    pure $ concat bindss
+
+doScc :: Env -> SCC CoreBind -> M [CoreBind]
+doScc env (AcyclicSCC bind) = do
+    (b,is_spec) <- doBind env bind
+    when is_spec $ addSpecBinders $ bindersOf b
+    pure [b]
+doScc env (CyclicSCC binds) = do
+    -- A bunch of binders which might refer to each other in a cyclic fashion via
+    -- something like rules. So we must put specrec on all of them.
+    (bs,is_specs) <- unzip <$> mapM (doBind env) binds
+    when (or is_specs) $ addSpecBinders $ bindersOfBinds bs
+    pure $ map setSpec bs
+
+addSpecBinders :: [Id] -> M ()
+addSpecBinders ids = do
+    mapM_ addSpec ids
+    mapM_ addDone ids
+    return ()
+
+doBind :: Env -> CoreBind -> M (CoreBind, Bool)
+doBind env bind = do
+    to_spec <- spec_set <$> get
+    let bs = bindersOf bind
+        rhss = rhssOfBind bind
+        spec_id = any idHasSpecRec bs
+        spec_rhs = any (is_spec_expr to_spec) rhss
+        spec_unf = any (is_spec_unf to_spec . realIdUnfolding) bs
+        spec_rules = any (is_spec_rules to_spec . idCoreRules) bs
+
+    if spec_id || spec_rhs || spec_unf || spec_rules
+        -- If the rhs, unfolding or a rule rhs mentions a spec-rec function
+        -- we must make the function itself spec-rec
+        then pure (setSpec bind, True)
+        -- Otherwise another binding might still become spec-rec in the future
+        else pure (bind, False)
+
+    where
+        is_spec_rules spec_set rules =
+            let fvs = rulesRhsSomeFVs (\v -> isId v && idHasSpecRec v || elemVarSet v spec_set) rules
+            in not (isEmptyVarSet fvs)
+
+        is_spec_expr :: VarSet -> CoreExpr -> Bool
+        is_spec_expr spec_set expr =
+            let fvs = exprSomeFreeVars (\v -> isId v && idHasSpecRec v || elemVarSet v spec_set) expr
+            in not (isEmptyVarSet fvs)
+
+        is_spec_unf :: VarSet -> Unfolding -> Bool
+        is_spec_unf spec_set unf = do
+            case maybeUnfoldingSource unf of
+                -- We already look at the rhs and the unf is the same
+                Just VanillaSrc -> False
+                _ -> case maybeUnfoldingTemplate unf of
+                    Just unf_tmpl -> do
+                        is_spec_expr spec_set unf_tmpl
+                    Nothing -> False
+
+        -- TODO: Properly set activation
+
+setSpec :: Bind Id -> Bind Id
+setSpec (NonRec b rhs) = NonRec (setHasSpecRec b (Just AlwaysActive)) rhs
+setSpec (Rec pairs) = Rec $ map (\(b,rhs) -> (setHasSpecRec b (Just AlwaysActive), rhs)) pairs
+
+
+-- doBind :: Env -> CoreBind -> M CoreBind
+-- doBind env (NonRec b rhs) = NonRec b <$> doExpr (addParent b env) rhs
+-- doBind env (Rec bs) = Rec <$> mapM doPair bs
+--   where
+--     doPair (b,rhs) = (b,) <$> doExpr (addParent b env) rhs
+
+-- doExpr :: Env -> CoreExpr -> M CoreExpr
+-- doExpr env e@(Var v)
+--   | needsCallSiteCostCentre env v = do
+--     let nameDoc :: SDoc
+--         nameDoc = withUserStyle alwaysQualify DefaultDepth $
+--           hcat (punctuate dot (map ppr (parents env))) <> parens (text "calling:" <> ppr v)
+
+--         ccName :: CcName
+--         ccName = mkFastString $ renderWithContext defaultSDocContext nameDoc
+--     ccIdx <- getCCIndex' ccName
+--     let count = countEntries env
+--         span = case revParents env of
+--           top:_ -> nameSrcSpan $ varName top
+--           _     -> noSrcSpan
+--         cc = NormalCC (ExprCC ccIdx) ccName (thisModule env) span
+--         tick :: CoreTickish
+--         tick = ProfNote cc count True
+--     pure $ Tick tick e
+--   | otherwise = pure e
+-- doExpr _env e@(Lit _)       = pure e
+-- doExpr env (f `App` x)      = App <$> doExpr env f <*> doExpr env x
+-- doExpr env (Lam b x)        = Lam b <$> doExpr env x
+-- doExpr env (Let b rhs)      = Let <$> doBind env b <*> doExpr env rhs
+-- doExpr env (Case scrut b ty alts) =
+--     Case <$> doExpr env scrut <*> pure b <*> pure ty <*> mapM doAlt alts
+--   where
+--     doAlt (Alt con bs rhs)  = Alt con bs <$> doExpr env rhs
+-- doExpr env (Cast expr co)   = Cast <$> doExpr env expr <*> pure co
+-- doExpr env (Tick t e)       = Tick t <$> doExpr env e
+-- doExpr _env e@(Type _)      = pure e
+-- doExpr _env e@(Coercion _)  = pure e
+
+data SpecState = SS
+    { spec_set :: !VarSet
+    , spec_done :: !VarSet
+    }
+
+type M = State SpecState
+
+addSpec :: Var -> M ()
+addSpec v = do
+    s <- get
+    put $! s { spec_set = extendVarSet (spec_set s) v }
+
+addDone :: Var -> M ()
+addDone v = do
+    s <- get
+    put $! s { spec_done = extendVarSet (spec_done s) v }
+
+
+
+data Env = Env
+  { thisModule  :: Module
+  , orphanRules :: RuleBase
+  }
+
diff --git a/compiler/GHC/Core/Opt/Specialise.hs b/compiler/GHC/Core/Opt/Specialise.hs
index 73024ed79b..10ab7f0717 100644
--- a/compiler/GHC/Core/Opt/Specialise.hs
+++ b/compiler/GHC/Core/Opt/Specialise.hs
@@ -34,6 +34,7 @@ import GHC.Core.Utils     ( exprIsTrivial
 import GHC.Core.FVs
 import GHC.Core.TyCo.FVs ( tyCoVarsOfTypeList )
 import GHC.Core.Opt.Arity( collectBindersPushingCo )
+import GHC.Core.Opt.SpecRec (transferSpecRecs)
 
 import GHC.Builtin.Types  ( unboxedUnitTy )
 
@@ -52,6 +53,7 @@ import GHC.Types.Var      ( PiTyBinder(..), isLocalVar, isInvisibleFunArg )
 import GHC.Types.Var.Set
 import GHC.Types.Var.Env
 import GHC.Types.Id
+import GHC.Types.Id.Info
 import GHC.Types.Error
 
 import GHC.Utils.Error ( mkMCDiagnostic )
@@ -636,10 +638,13 @@ Hence, the invariant is this:
 
 -- | Specialise calls to type-class overloaded functions occurring in a program.
 specProgram :: ModGuts -> CoreM ModGuts
-specProgram guts@(ModGuts { mg_module = this_mod
-                          , mg_rules  = local_rules
-                          , mg_binds  = binds })
+specProgram guts_in
   = do { dflags   <- getDynFlags
+
+       ; let guts@(ModGuts  { mg_module = this_mod
+                            , mg_rules  = local_rules
+                            , mg_binds  = binds }) = transferSpecRecs guts_in
+
        ; rule_env <- initRuleEnv guts
                      -- See Note [Fire rules in the specialiser]
 
@@ -1613,6 +1618,7 @@ specCalls spec_imp env dict_binds existing_rules calls_for_me fn rhs
   |  notNull calls_for_me               -- And there are some calls to specialise
   ,  not (isNeverActive inl_act)
   || idHasInlineable fn -- Explicit INLINEABLE pragma
+  || idHasSpecRec fn -- SpecRec
   || gopt Opt_SpecialiseAggressively dflags -- -fspecialise-aggressively
   , not (isOpaquePragma inl_prag)
   -- Don't specialise NOINLINE things by default.
@@ -1766,6 +1772,11 @@ specCalls spec_imp env dict_binds existing_rules calls_for_me fn rhs
                   | otherwise
                   = inl_prag
 
+                spec_inlineable = idHasInlineable fn
+                spec_rec = idSpecRec fn
+
+                spec_prag_info = mkPragInfo spec_inl_prag spec_inlineable spec_rec
+
                 --------------------------------------
                 -- Adding arity information just propagates it a bit faster
                 --      See Note [Arity decrease] in GHC.Core.Opt.Simplify
@@ -1773,7 +1784,7 @@ specCalls spec_imp env dict_binds existing_rules calls_for_me fn rhs
                 -- So if f has INLINE[1] so does spec_fn
                 arity_decr     = count isValArg rule_lhs_args - count isId spec_bndrs
                 spec_f_w_arity = spec_fn `setIdArity`      max 0 (fn_arity - arity_decr)
-                                         `setInlinePragma` spec_inl_prag
+                                         `setIdPragmaInfo` spec_prag_info
                                          `setIdUnfolding`  spec_unf
                                          `asJoinId_maybe`  spec_join_arity
 
diff --git a/compiler/GHC/Core/Opt/WorkWrap.hs b/compiler/GHC/Core/Opt/WorkWrap.hs
index 69ed8331f3..874e44e338 100644
--- a/compiler/GHC/Core/Opt/WorkWrap.hs
+++ b/compiler/GHC/Core/Opt/WorkWrap.hs
@@ -830,6 +830,7 @@ mkWWBindPair ww_opts fn_id fn_info fn_args fn_body work_uniq div
       -- inl_inline: copy from fn_id; see Note [Worker/wrapper for INLINABLE functions]
       -- inl_act:    see Note [Worker activation]
       -- inl_rule:   it does not make sense for workers to be constructorlike.
+    work_prag_info = mkPragInfo work_prag fn_has_inlineable fn_spec_rec
 
     work_join_arity | isJoinId fn_id = Just join_arity
                     | otherwise      = Nothing
@@ -844,8 +845,7 @@ mkWWBindPair ww_opts fn_id fn_info fn_args fn_body work_uniq div
                         -- Doesn't matter much, since we will simplify next, but
                         -- seems right-er to do so
 
-                `setInlinePragma` work_prag
-                `setHasInlineable` fn_has_inlineable
+                `setIdPragmaInfo` work_prag_info
 
                 `setIdUnfolding` mkWorkerUnfolding simpl_opts work_fn fn_unfolding
                         -- See Note [Worker/wrapper for INLINABLE functions]
@@ -874,6 +874,7 @@ mkWWBindPair ww_opts fn_id fn_info fn_args fn_body work_uniq div
 
     wrap_rhs  = wrap_fn work_id
     wrap_prag = mkStrWrapperInlinePrag fn_inl_prag fn_rules
+    wrap_prag_info = mkPragInfo wrap_prag fn_has_inlineable fn_spec_rec
     wrap_unf  = mkWrapperUnfolding (simpleOptExpr simpl_opts wrap_rhs) arity
 
     wrap_id   = fn_id `setIdUnfolding`  wrap_unf
@@ -881,7 +882,7 @@ mkWWBindPair ww_opts fn_id fn_info fn_args fn_body work_uniq div
                       `setIdOccInfo`    noOccInfo
                       -- We must keep hasInlineable to ensure wrappers can specialise
                       -- if they are NOINLINE[final]
-                      `setHasInlineable`fn_has_inlineable
+                      `setIdPragmaInfo` wrap_prag_info
                         -- Zap any loop-breaker-ness, to avoid bleating from Lint
                         -- about a loop breaker with an INLINE rule
 
@@ -890,6 +891,7 @@ mkWWBindPair ww_opts fn_id fn_info fn_args fn_body work_uniq div
     fn_unfolding    = realUnfoldingInfo fn_info
     fn_has_inlineable = inlineableInfo fn_info
     fn_rules        = ruleInfoRules (ruleInfo fn_info)
+    fn_spec_rec     = specRecInfo fn_info
 
 mkStrWrapperInlinePrag :: InlinePragma -> [CoreRule] -> InlinePragma
 mkStrWrapperInlinePrag (InlinePragma { inl_inline = fn_inl
diff --git a/compiler/GHC/Core/Ppr.hs b/compiler/GHC/Core/Ppr.hs
index 9f7bb747b3..85ae07fc6c 100644
--- a/compiler/GHC/Core/Ppr.hs
+++ b/compiler/GHC/Core/Ppr.hs
@@ -489,6 +489,9 @@ pprIdBndrInfo info
     has_inlineable = inlineableInfo info &&
                      isNoInlinePragma prag_info -- The flag is redundant
                                                 -- unless we have NOINLINE.
+    spec_rec = specRecInfo info
+    has_spec_rec = isJust spec_rec || True
+
 
     occ_info  = occInfo info
     dmd_info  = demandInfo info
@@ -502,6 +505,7 @@ pprIdBndrInfo info
     doc = showAttributes
           [ (has_prag, text "InlPrag=" <> pprInlineDebug prag_info)
           , (has_inlineable, text "Inlineable")
+          , (has_spec_rec, text "SpecRec:" <> ppr spec_rec)
           , (has_occ,  text "Occ=" <> ppr occ_info)
           , (has_dmd,  text "Dmd=" <> ppr dmd_info)
           , (has_lbv , text "OS=" <> ppr lbv_info)
@@ -512,6 +516,7 @@ instance Outputable IdInfo where
     [ (has_prag,         text "InlPrag=" <> pprInlineDebug prag_info)
     -- Todo: This is only interesting for NoInline pragmas
     , (has_inlineable,   text "Inlineable")
+    , (has_spec_rec,     text "SpecRec:" <> ppr spec_rec)
     , (has_occ,          text "Occ=" <> ppr occ_info)
     , (has_dmd,          text "Dmd=" <> ppr dmd_info)
     , (has_lbv ,         text "OS=" <> ppr lbv_info)
@@ -525,6 +530,9 @@ instance Outputable IdInfo where
     where
       prag_info = inlinePragInfo info
       has_prag  = not (isDefaultInlinePragma prag_info)
+      spec_rec = specRecInfo info
+      has_spec_rec = isJust spec_rec  || True
+
 
       occ_info  = occInfo info
       has_occ   = not (isManyOccs occ_info)