Split off the VHDL type generating code.
[matthijs/master-project/cλash.git] / NormalizeTools.hs
1 {-# LANGUAGE PackageImports #-}
2 -- 
3 -- This module provides functions for program transformations.
4 --
5 module NormalizeTools where
6 -- Standard modules
7 import Debug.Trace
8 import qualified List
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
15 import Data.Accessor
16
17 -- GHC API
18 import CoreSyn
19 import qualified UniqSupply
20 import qualified Unique
21 import qualified OccName
22 import qualified Name
23 import qualified Var
24 import qualified SrcLoc
25 import qualified Type
26 import qualified IdInfo
27 import qualified CoreUtils
28 import qualified CoreSubst
29 import qualified VarSet
30 import Outputable ( showSDoc, ppr, nest )
31
32 -- Local imports
33 import NormalizeTypes
34
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
41   uniq <- mkUnique
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
45
46 -- Create a new type variable with the given name and kind. A Unique is
47 -- appended to the given name, to ensure uniqueness (not strictly neccesary,
48 -- since the Unique is also stored in the name, but this ensures variable
49 -- names are unique in the output).
50 mkTypeVar :: String -> Type.Kind -> TransformMonad Var.Var
51 mkTypeVar str kind = do
52   uniq <- mkUnique
53   let occname = OccName.mkVarOcc (str ++ show uniq)
54   let name = Name.mkInternalName uniq occname SrcLoc.noSrcSpan
55   return $ Var.mkTyVar name kind
56
57 -- Creates a binder for the given expression with the given name. This
58 -- works for both value and type level expressions, so it can return a Var or
59 -- TyVar (which is just an alias for Var).
60 mkBinderFor :: CoreExpr -> String -> TransformMonad Var.Var
61 mkBinderFor (Type ty) string = mkTypeVar string (Type.typeKind ty)
62 mkBinderFor expr string = mkInternalVar string (CoreUtils.exprType expr)
63
64 -- Creates a reference to the given variable. This works for both a normal
65 -- variable as well as a type variable
66 mkReferenceTo :: Var.Var -> CoreExpr
67 mkReferenceTo var | Var.isTyVar var = (Type $ Type.mkTyVarTy var)
68                   | otherwise       = (Var var)
69
70 cloneVar :: Var.Var -> TransformMonad Var.Var
71 cloneVar v = do
72   uniq <- mkUnique
73   -- Swap out the unique, and reset the IdInfo (I'm not 100% sure what it
74   -- contains, but vannillaIdInfo is always correct, since it means "no info").
75   return $ Var.lazySetVarIdInfo (Var.setVarUnique v uniq) IdInfo.vanillaIdInfo
76
77 -- Creates a new function with the same name as the given binder (but with a
78 -- new unique) and with the given function body. Returns the new binder for
79 -- this function.
80 mkFunction :: CoreBndr -> CoreExpr -> TransformMonad CoreBndr
81 mkFunction bndr body = do
82   let ty = CoreUtils.exprType body
83   id <- cloneVar bndr
84   let newid = Var.setVarType id ty
85   Trans.lift $ addGlobalBind newid body
86   return newid
87
88 -- Apply the given transformation to all expressions in the given expression,
89 -- including the expression itself.
90 everywhere :: (String, Transform) -> Transform
91 everywhere trans = applyboth (subeverywhere (everywhere trans)) trans
92
93 -- Apply the first transformation, followed by the second transformation, and
94 -- keep applying both for as long as expression still changes.
95 applyboth :: Transform -> (String, Transform) -> Transform
96 applyboth first (name, second) expr  = do
97   -- Apply the first
98   expr' <- first expr
99   -- Apply the second
100   (expr'', changed) <- Writer.listen $ second expr'
101   if Monoid.getAny $
102   --      trace ("Trying to apply transform " ++ name ++ " to:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n") $
103         changed 
104     then 
105 --      trace ("Applying transform " ++ name ++ " to:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n") $
106  --     trace ("Result of applying " ++ name ++ ":\n" ++ showSDoc (nest 4 $ ppr expr'') ++ "\n" ++ "Type: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr'') ++ "\n" ) $
107       applyboth first (name, second) $
108         expr'' 
109     else 
110     --  trace ("No changes") $
111       return expr''
112
113 -- Apply the given transformation to all direct subexpressions (only), not the
114 -- expression itself.
115 subeverywhere :: Transform -> Transform
116 subeverywhere trans (App a b) = do
117   a' <- trans a
118   b' <- trans b
119   return $ App a' b'
120
121 subeverywhere trans (Let (Rec binds) expr) = do
122   expr' <- trans expr
123   binds' <- mapM transbind binds
124   return $ Let (Rec binds') expr'
125   where
126     transbind :: (CoreBndr, CoreExpr) -> TransformMonad (CoreBndr, CoreExpr)
127     transbind (b, e) = do
128       e' <- trans e
129       return (b, e')
130
131 subeverywhere trans (Lam x expr) = do
132   expr' <- trans expr
133   return $ Lam x expr'
134
135 subeverywhere trans (Case scrut b t alts) = do
136   scrut' <- trans scrut
137   alts' <- mapM transalt alts
138   return $ Case scrut' b t alts'
139   where
140     transalt :: CoreAlt -> TransformMonad CoreAlt
141     transalt (con, binders, expr) = do
142       expr' <- trans expr
143       return (con, binders, expr')
144       
145
146 subeverywhere trans expr = return expr
147
148 -- Apply the given transformation to all expressions, except for every first
149 -- argument of an application.
150 notapplied :: (String, Transform) -> Transform
151 notapplied trans = applyboth (subnotapplied trans) trans
152
153 -- Apply the given transformation to all (direct and indirect) subexpressions
154 -- (but not the expression itself), except for the first argument of an
155 -- applicfirst argument of an application
156 subnotapplied :: (String, Transform) -> Transform
157 subnotapplied trans (App a b) = do
158   a' <- subnotapplied trans a
159   b' <- notapplied trans b
160   return $ App a' b'
161
162 -- Let subeverywhere handle all other expressions
163 subnotapplied trans expr = subeverywhere (notapplied trans) expr
164
165 -- Runs each of the transforms repeatedly inside the State monad.
166 dotransforms :: [Transform] -> CoreExpr -> TransformSession CoreExpr
167 dotransforms transs expr = do
168   (expr', changed) <- Writer.runWriterT $ Monad.foldM (flip ($)) expr transs
169   if Monoid.getAny changed then dotransforms transs expr' else return expr'
170
171 -- Inline all let bindings that satisfy the given condition
172 inlinebind :: ((CoreBndr, CoreExpr) -> Bool) -> Transform
173 inlinebind condition (Let (Rec binds) expr) | not $ null replace =
174     change newexpr
175   where 
176     -- Find all simple bindings
177     (replace, others) = List.partition condition binds
178     -- Substitute the to be replaced binders with their expression
179     newexpr = substitute replace (Let (Rec others) expr)
180 -- Leave all other expressions unchanged
181 inlinebind _ expr = return expr
182
183 -- Sets the changed flag in the TransformMonad, to signify that some
184 -- transform has changed the result
185 setChanged :: TransformMonad ()
186 setChanged = Writer.tell (Monoid.Any True)
187
188 -- Sets the changed flag and returns the given value.
189 change :: a -> TransformMonad a
190 change val = do
191   setChanged
192   return val
193
194 -- Create a new Unique
195 mkUnique :: TransformMonad Unique.Unique
196 mkUnique = Trans.lift $ do
197     us <- getA tsUniqSupply 
198     let (us', us'') = UniqSupply.splitUniqSupply us
199     putA tsUniqSupply us'
200     return $ UniqSupply.uniqFromSupply us''
201
202 -- Replace each of the binders given with the coresponding expressions in the
203 -- given expression.
204 substitute :: [(CoreBndr, CoreExpr)] -> CoreExpr -> CoreExpr
205 substitute replace expr = CoreSubst.substExpr subs expr
206     where subs = foldl (\s (b, e) -> CoreSubst.extendSubst s b e) CoreSubst.emptySubst replace
207
208 -- Run a given TransformSession. Used mostly to setup the right calls and
209 -- an initial state.
210 runTransformSession :: UniqSupply.UniqSupply -> TransformSession a -> a
211 runTransformSession uniqSupply session = State.evalState session initState
212                        where initState = TransformState uniqSupply Map.empty VarSet.emptyVarSet