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.Arrow as Arrow
11 import qualified Control.Monad as Monad
12 import qualified Control.Monad.Trans.State as State
13 import qualified Control.Monad.Trans.Writer as Writer
14 import qualified "transformers" Control.Monad.Trans as Trans
15 import qualified Data.Map as Map
20 import qualified UniqSupply
21 import qualified Unique
22 import qualified OccName
25 import qualified SrcLoc
27 import qualified IdInfo
28 import qualified CoreUtils
29 import qualified CoreSubst
30 import qualified VarSet
31 import Outputable ( showSDoc, ppr, nest )
36 -- Create a new internal var with the given name and type. A Unique is
37 -- appended to the given name, to ensure uniqueness (not strictly neccesary,
38 -- since the Unique is also stored in the name, but this ensures variable
39 -- names are unique in the output).
40 mkInternalVar :: String -> Type.Type -> TransformMonad Var.Var
41 mkInternalVar str ty = do
43 let occname = OccName.mkVarOcc (str ++ show uniq)
44 let name = Name.mkInternalName uniq occname SrcLoc.noSrcSpan
45 return $ Var.mkLocalIdVar name ty IdInfo.vanillaIdInfo
47 -- Create a new type variable with the given name and kind. A Unique is
48 -- appended to the given name, to ensure uniqueness (not strictly neccesary,
49 -- since the Unique is also stored in the name, but this ensures variable
50 -- names are unique in the output).
51 mkTypeVar :: String -> Type.Kind -> TransformMonad Var.Var
52 mkTypeVar str kind = do
54 let occname = OccName.mkVarOcc (str ++ show uniq)
55 let name = Name.mkInternalName uniq occname SrcLoc.noSrcSpan
56 return $ Var.mkTyVar name kind
58 -- Creates a binder for the given expression with the given name. This
59 -- works for both value and type level expressions, so it can return a Var or
60 -- TyVar (which is just an alias for Var).
61 mkBinderFor :: CoreExpr -> String -> TransformMonad Var.Var
62 mkBinderFor (Type ty) string = mkTypeVar string (Type.typeKind ty)
63 mkBinderFor expr string = mkInternalVar string (CoreUtils.exprType expr)
65 -- Creates a reference to the given variable. This works for both a normal
66 -- variable as well as a type variable
67 mkReferenceTo :: Var.Var -> CoreExpr
68 mkReferenceTo var | Var.isTyVar var = (Type $ Type.mkTyVarTy var)
69 | otherwise = (Var var)
71 cloneVar :: Var.Var -> TransformMonad Var.Var
74 -- Swap out the unique, and reset the IdInfo (I'm not 100% sure what it
75 -- contains, but vannillaIdInfo is always correct, since it means "no info").
76 return $ Var.lazySetVarIdInfo (Var.setVarUnique v uniq) IdInfo.vanillaIdInfo
78 -- Creates a new function with the same name as the given binder (but with a
79 -- new unique) and with the given function body. Returns the new binder for
81 mkFunction :: CoreBndr -> CoreExpr -> TransformMonad CoreBndr
82 mkFunction bndr body = do
83 let ty = CoreUtils.exprType body
85 let newid = Var.setVarType id ty
86 Trans.lift $ addGlobalBind newid body
89 -- Apply the given transformation to all expressions in the given expression,
90 -- including the expression itself.
91 everywhere :: (String, Transform) -> Transform
92 everywhere trans = applyboth (subeverywhere (everywhere trans)) trans
94 -- Apply the first transformation, followed by the second transformation, and
95 -- keep applying both for as long as expression still changes.
96 applyboth :: Transform -> (String, Transform) -> Transform
97 applyboth first (name, second) expr = do
101 (expr'', changed) <- Writer.listen $ second expr'
103 -- trace ("Trying to apply transform " ++ name ++ " to:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n") $
106 -- trace ("Applying transform " ++ name ++ " to:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n") $
107 -- trace ("Result of applying " ++ name ++ ":\n" ++ showSDoc (nest 4 $ ppr expr'') ++ "\n" ++ "Type: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr'') ++ "\n" ) $
108 applyboth first (name, second) $
111 -- trace ("No changes") $
114 -- Apply the given transformation to all direct subexpressions (only), not the
115 -- expression itself.
116 subeverywhere :: Transform -> Transform
117 subeverywhere trans (App a b) = do
122 subeverywhere trans (Let (NonRec b bexpr) expr) = do
123 bexpr' <- trans bexpr
125 return $ Let (NonRec b bexpr') expr'
127 subeverywhere trans (Let (Rec binds) expr) = do
129 binds' <- mapM transbind binds
130 return $ Let (Rec binds') expr'
132 transbind :: (CoreBndr, CoreExpr) -> TransformMonad (CoreBndr, CoreExpr)
133 transbind (b, e) = do
137 subeverywhere trans (Lam x expr) = do
141 subeverywhere trans (Case scrut b t alts) = do
142 scrut' <- trans scrut
143 alts' <- mapM transalt alts
144 return $ Case scrut' b t alts'
146 transalt :: CoreAlt -> TransformMonad CoreAlt
147 transalt (con, binders, expr) = do
149 return (con, binders, expr')
151 subeverywhere trans (Var x) = return $ Var x
152 subeverywhere trans (Lit x) = return $ Lit x
153 subeverywhere trans (Type x) = return $ Type x
155 subeverywhere trans (Cast expr ty) = do
157 return $ Cast expr' ty
159 subeverywhere trans expr = error $ "\nNormalizeTools.subeverywhere: Unsupported expression: " ++ show expr
161 -- Apply the given transformation to all expressions, except for direct
162 -- arguments of an application
163 notappargs :: (String, Transform) -> Transform
164 notappargs trans = applyboth (subnotappargs trans) trans
166 -- Apply the given transformation to all (direct and indirect) subexpressions
167 -- (but not the expression itself), except for direct arguments of an
169 subnotappargs :: (String, Transform) -> Transform
170 subnotappargs trans (App a b) = do
171 a' <- subnotappargs trans a
172 b' <- subnotappargs trans b
175 -- Let subeverywhere handle all other expressions
176 subnotappargs trans expr = subeverywhere (notappargs trans) expr
178 -- Runs each of the transforms repeatedly inside the State monad.
179 dotransforms :: [Transform] -> CoreExpr -> TransformSession CoreExpr
180 dotransforms transs expr = do
181 (expr', changed) <- Writer.runWriterT $ Monad.foldM (flip ($)) expr transs
182 if Monoid.getAny changed then dotransforms transs expr' else return expr'
184 -- Inline all let bindings that satisfy the given condition
185 inlinebind :: ((CoreBndr, CoreExpr) -> Bool) -> Transform
186 inlinebind condition (Let (Rec binds) expr) | not $ null replace =
189 -- Find all simple bindings
190 (replace, others) = List.partition condition binds
191 -- Substitute the to be replaced binders with their expression
192 newexpr = substitute replace (Let (Rec others) expr)
193 -- Leave all other expressions unchanged
194 inlinebind _ expr = return expr
196 -- Sets the changed flag in the TransformMonad, to signify that some
197 -- transform has changed the result
198 setChanged :: TransformMonad ()
199 setChanged = Writer.tell (Monoid.Any True)
201 -- Sets the changed flag and returns the given value.
202 change :: a -> TransformMonad a
207 -- Create a new Unique
208 mkUnique :: TransformMonad Unique.Unique
209 mkUnique = Trans.lift $ do
210 us <- getA tsUniqSupply
211 let (us', us'') = UniqSupply.splitUniqSupply us
212 putA tsUniqSupply us'
213 return $ UniqSupply.uniqFromSupply us''
215 -- Replace each of the binders given with the coresponding expressions in the
217 substitute :: [(CoreBndr, CoreExpr)] -> CoreExpr -> CoreExpr
218 substitute [] expr = expr
219 -- Apply one substitution on the expression, but also on any remaining
220 -- substitutions. This seems to be the only way to handle substitutions like
221 -- [(b, c), (a, b)]. This means we reuse a substitution, which is not allowed
222 -- according to CoreSubst documentation (but it doesn't seem to be a problem).
223 -- TODO: Find out how this works, exactly.
224 substitute ((b, e):subss) expr = substitute subss' expr'
227 subs = (CoreSubst.extendSubst CoreSubst.emptySubst b e)
228 -- Apply this substitution to the main expression
229 expr' = CoreSubst.substExpr subs expr
230 -- Apply this substitution on all the expressions in the remaining
232 subss' = map (Arrow.second (CoreSubst.substExpr subs)) subss
234 -- Run a given TransformSession. Used mostly to setup the right calls and
236 runTransformSession :: UniqSupply.UniqSupply -> TransformSession a -> a
237 runTransformSession uniqSupply session = State.evalState session initState
238 where initState = TransformState uniqSupply Map.empty VarSet.emptyVarSet