1{-# LANGUAGE CPP #-} 2 3-- | (Mostly) textbook instance of the lambda lifting transformation, 4-- selecting which bindings to lambda lift by consulting 'goodToLift'. 5module StgLiftLams.Transformation (stgLiftLams) where 6 7#include "HsVersions.h" 8 9import GhcPrelude 10 11import BasicTypes 12import DynFlags 13import Id 14import IdInfo 15import StgFVs ( annBindingFreeVars ) 16import StgLiftLams.Analysis 17import StgLiftLams.LiftM 18import StgSyn 19import Outputable 20import UniqSupply 21import Util 22import VarSet 23import Control.Monad ( when ) 24import Data.Maybe ( isNothing ) 25 26-- | Lambda lifts bindings to top-level deemed worth lifting (see 'goodToLift'). 27stgLiftLams :: DynFlags -> UniqSupply -> [InStgTopBinding] -> [OutStgTopBinding] 28stgLiftLams dflags us = runLiftM dflags us . foldr liftTopLvl (pure ()) 29 30liftTopLvl :: InStgTopBinding -> LiftM () -> LiftM () 31liftTopLvl (StgTopStringLit bndr lit) rest = withSubstBndr bndr $ \bndr' -> do 32 addTopStringLit bndr' lit 33 rest 34liftTopLvl (StgTopLifted bind) rest = do 35 let is_rec = isRec $ fst $ decomposeStgBinding bind 36 when is_rec startBindingGroup 37 let bind_w_fvs = annBindingFreeVars bind 38 withLiftedBind TopLevel (tagSkeletonTopBind bind_w_fvs) NilSk $ \mb_bind' -> do 39 -- We signal lifting of a binding through returning Nothing. 40 -- Should never happen for a top-level binding, though, since we are already 41 -- at top-level. 42 case mb_bind' of 43 Nothing -> pprPanic "StgLiftLams" (text "Lifted top-level binding") 44 Just bind' -> addLiftedBinding bind' 45 when is_rec endBindingGroup 46 rest 47 48withLiftedBind 49 :: TopLevelFlag 50 -> LlStgBinding 51 -> Skeleton 52 -> (Maybe OutStgBinding -> LiftM a) 53 -> LiftM a 54withLiftedBind top_lvl bind scope k 55 | isTopLevel top_lvl 56 = withCaffyness (is_caffy pairs) go 57 | otherwise 58 = go 59 where 60 (rec, pairs) = decomposeStgBinding bind 61 is_caffy = any (mayHaveCafRefs . idCafInfo . binderInfoBndr . fst) 62 go = withLiftedBindPairs top_lvl rec pairs scope (k . fmap (mkStgBinding rec)) 63 64withLiftedBindPairs 65 :: TopLevelFlag 66 -> RecFlag 67 -> [(BinderInfo, LlStgRhs)] 68 -> Skeleton 69 -> (Maybe [(Id, OutStgRhs)] -> LiftM a) 70 -> LiftM a 71withLiftedBindPairs top rec pairs scope k = do 72 let (infos, rhss) = unzip pairs 73 let bndrs = map binderInfoBndr infos 74 expander <- liftedIdsExpander 75 dflags <- getDynFlags 76 case goodToLift dflags top rec expander pairs scope of 77 -- @abs_ids@ is the set of all variables that need to become parameters. 78 Just abs_ids -> withLiftedBndrs abs_ids bndrs $ \bndrs' -> do 79 -- Within this block, all binders in @bndrs@ will be noted as lifted, so 80 -- that the return value of @liftedIdsExpander@ in this context will also 81 -- expand the bindings in @bndrs@ to their free variables. 82 -- Now we can recurse into the RHSs and see if we can lift any further 83 -- bindings. We pass the set of expanded free variables (thus OutIds) on 84 -- to @liftRhs@ so that it can add them as parameter binders. 85 when (isRec rec) startBindingGroup 86 rhss' <- traverse (liftRhs (Just abs_ids)) rhss 87 let pairs' = zip bndrs' rhss' 88 addLiftedBinding (mkStgBinding rec pairs') 89 when (isRec rec) endBindingGroup 90 k Nothing 91 Nothing -> withSubstBndrs bndrs $ \bndrs' -> do 92 -- Don't lift the current binding, but possibly some bindings in their 93 -- RHSs. 94 rhss' <- traverse (liftRhs Nothing) rhss 95 let pairs' = zip bndrs' rhss' 96 k (Just pairs') 97 98liftRhs 99 :: Maybe (DIdSet) 100 -- ^ @Just former_fvs@ <=> this RHS was lifted and we have to add @former_fvs@ 101 -- as lambda binders, discarding all free vars. 102 -> LlStgRhs 103 -> LiftM OutStgRhs 104liftRhs mb_former_fvs rhs@(StgRhsCon ccs con args) 105 = ASSERT2(isNothing mb_former_fvs, text "Should never lift a constructor" $$ ppr rhs) 106 StgRhsCon ccs con <$> traverse liftArgs args 107liftRhs Nothing (StgRhsClosure _ ccs upd infos body) = do 108 -- This RHS wasn't lifted. 109 withSubstBndrs (map binderInfoBndr infos) $ \bndrs' -> 110 StgRhsClosure noExtFieldSilent ccs upd bndrs' <$> liftExpr body 111liftRhs (Just former_fvs) (StgRhsClosure _ ccs upd infos body) = do 112 -- This RHS was lifted. Insert extra binders for @former_fvs@. 113 withSubstBndrs (map binderInfoBndr infos) $ \bndrs' -> do 114 let bndrs'' = dVarSetElems former_fvs ++ bndrs' 115 StgRhsClosure noExtFieldSilent ccs upd bndrs'' <$> liftExpr body 116 117liftArgs :: InStgArg -> LiftM OutStgArg 118liftArgs a@(StgLitArg _) = pure a 119liftArgs (StgVarArg occ) = do 120 ASSERTM2( not <$> isLifted occ, text "StgArgs should never be lifted" $$ ppr occ ) 121 StgVarArg <$> substOcc occ 122 123liftExpr :: LlStgExpr -> LiftM OutStgExpr 124liftExpr (StgLit lit) = pure (StgLit lit) 125liftExpr (StgTick t e) = StgTick t <$> liftExpr e 126liftExpr (StgApp f args) = do 127 f' <- substOcc f 128 args' <- traverse liftArgs args 129 fvs' <- formerFreeVars f 130 let top_lvl_args = map StgVarArg fvs' ++ args' 131 pure (StgApp f' top_lvl_args) 132liftExpr (StgConApp con args tys) = StgConApp con <$> traverse liftArgs args <*> pure tys 133liftExpr (StgOpApp op args ty) = StgOpApp op <$> traverse liftArgs args <*> pure ty 134liftExpr (StgLam _ _) = pprPanic "stgLiftLams" (text "StgLam") 135liftExpr (StgCase scrut info ty alts) = do 136 scrut' <- liftExpr scrut 137 withSubstBndr (binderInfoBndr info) $ \bndr' -> do 138 alts' <- traverse liftAlt alts 139 pure (StgCase scrut' bndr' ty alts') 140liftExpr (StgLet scope bind body) 141 = withLiftedBind NotTopLevel bind scope $ \mb_bind' -> do 142 body' <- liftExpr body 143 case mb_bind' of 144 Nothing -> pure body' -- withLiftedBindPairs decided to lift it and already added floats 145 Just bind' -> pure (StgLet noExtFieldSilent bind' body') 146liftExpr (StgLetNoEscape scope bind body) 147 = withLiftedBind NotTopLevel bind scope $ \mb_bind' -> do 148 body' <- liftExpr body 149 case mb_bind' of 150 Nothing -> pprPanic "stgLiftLams" (text "Should never decide to lift LNEs") 151 Just bind' -> pure (StgLetNoEscape noExtFieldSilent bind' body') 152 153liftAlt :: LlStgAlt -> LiftM OutStgAlt 154liftAlt (con, infos, rhs) = withSubstBndrs (map binderInfoBndr infos) $ \bndrs' -> 155 (,,) con bndrs' <$> liftExpr rhs 156