packages feed

strict-impl-params 1.1.2 → 1.2.0

raw patch · 4 files changed

+282/−72 lines, 4 filesPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

Files

StrictImplParams.hs view
@@ -1,36 +1,36 @@ -{-# language LambdaCase, Strict, TupleSections, CPP #-}+{-# language LambdaCase, Strict, TupleSections, CPP, MagicHash #-} {-# options_ghc -Wincomplete-patterns -Wunused-imports #-}  module StrictImplParams (plugin) where -import Data.Foldable import Data.Maybe import GHC.Core.Class (classMethods) import GHC.Core.Predicate import GHC.Plugins  import qualified GHC.Core.TyCo.Rep as GHC+import GHC.Exts (isTrue#, reallyUnsafePtrEquality#) +#if __GLASGOW_HASKELL__ >= 914+import GHC.Types.Demand+#endif+ plugin :: Plugin plugin = defaultPlugin {-  installCoreToDos = \_ todo -> pure (CoreDoPluginPass "Strict Implicit Params" pass : todo),+  installCoreToDos = \_ todo -> pure (CoreDoPluginPass "Strict Implicit Params" pass : todo+#if __GLASGOW_HASKELL__ >= 914+                                       ++ [CoreDoPluginPass "Strict Implicit Params (call sites)" callPass]+#endif+                                      ),   pluginRecompile  = purePlugin   } -{-# inline ($$!) #-}-($$!) :: (a -> b) -> a -> b-f $$! x = f x-infixl 8 $$!--{-# inline ($$~) #-}-($$~) :: (a -> b) -> a -> b-f $$~ ~x = f x-infixl 8 $$~--map' :: (a -> b) -> [a] -> [b]-map' f = foldr' (\a bs -> (:) $$! f a $$! bs) []-{-# inline map' #-}+-- | Pointer equality. We use it to return the original node when none of its children changed,+--   so that unchanged Core is not reallocated. False negatives only cause reallocation.+same :: a -> a -> Bool+same x y = isTrue# (reallyUnsafePtrEquality# x y)+{-# inline same #-}  #if __GLASGOW_HASKELL__ <= 904 manyType :: Mult@@ -59,11 +59,12 @@ --   i.e. @(ip \@x \@a d, a)@ for @d :: IP x a@. ipValue :: Var -> Maybe (CoreExpr, Type) ipValue x = do-  (tc, args) <- splitTyConApp_maybe (varType x)-  cls        <- tyConClass_maybe tc-  valTy      <- case args of [_sym, valTy] -> Just valTy; _ -> Nothing-  sel        <- case classMethods cls of [sel] -> Just sel; _ -> Nothing-  pure (mkCoreApps (Var sel) (foldr' (\a as -> (Type a:) $! as) [Var x] args), valTy)+  (tc, args)   <- splitTyConApp_maybe (varType x)+  cls          <- tyConClass_maybe tc+  (sym, valTy) <- case args of [sym, valTy] -> Just (sym, valTy); _ -> Nothing+  sel          <- case classMethods cls of [sel] -> Just sel; _ -> Nothing+  let val = mkCoreApps (Var sel) [Type sym, Type valTy, Var x]+  pure (val, valTy)  -- | Force the value of an implicit param, continue with CoreExpr body that has Type --   type.@@ -72,66 +73,260 @@ --   to it. forceVar :: Var -> CoreExpr -> Type -> CoreExpr forceVar x body bodyTy = case ipValue x of-  Just (val, valTy) -> mkWildCase val (GHC.Scaled manyType valTy) bodyTy [Alt DEFAULT [] body]+  Just (val, valTy) -> let wild = mkWildValBinder manyType valTy in+                       Case val wild bodyTy [Alt DEFAULT [] body]   Nothing           -> error "forceVar: not an implicit parameter dictionary" +-- | Force the implicit params in the list, continue with CoreExpr body that has Type type.+forceVars :: [Var] -> CoreExpr -> Type -> CoreExpr+forceVars vars body bodyTy = case vars of+  []       -> body+  x : vars -> forceVars vars (forceVar x body bodyTy) bodyTy+ pass :: ModGuts -> CoreM ModGuts pass guts = do-  dflags <- getDynFlags+  let binds' = goBinds (mg_binds guts)+  pure $! guts {mg_binds = binds'} -  let dbg :: Outputable a => a -> String-      dbg x = showSDoc dflags (ppr x)+goBinds :: [CoreBind] -> [CoreBind]+goBinds bbs = case bbs of+  []   -> bbs+  b:bs -> let b' = goBind b; bs' = goBinds bs in+          if same b b' && same bs bs' then bbs else b':bs' -  let goBind :: CoreBind -> CoreBind-      goBind b = let+goBind :: CoreBind -> CoreBind+goBind b = case b of+  NonRec x t -> let t' = goDef [] t (varType x) in+                if same t t' then b else NonRec x t'+  Rec defs   -> let defs' = goDefs defs in+                if same defs defs' then b else Rec defs' -        go :: [Var] -> CoreExpr -> Type -> CoreExpr-        go vars t a = case t of-          Lam x t -> case coreFullView a of-            -- Note: a Core type lambda binds its own type variable, which is not necessarily the-            -- one bound by the ForAllTy in the definition's type. Since we use the def type to-            -- generate the Core of the forcing, we need to rename the type binder to match the-            -- lambda binder. We sincerely hope that Core names are unique and this renaming doesn't-            -- introduce shadowing.-            GHC.ForAllTy bndr a -> let tv = binderVar bndr-                                       a' | isTyVar x, tv /= x = substTyWith [tv] [mkTyVarTy x] a-                                          | otherwise          = a-                                   in Lam x $! go vars t a'-            -- Precaution: we add an occurrence of the binder, so whatever occurrence info it-            -- carries from the desugarer should be outdated.-            GHC.FunTy _ _ a b | isImplicitParamTy a -> let x' = zapIdOccInfo x in-                                                       Lam x' $! go (x':vars) t b-                              | otherwise           -> Lam x $! go vars t b-            _ -> error $ "unexpected type for lambda expression: " ++ dbg a-          t ->-            foldl' (\acc x -> forceVar x acc a) (goExpr t) vars+goDefs :: [(Id, CoreExpr)] -> [(Id, CoreExpr)]+goDefs ddefs = case ddefs of+  []              -> ddefs+  def@(x, t):defs -> let t' = goDef [] t (varType x); defs' = goDefs defs in+                     if same t t' then+                       if same defs defs' then ddefs else def:defs'+                     else (x, t'):defs' -        in case b of-          NonRec b t -> NonRec b $! go [] t (varType b)-          Rec defs   -> Rec $! map' (\(b, t) -> (b,) $! go [] t (varType b)) defs+-- | Traverse a definition with the given type, forcing the implicit params bound by the+--   leading lambdas. The list contains the implicit param binders bound so far.+goDef :: [Var] -> CoreExpr -> Type -> CoreExpr+goDef vars e a = case e of+  Lam x t -> case coreFullView a of+    -- Note: a Core type lambda binds its own type variable, which is not necessarily the+    -- one bound by the ForAllTy in the definition's type. Since we use the def type to+    -- generate the Core of the forcing, we need to rename the type binder to match the+    -- lambda binder. We sincerely hope that Core names are unique and this renaming doesn't+    -- introduce shadowing.+    GHC.ForAllTy bndr a -> let tv = binderVar bndr+                               a' | isTyVar x, tv /= x = substTyWith [tv] [mkTyVarTy x] a+                                  | otherwise          = a+                               t' = goDef vars t a'+                           in if same t t' then e else Lam x t'+    -- Precaution: we add an occurrence of the binder, so whatever occurrence info it+    -- carries from the desugarer should be outdated.+    GHC.FunTy _ _ a b | isImplicitParamTy a -> let x' = zapIdOccInfo x; t' = goDef (x':vars) t b in+                                               Lam x' t'+                      | otherwise           -> let t' = goDef vars t b in+                                               if same t t' then e else Lam x t'+    _ -> pprPanic "unexpected type for lambda expression" (ppr a)+  _ ->+    forceVars vars (goExpr e) a -      goExpr :: CoreExpr -> CoreExpr-      goExpr t = case t of-        Var{}                  -> t-        Lit{}                  -> t-        App t u                -> App $$! goExpr t $$! goExpr u-        Case t scr bodyty alts -> Case $$! goExpr t $$! scr $$! bodyty $$! map' goAlt alts-        Cast t coe             -> Cast $$! goExpr t $$~ coe-        Tick tck t             -> Tick tck $$! goExpr t-        Type{}                 -> t-        Coercion{}             -> t+goExpr :: CoreExpr -> CoreExpr+goExpr e = case e of+  Var{}                  -> e+  Lit{}                  -> e+  Type{}                 -> e+  Coercion{}             -> e+  App t u                -> let t' = goExpr t; u' = goExpr u in+                            if same t t' && same u u' then e else App t' u'+  Case t scr bodyty alts -> let t' = goExpr t; alts' = goAlts alts in+                            if same t t' && same alts alts' then e else Case t' scr bodyty alts'+  Cast t coe             -> let t' = goExpr t in+                            if same t t' then e else Cast t' coe+  Tick tck t             -> let t' = goExpr t in+                            if same t t' then e else Tick tck t' -        -- Note: the bound var of a free-standing lambda-        -- does not get forced!-        Lam x t -> Lam x $! goExpr t+  -- Note: the bound var of a free-standing lambda+  -- does not get forced!+  Lam x t                -> let t' = goExpr t in+                            if same t t' then e else Lam x t' -        -- Neither does an implicit let binder!-        -- In both cases the issue is that we don't know the type-        -- of the expr body and I don't like the idea of recomputing it!-        Let b t -> Let $$! goBind b $$! goExpr t+  -- Neither does an implicit let binder!+  -- In both cases the issue is that we don't know the type+  -- of the expr body and I don't like the idea of recomputing it!+  Let b t                -> let b' = goBind b; t' = goExpr t in+                            if same b b' && same t t' then e else Let b' t' -      goAlt :: Alt CoreBndr -> Alt CoreBndr-      goAlt (Alt con bs body) = Alt con bs $! goExpr body+goAlts :: [CoreAlt] -> [CoreAlt]+goAlts aalts = case aalts of+  []                        -> aalts+  alt@(Alt con bs rhs):alts -> let rhs' = goExpr rhs; alts' = goAlts alts in+                               if same rhs rhs' then+                                 if same alts alts' then aalts else alt:alts'+                               else Alt con bs rhs':alts' -  let mg_binds' = map' goBind (mg_binds guts)-  pure $! guts {mg_binds = mg_binds'}+#if __GLASGOW_HASKELL__ >= 914+-- | Starting from GHC 9.14, implicit parameter dictionaries are represented in Core by a+--   constructor application @C:IP \@x \@a e@ (which is erased in STG), instead of a newtype+--   cast. Since a constructor application is a value, CorePrep does not evaluate @e@ at+--   strict call sites, but allocates a thunk for it, even when the callee's demand signature+--   says that it strictly uses the field (which is the case for all functions strictified+--   by 'pass').+--+--   This pass runs at the end of the Core pipeline, and rewrites calls+--   @f .. (C:IP \@x \@a e) ..@ to @case e of v { DEFAULT -> f .. (C:IP \@x \@a v) .. }@+--   whenever @f@'s demand signature is strict in the dictionary field.+--   We only do this for non-trivial @e@.+callPass :: ModGuts -> CoreM ModGuts+callPass guts = do+  tag <- getUniqTag+  liftIO $ do binds <- cpBinds tag (mg_binds guts)+              pure $! guts {mg_binds = binds}++cpBinds :: Char -> [CoreBind] -> IO [CoreBind]+cpBinds tag bbs = case bbs of+  []   -> pure bbs+  b:bs -> do b' <- cpBind tag b; bs' <- cpBinds tag bs+             if same b b' && same bs bs' then pure bbs else pure (b':bs')++cpBind :: Char -> CoreBind -> IO CoreBind+cpBind tag b = case b of+  NonRec x t -> do t' <- cpExpr tag t+                   if same t t' then pure b else pure (NonRec x t')+  Rec defs   -> do defs' <- cpDefs tag defs+                   if same defs defs' then pure b else pure (Rec defs')++cpDefs :: Char -> [(Id, CoreExpr)] -> IO [(Id, CoreExpr)]+cpDefs tag ddefs = case ddefs of+  []                 -> pure ddefs+  def@(x, t):defs -> do t' <- cpExpr tag t; defs' <- cpDefs tag defs+                        if same t t' then+                          if same defs defs' then pure ddefs else pure (def:defs')+                        else pure ((x, t'):defs')++cpAlts :: Char -> [CoreAlt] -> IO [CoreAlt]+cpAlts tag aalts = case aalts of+  []                      -> pure aalts+  alt@(Alt con bs rhs):alts -> do rhs' <- cpExpr tag rhs; alts' <- cpAlts tag alts+                                  if same rhs rhs' then+                                    if same alts alts' then pure aalts else pure (alt:alts')+                                  else pure (Alt con bs rhs':alts')++cpExpr :: Char -> CoreExpr -> IO CoreExpr+cpExpr tag e = case e of+  Var{}            -> pure e+  Lit{}            -> pure e+  Type{}           -> pure e+  Coercion{}       -> pure e+  App{}            -> case appHead e of+    -- Demands only apply to saturated calls. The rightmost value arg has index n - 1.+    Var f | dmds <- fst (splitDmdSig (idDmdSig f)), not (null dmds)+          , n <- spineValArgs e 0, n >= length dmds+          , needsEval dmds (n - 1) e+          -> cpCallEval tag dmds (n - 1) e+    _     -> cpSpine tag e+  Lam x t          -> do t' <- cpExpr tag t+                         if same t t' then pure e else pure (Lam x t')+  Let b t          -> do b' <- cpBind tag b; t' <- cpExpr tag t+                         if same b b' && same t t' then pure e else pure (Let b' t')+  Case t x ty alts -> do t' <- cpExpr tag t; alts' <- cpAlts tag alts+                         if same t t' && same alts alts' then pure e else pure (Case t' x ty alts')+  Cast t co        -> do t' <- cpExpr tag t+                         if same t t' then pure e else pure (Cast t' co)+  Tick tck t       -> do t' <- cpExpr tag t+                         if same t t' then pure e else pure (Tick tck t')++-- | Rebuild an application spine without rewriting it. We must not call 'cpExpr' on the+--   function part of an 'App', since that would treat a partial application as a call.+cpSpine :: Char -> CoreExpr -> IO CoreExpr+cpSpine tag e = case e of+  App f a -> do f' <- cpSpine tag f; a' <- cpExpr tag a+                if same f f' && same a a' then pure e else pure (App f' a')+  _       -> cpExpr tag e++-- | Rebuild a call spine where at least one argument must be evaluated before the call.+cpCallEval :: Char -> [Demand] -> Int -> CoreExpr -> IO CoreExpr+cpCallEval tag dmds i t = do+  (call, evals) <- cpSpineEval tag dmds i t []+  let ty = exprType call+  seqType ty `seq` (pure $! wrapEvals ty evals call)++-- | Rebuild a call spine, returning the (scrutinee, binder) pairs to evaluate before the call.+--   @i@ is the index of the rightmost value arg in the spine.+cpSpineEval :: Char -> [Demand] -> Int -> CoreExpr -> [(CoreExpr, Id)]+            -> IO (CoreExpr, [(CoreExpr, Id)])+cpSpineEval tag dmds i t evals = case t of+  App f a+    | isTypeArg a -> do+        (f', evals') <- cpSpineEval tag dmds i f evals+        pure (App f' a, evals')+    | App con e <- a, isEvalArg dmds i a -> do+        e' <- cpExpr tag e+        u  <- uniqFromTag tag+        let ty = ipDictValTy con+            v  = mkSysLocal (fsLit "ipv") u manyType ty+        (f', evals') <- cpSpineEval tag dmds (i - 1) f ((e', v):evals)+        pure (App f' (App con (Var v)), evals')+    | otherwise -> do+        a' <- cpExpr tag a+        (f', evals') <- cpSpineEval tag dmds (i - 1) f evals+        pure (App f' a', evals')+  _ -> do t' <- cpExpr tag t; pure (t', evals)++wrapEvals :: Type -> [(CoreExpr, Id)] -> CoreExpr -> CoreExpr+wrapEvals ty evals body = case evals of+  []           -> body+  (e, v):evals -> wrapEvals ty evals (Case e v ty [Alt DEFAULT [] body])++appHead :: CoreExpr -> CoreExpr+appHead = \case+  App f _ -> appHead f+  t       -> t++spineValArgs :: CoreExpr -> Int -> Int+spineValArgs t acc = case t of+  App f a | isTypeArg a -> spineValArgs f acc+          | otherwise   -> spineValArgs f (acc + 1)+  _                     -> acc++-- | Does any value arg of the spine have to be evaluated before the call? @i@ is the index of+--   the rightmost value arg.+needsEval :: [Demand] -> Int -> CoreExpr -> Bool+needsEval dmds i = \case+  App f a | isTypeArg a         -> needsEval dmds i f+          | isEvalArg dmds i a  -> True+          | otherwise           -> needsEval dmds (i - 1) f+  _                             -> False++-- | Is the arg @C:IP \@x \@a e@ with non-trivial and unevaluated @e@, which is strictly used+--   by the @i@-th demand?+isEvalArg :: [Demand] -> Int -> CoreArg -> Bool+isEvalArg dmds i = \case+  App con e | isIPDictCon con, not (exprIsTrivial e), not (exprIsHNF e) -> case drop i dmds of+    d:_ -> strictField d+    []  -> False+  _ -> False++-- | Is the (single) field of a dictionary strictly used by the demand?+--   Note: we don't use 'viewProd' because it allocates field demands for 'Poly'.+strictField :: Demand -> Bool+strictField (_ :* sd) = case sd of+  Prod _ [fd] -> isStrUsedDmd fd+  Poly _ n    -> isStrict n && not (isAbs n) -- every field has cardinality n+  _           -> False++-- | Match @C:IP \@x \@a@.+isIPDictCon :: CoreExpr -> Bool+isIPDictCon = \case+  App (App (Var dc) (Type _)) (Type _) -> isDataConWorkId dc && isIPTyCon (dataConTyCon (idDataCon dc))+  _                                    -> False++-- | The value type @a@ of @C:IP \@x \@a@.+ipDictValTy :: CoreExpr -> Type+ipDictValTy = \case+  App _ (Type a) -> a+  _              -> error "ipDictValTy: not an implicit parameter dictionary constructor"+#endif
strict-impl-params.cabal view
@@ -1,7 +1,7 @@ name: strict-impl-params cabal-version: 1.12 build-type: Simple-version: 1.1.2+version: 1.2.0 synopsis: Plugin for making implicit parameters strict homepage: https://github.com/AndrasKovacs/ghc-strict-implicit-params bug-reports: https://github.com/AndrasKovacs/ghc-strict-implicit-params/issues
test/Defs.hs view
@@ -97,3 +97,13 @@ closure :: (?x :: Int) => Int -> (Int -> Int) closure n = \k -> k + n + ?x {-# noinline closure #-}++-- Single-use dictionary with a non-trivial field: on GHC >= 9.14 the field is evaluated+-- at the call site of g.+callSite :: (Int -> Int) -> Int -> Int+callSite k n = let ?x = k n in g 20+{-# noinline callSite #-}++callSiteBranch :: (Int -> Int) -> Int -> Int+callSiteBranch k n = let ?x = k n in if n > 0 then g 20 else dead 30+{-# noinline callSiteBranch #-}
test/Main.hs view
@@ -29,6 +29,9 @@     let ?y = 100 in eq ref "h" (h 5) 131    eq ref "localIP" (localIP 1) 17+  eq ref "callSite"        (callSite (* 2) 5)       30+  eq ref "callSiteBranchT" (callSiteBranch (* 2) 5) 30+  eq ref "callSiteBranchF" (callSiteBranch (* 2) 0) 31    let ?s = "!" in do     eq ref "str"      (str 7)  "7!"@@ -55,6 +58,8 @@   strict ref "closure"  (\_ -> let ?x = boom in closure 1 2)   strict ref "cfgUse"   (\_ -> let ?cfg = error "boom" in cfgUse 1)   strict ref "fnUse"    (\_ -> let ?k = error "boom" in fnUse 1)+  strict ref "callSite" (\_ -> callSite (\_ -> boom) 1)+  strict ref "callSiteBranchF" (\_ -> callSiteBranch (\_ -> boom) 0)    failures <- readIORef ref   if failures == 0