1 {-# LANGUAGE PackageImports #-}
3 -- This module provides functions for program transformations.
5 module NormalizeTools where
9 import qualified Data.Monoid as Monoid
10 import qualified Control.Monad as Monad
11 import qualified Control.Monad.Trans.State as State
12 import qualified Control.Monad.Trans.Writer as Writer
13 import qualified "transformers" Control.Monad.Trans as Trans
14 import qualified Data.Map as Map
19 import qualified UniqSupply
20 import qualified Unique
21 import qualified OccName
24 import qualified SrcLoc
26 import qualified IdInfo
27 import qualified CoreUtils
28 import qualified CoreSubst
29 import qualified VarSet
30 import Outputable ( showSDoc, ppr, nest )
35 -- Create a new internal var with the given name and type. A Unique is
36 -- appended to the given name, to ensure uniqueness (not strictly neccesary,
37 -- since the Unique is also stored in the name, but this ensures variable
38 -- names are unique in the output).
39 mkInternalVar :: String -> Type.Type -> TransformMonad Var.Var
40 mkInternalVar str ty = do
42 let occname = OccName.mkVarOcc (str ++ show uniq)
43 let name = Name.mkInternalName uniq occname SrcLoc.noSrcSpan
44 return $ Var.mkLocalIdVar name ty IdInfo.vanillaIdInfo
46 cloneVar :: Var.Var -> TransformMonad Var.Var
49 -- Swap out the unique, and reset the IdInfo (I'm not 100% sure what it
50 -- contains, but vannillaIdInfo is always correct, since it means "no info").
51 return $ Var.lazySetVarIdInfo (Var.setVarUnique v uniq) IdInfo.vanillaIdInfo
53 -- Apply the given transformation to all expressions in the given expression,
54 -- including the expression itself.
55 everywhere :: (String, Transform) -> Transform
56 everywhere trans = applyboth (subeverywhere (everywhere trans)) trans
58 -- Apply the first transformation, followed by the second transformation, and
59 -- keep applying both for as long as expression still changes.
60 applyboth :: Transform -> (String, Transform) -> Transform
61 applyboth first (name, second) expr = do
65 (expr'', changed) <- Writer.listen $ second expr'
67 -- trace ("Trying to apply transform " ++ name ++ " to:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n") $
70 -- trace ("Applying transform " ++ name ++ " to:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n") $
71 -- trace ("Result of applying " ++ name ++ ":\n" ++ showSDoc (nest 4 $ ppr expr'') ++ "\n" ++ "Type: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr'') ++ "\n" ) $
72 applyboth first (name, second) $
75 -- trace ("No changes") $
78 -- Apply the given transformation to all direct subexpressions (only), not the
80 subeverywhere :: Transform -> Transform
81 subeverywhere trans (App a b) = do
86 subeverywhere trans (Let (Rec binds) expr) = do
88 binds' <- mapM transbind binds
89 return $ Let (Rec binds') expr'
91 transbind :: (CoreBndr, CoreExpr) -> TransformMonad (CoreBndr, CoreExpr)
96 subeverywhere trans (Lam x expr) = do
100 subeverywhere trans (Case scrut b t alts) = do
101 scrut' <- trans scrut
102 alts' <- mapM transalt alts
103 return $ Case scrut' b t alts'
105 transalt :: CoreAlt -> TransformMonad CoreAlt
106 transalt (con, binders, expr) = do
108 return (con, binders, expr')
111 subeverywhere trans expr = return expr
113 -- Apply the given transformation to all expressions, except for every first
114 -- argument of an application.
115 notapplied :: (String, Transform) -> Transform
116 notapplied trans = applyboth (subnotapplied trans) trans
118 -- Apply the given transformation to all (direct and indirect) subexpressions
119 -- (but not the expression itself), except for the first argument of an
120 -- applicfirst argument of an application
121 subnotapplied :: (String, Transform) -> Transform
122 subnotapplied trans (App a b) = do
123 a' <- subnotapplied trans a
124 b' <- notapplied trans b
127 -- Let subeverywhere handle all other expressions
128 subnotapplied trans expr = subeverywhere (notapplied trans) expr
130 -- Runs each of the transforms repeatedly inside the State monad.
131 dotransforms :: [Transform] -> CoreExpr -> TransformSession CoreExpr
132 dotransforms transs expr = do
133 (expr', changed) <- Writer.runWriterT $ Monad.foldM (flip ($)) expr transs
134 if Monoid.getAny changed then dotransforms transs expr' else return expr'
136 -- Inline all let bindings that satisfy the given condition
137 inlinebind :: ((CoreBndr, CoreExpr) -> Bool) -> Transform
138 inlinebind condition (Let (Rec binds) expr) | not $ null replace =
141 -- Find all simple bindings
142 (replace, others) = List.partition condition binds
143 -- Substitute the to be replaced binders with their expression
144 newexpr = substitute replace (Let (Rec others) expr)
145 -- Leave all other expressions unchanged
146 inlinebind _ expr = return expr
148 -- Sets the changed flag in the TransformMonad, to signify that some
149 -- transform has changed the result
150 setChanged :: TransformMonad ()
151 setChanged = Writer.tell (Monoid.Any True)
153 -- Sets the changed flag and returns the given value.
154 change :: a -> TransformMonad a
159 -- Create a new Unique
160 mkUnique :: TransformMonad Unique.Unique
161 mkUnique = Trans.lift $ do
162 us <- getA tsUniqSupply
163 let (us', us'') = UniqSupply.splitUniqSupply us
164 putA tsUniqSupply us'
165 return $ UniqSupply.uniqFromSupply us''
167 -- Replace each of the binders given with the coresponding expressions in the
169 substitute :: [(CoreBndr, CoreExpr)] -> CoreExpr -> CoreExpr
170 substitute replace expr = CoreSubst.substExpr subs expr
171 where subs = foldl (\s (b, e) -> CoreSubst.extendSubst s b e) CoreSubst.emptySubst replace
173 -- Run a given TransformSession. Used mostly to setup the right calls and
175 runTransformSession :: UniqSupply.UniqSupply -> TransformSession a -> a
176 runTransformSession uniqSupply session = State.evalState session initState
177 where initState = TransformState uniqSupply Map.empty VarSet.emptyVarSet