Remove two old debug traces.
[matthijs/master-project/cλash.git] / Normalize.hs
1 {-# LANGUAGE PackageImports #-}
2 --
3 -- Functions to bring a Core expression in normal form. This module provides a
4 -- top level function "normalize", and defines the actual transformation passes that
5 -- are performed.
6 --
7 module Normalize (normalizeModule) where
8
9 -- Standard modules
10 import Debug.Trace
11 import qualified Maybe
12 import qualified "transformers" Control.Monad.Trans as Trans
13 import qualified Control.Monad as Monad
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 CoreUtils
21 import qualified Type
22 import qualified Id
23 import qualified Var
24 import qualified VarSet
25 import qualified CoreFVs
26 import Outputable ( showSDoc, ppr, nest )
27
28 -- Local imports
29 import NormalizeTypes
30 import NormalizeTools
31 import CoreTools
32
33 --------------------------------
34 -- Start of transformations
35 --------------------------------
36
37 --------------------------------
38 -- η abstraction
39 --------------------------------
40 eta, etatop :: Transform
41 eta expr | is_fun expr && not (is_lam expr) = do
42   let arg_ty = (fst . Type.splitFunTy . CoreUtils.exprType) expr
43   id <- mkInternalVar "param" arg_ty
44   change (Lam id (App expr (Var id)))
45 -- Leave all other expressions unchanged
46 eta e = return e
47 etatop = notapplied ("eta", eta)
48
49 --------------------------------
50 -- β-reduction
51 --------------------------------
52 beta, betatop :: Transform
53 -- Substitute arg for x in expr
54 beta (App (Lam x expr) arg) = change $ substitute [(x, arg)] expr
55 -- Propagate the application into the let
56 beta (App (Let binds expr) arg) = change $ Let binds (App expr arg)
57 -- Propagate the application into each of the alternatives
58 beta (App (Case scrut b ty alts) arg) = change $ Case scrut b ty' alts'
59   where 
60     alts' = map (\(con, bndrs, expr) -> (con, bndrs, (App expr arg))) alts
61     (_, ty') = Type.splitFunTy ty
62 -- Leave all other expressions unchanged
63 beta expr = return expr
64 -- Perform this transform everywhere
65 betatop = everywhere ("beta", beta)
66
67 --------------------------------
68 -- let recursification
69 --------------------------------
70 letrec, letrectop :: Transform
71 letrec (Let (NonRec b expr) res) = change $ Let (Rec [(b, expr)]) res
72 -- Leave all other expressions unchanged
73 letrec expr = return expr
74 -- Perform this transform everywhere
75 letrectop = everywhere ("letrec", letrec)
76
77 --------------------------------
78 -- let simplification
79 --------------------------------
80 letsimpl, letsimpltop :: Transform
81 -- Don't simplifiy lets that are already simple
82 letsimpl expr@(Let _ (Var _)) = return expr
83 -- Put the "in ..." value of a let in its own binding, but not when the
84 -- expression has a function type (to prevent loops with inlinefun).
85 letsimpl (Let (Rec binds) expr) | not $ is_fun expr = do
86   id <- mkInternalVar "foo" (CoreUtils.exprType expr)
87   let bind = (id, expr)
88   change $ Let (Rec (bind:binds)) (Var id)
89 -- Leave all other expressions unchanged
90 letsimpl expr = return expr
91 -- Perform this transform everywhere
92 letsimpltop = everywhere ("letsimpl", letsimpl)
93
94 --------------------------------
95 -- let flattening
96 --------------------------------
97 letflat, letflattop :: Transform
98 letflat (Let (Rec binds) expr) = do
99   -- Turn each binding into a list of bindings (possibly containing just one
100   -- element, of course)
101   bindss <- Monad.mapM flatbind binds
102   -- Concat all the bindings
103   let binds' = concat bindss
104   -- Return the new let. We don't use change here, since possibly nothing has
105   -- changed. If anything has changed, flatbind has already flagged that
106   -- change.
107   return $ Let (Rec binds') expr
108   where
109     -- Turns a binding of a let into a multiple bindings, or any other binding
110     -- into a list with just that binding
111     flatbind :: (CoreBndr, CoreExpr) -> TransformMonad [(CoreBndr, CoreExpr)]
112     flatbind (b, Let (Rec binds) expr) = change ((b, expr):binds)
113     flatbind (b, expr) = return [(b, expr)]
114 -- Leave all other expressions unchanged
115 letflat expr = return expr
116 -- Perform this transform everywhere
117 letflattop = everywhere ("letflat", letflat)
118
119 --------------------------------
120 -- Simple let binding removal
121 --------------------------------
122 -- Remove a = b bindings from let expressions everywhere
123 letremovetop :: Transform
124 letremovetop = everywhere ("letremove", inlinebind (\(b, e) -> case e of (Var v) -> True; otherwise -> False))
125
126 --------------------------------
127 -- Function inlining
128 --------------------------------
129 -- Remove a = B bindings, with B :: a -> b, from let expressions everywhere.
130 -- This is a tricky function, which is prone to create loops in the
131 -- transformations. To fix this, we make sure that no transformation will
132 -- create a new let binding with a function type. These other transformations
133 -- will just not work on those function-typed values at first, but the other
134 -- transformations (in particular β-reduction) should make sure that the type
135 -- of those values eventually becomes primitive.
136 inlinefuntop :: Transform
137 inlinefuntop = everywhere ("inlinefun", inlinebind (Type.isFunTy . CoreUtils.exprType . snd))
138
139 --------------------------------
140 -- Scrutinee simplification
141 --------------------------------
142 scrutsimpl,scrutsimpltop :: Transform
143 -- Don't touch scrutinees that are already simple
144 scrutsimpl expr@(Case (Var _) _ _ _) = return expr
145 -- Replace all other cases with a let that binds the scrutinee and a new
146 -- simple scrutinee, but not when the scrutinee is a function type (to prevent
147 -- loops with inlinefun, though I don't think a scrutinee can have a function
148 -- type...)
149 scrutsimpl (Case scrut b ty alts) | not $ is_fun scrut = do
150   id <- mkInternalVar "scrut" (CoreUtils.exprType scrut)
151   change $ Let (Rec [(id, scrut)]) (Case (Var id) b ty alts)
152 -- Leave all other expressions unchanged
153 scrutsimpl expr = return expr
154 -- Perform this transform everywhere
155 scrutsimpltop = everywhere ("scrutsimpl", scrutsimpl)
156
157 --------------------------------
158 -- Case binder wildening
159 --------------------------------
160 casewild, casewildtop :: Transform
161 casewild expr@(Case scrut b ty alts) = do
162   (bindingss, alts') <- (Monad.liftM unzip) $ mapM doalt alts
163   let bindings = concat bindingss
164   -- Replace the case with a let with bindings and a case
165   let newlet = (Let (Rec bindings) (Case scrut b ty alts'))
166   -- If there are no non-wild binders, or this case is already a simple
167   -- selector (i.e., a single alt with exactly one binding), already a simple
168   -- selector altan no bindings (i.e., no wild binders in the original case),
169   -- don't change anything, otherwise, replace the case.
170   if null bindings || length alts == 1 && length bindings == 1 then return expr else change newlet 
171   where
172   -- Generate a single wild binder, since they are all the same
173   wild = Id.mkWildId
174   -- Wilden the binders of one alt, producing a list of bindings as a
175   -- sideeffect.
176   doalt :: CoreAlt -> TransformMonad ([(CoreBndr, CoreExpr)], CoreAlt)
177   doalt (con, bndrs, expr) = do
178     bindings_maybe <- Monad.zipWithM mkextracts bndrs [0..]
179     let bindings = Maybe.catMaybes bindings_maybe
180     -- We replace the binders with wild binders only. We can leave expr
181     -- unchanged, since the new bindings bind the same vars as the original
182     -- did.
183     let newalt = (con, wildbndrs, expr)
184     return (bindings, newalt)
185     where
186       -- Make all binders wild
187       wildbndrs = map (\bndr -> Id.mkWildId (Id.idType bndr)) bndrs
188       -- Creates a case statement to retrieve the ith element from the scrutinee
189       -- and binds that to b.
190       mkextracts :: CoreBndr -> Int -> TransformMonad (Maybe (CoreBndr, CoreExpr))
191       mkextracts b i =
192         if is_wild b || Type.isFunTy (Id.idType b) 
193           -- Don't create extra bindings for binders that are already wild, or
194           -- for binders that bind function types (to prevent loops with
195           -- inlinefun).
196           then return Nothing
197           else do
198             -- Create on new binder that will actually capture a value in this
199             -- case statement, and return it
200             let bty = (Id.idType b)
201             id <- mkInternalVar "sel" bty
202             let binders = take i wildbndrs ++ [id] ++ drop (i+1) wildbndrs
203             return $ Just (b, Case scrut b bty [(con, binders, Var id)])
204 -- Leave all other expressions unchanged
205 casewild expr = return expr
206 -- Perform this transform everywhere
207 casewildtop = everywhere ("casewild", casewild)
208
209 --------------------------------
210 -- Case value simplification
211 --------------------------------
212 casevalsimpl, casevalsimpltop :: Transform
213 casevalsimpl expr@(Case scrut b ty alts) = do
214   -- Try to simplify each alternative, resulting in an optional binding and a
215   -- new alternative.
216   (bindings_maybe, alts') <- (Monad.liftM unzip) $ mapM doalt alts
217   let bindings = Maybe.catMaybes bindings_maybe
218   -- Create a new let around the case, that binds of the cases values.
219   let newlet = Let (Rec bindings) (Case scrut b ty alts')
220   -- If there were no values that needed and allowed simplification, don't
221   -- change the case.
222   if null bindings then return expr else change newlet 
223   where
224     doalt :: CoreAlt -> TransformMonad (Maybe (CoreBndr, CoreExpr), CoreAlt)
225     -- Don't simplify values that are already simple
226     doalt alt@(con, bndrs, Var _) = return (Nothing, alt)
227     -- Simplify each alt by creating a new id, binding the case value to it and
228     -- replacing the case value with that id. Only do this when the case value
229     -- does not use any of the binders bound by this alternative, for that would
230     -- cause those binders to become unbound when moving the value outside of
231     -- the case statement. Also, don't create a binding for function-typed
232     -- expressions, to prevent loops with inlinefun.
233     doalt (con, bndrs, expr) | (not usesvars) && (not $ is_fun expr) = do
234       id <- mkInternalVar "caseval" (CoreUtils.exprType expr)
235       -- We don't flag a change here, since casevalsimpl will do that above
236       -- based on Just we return here.
237       return $ (Just (id, expr), (con, bndrs, Var id))
238       -- Find if any of the binders are used by expr
239       where usesvars = (not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))) expr
240     -- Don't simplify anything else
241     doalt alt = return (Nothing, alt)
242 -- Leave all other expressions unchanged
243 casevalsimpl expr = return expr
244 -- Perform this transform everywhere
245 casevalsimpltop = everywhere ("casevalsimpl", casevalsimpl)
246
247 --------------------------------
248 -- Case removal
249 --------------------------------
250 -- Remove case statements that have only a single alternative and only wild
251 -- binders.
252 caseremove, caseremovetop :: Transform
253 -- Replace a useless case by the value of its single alternative
254 caseremove (Case scrut b ty [(con, bndrs, expr)]) | not usesvars = change expr
255     -- Find if any of the binders are used by expr
256     where usesvars = (not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))) expr
257 -- Leave all other expressions unchanged
258 caseremove expr = return expr
259 -- Perform this transform everywhere
260 caseremovetop = everywhere ("caseremove", caseremove)
261
262 --------------------------------
263 -- Application simplification
264 --------------------------------
265 -- Make sure that all arguments in an application are simple variables.
266 appsimpl, appsimpltop :: Transform
267 -- Don't simplify arguments that are already simple
268 appsimpl expr@(App f (Var _)) = return expr
269 -- Simplify all arguments that do not have a function type (to prevent loops
270 -- with inlinefun) and is not a type argument. Do this by introducing a new
271 -- Let that binds the argument and passing the new binder in the application.
272 appsimpl (App f expr) | (not $ is_fun expr) && (not $ CoreSyn.isTypeArg expr) = do
273   id <- mkInternalVar "arg" (CoreUtils.exprType expr)
274   change $ Let (Rec [(id, expr)]) (App f (Var id))
275 -- Leave all other expressions unchanged
276 appsimpl expr = return expr
277 -- Perform this transform everywhere
278 appsimpltop = everywhere ("appsimpl", appsimpl)
279
280
281 --------------------------------
282 -- Type argument propagation
283 --------------------------------
284 -- Remove all applications to type arguments, by duplicating the function
285 -- called with the type application in its new definition. We leave
286 -- dictionaries that might be associated with the type untouched, the funprop
287 -- transform should propagate these later on.
288 typeprop, typeproptop :: Transform
289 -- Transform any function that is applied to a type argument. Since type
290 -- arguments are always the first ones to apply and we'll remove all type
291 -- arguments, we can simply do them one by one.
292 typeprop expr@(App (Var f) (Type ty)) = do
293   id <- cloneVar f
294   let newty = Type.applyTy (Id.idType f) ty
295   let newf = Var.setVarType id newty
296   body_maybe <- Trans.lift $ getGlobalBind f
297   case body_maybe of
298     Just body -> do
299       let newbody = App body (Type ty)
300       Trans.lift $ addGlobalBind newf newbody
301       change (Var newf)
302     -- If we don't have a body for the function called, leave it unchanged (it
303     -- should be a primitive function then).
304     Nothing -> return expr
305 -- Leave all other expressions unchanged
306 typeprop expr = return expr
307 -- Perform this transform everywhere
308 typeproptop = everywhere ("typeprop", typeprop)
309
310 -- TODO: introduce top level let if needed?
311
312 --------------------------------
313 -- End of transformations
314 --------------------------------
315
316
317
318
319 -- What transforms to run?
320 transforms = [typeproptop, etatop, betatop, letremovetop, letrectop, letsimpltop, letflattop, casewildtop, scrutsimpltop, casevalsimpltop, caseremovetop, inlinefuntop, appsimpltop]
321
322 -- Turns the given bind into VHDL
323 normalizeModule :: 
324   UniqSupply.UniqSupply -- ^ A UniqSupply we can use
325   -> [(CoreBndr, CoreExpr)]  -- ^ All bindings we know (i.e., in the current module)
326   -> [CoreBndr]  -- ^ The bindings to generate VHDL for (i.e., the top level bindings)
327   -> [Bool] -- ^ For each of the bindings to generate VHDL for, if it is stateful
328   -> [(CoreBndr, CoreExpr)] -- ^ The resulting VHDL
329
330 normalizeModule uniqsupply bindings generate_for statefuls = runTransformSession uniqsupply $ do
331   -- Put all the bindings in this module in the tsBindings map
332   putA tsBindings (Map.fromList bindings)
333   -- (Recursively) normalize each of the requested bindings
334   mapM normalizeBind generate_for
335   -- Get all initial bindings and the ones we produced
336   bindings_map <- getA tsBindings
337   let bindings = Map.assocs bindings_map
338   normalized_bindings <- getA tsNormalized
339   -- But return only the normalized bindings
340   return $ filter ((flip VarSet.elemVarSet normalized_bindings) . fst) bindings
341
342 normalizeBind :: CoreBndr -> TransformSession ()
343 normalizeBind bndr = do
344   normalized_funcs <- getA tsNormalized
345   -- See if this function was normalized already
346   if VarSet.elemVarSet bndr normalized_funcs
347     then
348       -- Yup, don't do it again
349       return ()
350     else do
351       -- Nope, note that it has been and do it.
352       modA tsNormalized (flip VarSet.extendVarSet bndr)
353       expr_maybe <- getGlobalBind bndr
354       case expr_maybe of 
355         Just expr -> do
356           -- Normalize this expression
357           expr' <- dotransforms transforms expr
358           let expr'' = trace ("Before:\n\n" ++ showSDoc ( ppr expr ) ++ "\n\nAfter:\n\n" ++ showSDoc ( ppr expr')) expr'
359           -- And store the normalized version in the session
360           modA tsBindings (Map.insert bndr expr'')
361           -- Find all vars used with a function type. All of these should be global
362           -- binders (i.e., functions used), since any local binders with a function
363           -- type should have been inlined already.
364           let used_funcs_set = CoreFVs.exprSomeFreeVars (\v -> (Type.isFunTy . snd . Type.splitForAllTys . Id.idType) v) expr''
365           let used_funcs = VarSet.varSetElems used_funcs_set
366           -- Process each of the used functions recursively
367           mapM normalizeBind used_funcs
368           return ()
369         -- We don't have a value for this binder, let's assume this is a builtin
370         -- function. This might need some extra checking and a nice error
371         -- message).
372         Nothing -> return ()