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
18 import qualified UniqSupply
19 import qualified Unique
20 import qualified OccName
23 import qualified SrcLoc
25 import qualified IdInfo
26 import qualified CoreUtils
27 import qualified CoreSubst
28 import Outputable ( showSDoc, ppr, nest )
33 -- Create a new internal var with the given name and type. A Unique is
34 -- appended to the given name, to ensure uniqueness (not strictly neccesary,
35 -- since the Unique is also stored in the name, but this ensures variable
36 -- names are unique in the output).
37 mkInternalVar :: String -> Type.Type -> TransformMonad Var.Var
38 mkInternalVar str ty = do
40 let occname = OccName.mkVarOcc (str ++ show uniq)
41 let name = Name.mkInternalName uniq occname SrcLoc.noSrcSpan
42 return $ Var.mkLocalIdVar name ty IdInfo.vanillaIdInfo
44 -- Apply the given transformation to all expressions in the given expression,
45 -- including the expression itself.
46 everywhere :: (String, Transform) -> Transform
47 everywhere trans = applyboth (subeverywhere (everywhere trans)) trans
49 -- Apply the first transformation, followed by the second transformation, and
50 -- keep applying both for as long as expression still changes.
51 applyboth :: Transform -> (String, Transform) -> Transform
52 applyboth first (name, second) expr = do
56 (expr'', changed) <- Writer.listen $ second expr'
57 if Monoid.getAny changed
59 trace ("Transform " ++ name ++ " changed from:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n" ++ "\nTo:\n" ++ showSDoc (nest 4 $ ppr expr'') ++ "\n" ++ "Type: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr'') ++ "\n" ) $
60 applyboth first (name, second) expr''
64 -- Apply the given transformation to all direct subexpressions (only), not the
66 subeverywhere :: Transform -> Transform
67 subeverywhere trans (App a b) = do
72 subeverywhere trans (Let (Rec binds) expr) = do
74 binds' <- mapM transbind binds
75 return $ Let (Rec binds') expr'
77 transbind :: (CoreBndr, CoreExpr) -> TransformMonad (CoreBndr, CoreExpr)
82 subeverywhere trans (Lam x expr) = do
86 subeverywhere trans (Case scrut b t alts) = do
88 alts' <- mapM transalt alts
89 return $ Case scrut' b t alts'
91 transalt :: CoreAlt -> TransformMonad CoreAlt
92 transalt (con, binders, expr) = do
94 return (con, binders, expr')
97 subeverywhere trans expr = return expr
99 -- Apply the given transformation to all expressions, except for every first
100 -- argument of an application.
101 notapplied :: (String, Transform) -> Transform
102 notapplied trans = applyboth (subnotapplied trans) trans
104 -- Apply the given transformation to all (direct and indirect) subexpressions
105 -- (but not the expression itself), except for the first argument of an
106 -- applicfirst argument of an application
107 subnotapplied :: (String, Transform) -> Transform
108 subnotapplied trans (App a b) = do
109 a' <- subnotapplied trans a
110 b' <- notapplied trans b
113 -- Let subeverywhere handle all other expressions
114 subnotapplied trans expr = subeverywhere (notapplied trans) expr
116 -- Run the given transforms over the given expression
117 dotransforms :: [Transform] -> UniqSupply.UniqSupply -> CoreExpr -> CoreExpr
118 dotransforms transs uniqSupply = (flip State.evalState initState) . (dotransforms' transs)
119 where initState = TransformState uniqSupply
121 -- Runs each of the transforms repeatedly inside the State monad.
122 dotransforms' :: [Transform] -> CoreExpr -> State.State TransformState CoreExpr
123 dotransforms' transs expr = do
124 (expr', changed) <- Writer.runWriterT $ Monad.foldM (flip ($)) expr transs
125 if Monoid.getAny changed then dotransforms' transs expr' else return expr'
127 -- Inline all let bindings that satisfy the given condition
128 inlinebind :: ((CoreBndr, CoreExpr) -> Bool) -> Transform
129 inlinebind condition (Let (Rec binds) expr) | not $ null replace =
132 -- Find all simple bindings
133 (replace, others) = List.partition condition binds
134 -- Substitute the to be replaced binders with their expression
135 newexpr = substitute replace (Let (Rec others) expr)
136 -- Leave all other expressions unchanged
137 inlinebind _ expr = return expr
139 -- Sets the changed flag in the TransformMonad, to signify that some
140 -- transform has changed the result
141 setChanged :: TransformMonad ()
142 setChanged = Writer.tell (Monoid.Any True)
144 -- Sets the changed flag and returns the given value.
145 change :: a -> TransformMonad a
150 -- Create a new Unique
151 mkUnique :: TransformMonad Unique.Unique
152 mkUnique = Trans.lift $ do
153 us <- getA tsUniqSupply
154 let (us', us'') = UniqSupply.splitUniqSupply us
155 putA tsUniqSupply us'
156 return $ UniqSupply.uniqFromSupply us''
158 -- Replace each of the binders given with the coresponding expressions in the
160 substitute :: [(CoreBndr, CoreExpr)] -> CoreExpr -> CoreExpr
161 substitute replace expr = CoreSubst.substExpr subs expr
162 where subs = foldl (\s (b, e) -> CoreSubst.extendIdSubst s b e) CoreSubst.emptySubst replace