1 {-# LANGUAGE PackageImports #-}
3 -- This module provides functions for program transformations.
5 module CLasH.Normalize.NormalizeTools where
9 import qualified Data.Monoid as Monoid
10 import qualified Data.Either as Either
11 import qualified Control.Arrow as Arrow
12 import qualified Control.Monad as Monad
13 import qualified Control.Monad.Trans.State as State
14 import qualified Control.Monad.Trans.Writer as Writer
15 import qualified "transformers" Control.Monad.Trans as Trans
16 import qualified Data.Map as Map
18 import Data.Accessor.MonadState as MonadState
22 import qualified UniqSupply
23 import qualified Unique
24 import qualified OccName
27 import qualified SrcLoc
29 import qualified IdInfo
30 import qualified CoreUtils
31 import qualified CoreSubst
32 import qualified VarSet
33 import qualified HscTypes
34 import Outputable ( showSDoc, ppr, nest )
37 import CLasH.Normalize.NormalizeTypes
38 import CLasH.Translator.TranslatorTypes
39 import CLasH.Utils.Pretty
40 import CLasH.VHDL.VHDLTypes
41 import qualified CLasH.VHDL.VHDLTools as VHDLTools
43 -- Create a new internal var with the given name and type. A Unique is
44 -- appended to the given name, to ensure uniqueness (not strictly neccesary,
45 -- since the Unique is also stored in the name, but this ensures variable
46 -- names are unique in the output).
47 mkInternalVar :: String -> Type.Type -> TransformMonad Var.Var
48 mkInternalVar str ty = Trans.lift (mkInternalVar' str ty)
50 mkInternalVar' :: String -> Type.Type -> TransformSession Var.Var
51 mkInternalVar' str ty = do
53 let occname = OccName.mkVarOcc (str ++ show uniq)
54 let name = Name.mkInternalName uniq occname SrcLoc.noSrcSpan
55 return $ Var.mkLocalVar IdInfo.VanillaId name ty IdInfo.vanillaIdInfo
57 -- Create a new type variable with the given name and kind. A Unique is
58 -- appended to the given name, to ensure uniqueness (not strictly neccesary,
59 -- since the Unique is also stored in the name, but this ensures variable
60 -- names are unique in the output).
61 mkTypeVar :: String -> Type.Kind -> TransformMonad Var.Var
62 mkTypeVar str kind = Trans.lift (mkTypeVar' str kind)
64 mkTypeVar' :: String -> Type.Kind -> TransformSession Var.Var
65 mkTypeVar' str kind = do
67 let occname = OccName.mkVarOcc (str ++ show uniq)
68 let name = Name.mkInternalName uniq occname SrcLoc.noSrcSpan
69 return $ Var.mkTyVar name kind
71 -- Creates a binder for the given expression with the given name. This
72 -- works for both value and type level expressions, so it can return a Var or
73 -- TyVar (which is just an alias for Var).
74 mkBinderFor :: CoreExpr -> String -> TransformMonad Var.Var
75 mkBinderFor expr string = Trans.lift (mkBinderFor' expr string)
77 mkBinderFor' :: CoreExpr -> String -> TransformSession Var.Var
78 mkBinderFor' (Type ty) string = mkTypeVar' string (Type.typeKind ty)
79 mkBinderFor' expr string = mkInternalVar' string (CoreUtils.exprType expr)
81 -- Creates a reference to the given variable. This works for both a normal
82 -- variable as well as a type variable
83 mkReferenceTo :: Var.Var -> CoreExpr
84 mkReferenceTo var | Var.isTyVar var = (Type $ Type.mkTyVarTy var)
85 | otherwise = (Var var)
87 cloneVar :: Var.Var -> TransformMonad Var.Var
90 -- Swap out the unique, and reset the IdInfo (I'm not 100% sure what it
91 -- contains, but vannillaIdInfo is always correct, since it means "no info").
92 return $ Var.lazySetIdInfo (Var.setVarUnique v uniq) IdInfo.vanillaIdInfo
94 -- Creates a new function with the same name as the given binder (but with a
95 -- new unique) and with the given function body. Returns the new binder for
97 mkFunction :: CoreBndr -> CoreExpr -> TransformMonad CoreBndr
98 mkFunction bndr body = do
99 let ty = CoreUtils.exprType body
101 let newid = Var.setVarType id ty
102 Trans.lift $ addGlobalBind newid body
105 -- Apply the given transformation to all expressions in the given expression,
106 -- including the expression itself.
107 everywhere :: (String, Transform) -> Transform
108 everywhere trans = applyboth (subeverywhere (everywhere trans)) trans
110 -- Apply the first transformation, followed by the second transformation, and
111 -- keep applying both for as long as expression still changes.
112 applyboth :: Transform -> (String, Transform) -> Transform
113 applyboth first (name, second) expr = do
117 (expr'', changed) <- Writer.listen $ second expr'
119 -- trace ("Trying to apply transform " ++ name ++ " to:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n") $
122 -- trace ("Applying transform " ++ name ++ " to:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n") $
123 -- trace ("Result of applying " ++ name ++ ":\n" ++ showSDoc (nest 4 $ ppr expr'') ++ "\n" ++ "Type: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr'') ++ "\n" ) $
124 applyboth first (name, second) $
127 -- trace ("No changes") $
130 -- Apply the given transformation to all direct subexpressions (only), not the
131 -- expression itself.
132 subeverywhere :: Transform -> Transform
133 subeverywhere trans (App a b) = do
138 subeverywhere trans (Let (NonRec b bexpr) expr) = do
139 bexpr' <- trans bexpr
141 return $ Let (NonRec b bexpr') expr'
143 subeverywhere trans (Let (Rec binds) expr) = do
145 binds' <- mapM transbind binds
146 return $ Let (Rec binds') expr'
148 transbind :: (CoreBndr, CoreExpr) -> TransformMonad (CoreBndr, CoreExpr)
149 transbind (b, e) = do
153 subeverywhere trans (Lam x expr) = do
157 subeverywhere trans (Case scrut b t alts) = do
158 scrut' <- trans scrut
159 alts' <- mapM transalt alts
160 return $ Case scrut' b t alts'
162 transalt :: CoreAlt -> TransformMonad CoreAlt
163 transalt (con, binders, expr) = do
165 return (con, binders, expr')
167 subeverywhere trans (Var x) = return $ Var x
168 subeverywhere trans (Lit x) = return $ Lit x
169 subeverywhere trans (Type x) = return $ Type x
171 subeverywhere trans (Cast expr ty) = do
173 return $ Cast expr' ty
175 subeverywhere trans expr = error $ "\nNormalizeTools.subeverywhere: Unsupported expression: " ++ show expr
177 -- Apply the given transformation to all expressions, except for direct
178 -- arguments of an application
179 notappargs :: (String, Transform) -> Transform
180 notappargs trans = applyboth (subnotappargs trans) trans
182 -- Apply the given transformation to all (direct and indirect) subexpressions
183 -- (but not the expression itself), except for direct arguments of an
185 subnotappargs :: (String, Transform) -> Transform
186 subnotappargs trans (App a b) = do
187 a' <- subnotappargs trans a
188 b' <- subnotappargs trans b
191 -- Let subeverywhere handle all other expressions
192 subnotappargs trans expr = subeverywhere (notappargs trans) expr
194 -- Runs each of the transforms repeatedly inside the State monad.
195 dotransforms :: [Transform] -> CoreExpr -> TransformSession CoreExpr
196 dotransforms transs expr = do
197 (expr', changed) <- Writer.runWriterT $ Monad.foldM (flip ($)) expr transs
198 if Monoid.getAny changed then dotransforms transs expr' else return expr'
200 -- Inline all let bindings that satisfy the given condition
201 inlinebind :: ((CoreBndr, CoreExpr) -> TransformMonad Bool) -> Transform
202 inlinebind condition expr@(Let (Rec binds) res) = do
203 -- Find all bindings that adhere to the condition
204 res_eithers <- mapM docond binds
205 case Either.partitionEithers res_eithers of
206 -- No replaces? No change
207 ([], _) -> return expr
208 (replace, others) -> do
209 -- Substitute the to be replaced binders with their expression
210 let newexpr = substitute replace (Let (Rec others) res)
213 docond :: (CoreBndr, CoreExpr) -> TransformMonad (Either (CoreBndr, CoreExpr) (CoreBndr, CoreExpr))
216 return $ case res of True -> Left b; False -> Right b
218 -- Leave all other expressions unchanged
219 inlinebind _ expr = return expr
221 -- Sets the changed flag in the TransformMonad, to signify that some
222 -- transform has changed the result
223 setChanged :: TransformMonad ()
224 setChanged = Writer.tell (Monoid.Any True)
226 -- Sets the changed flag and returns the given value.
227 change :: a -> TransformMonad a
232 -- Create a new Unique
233 mkUnique :: TransformMonad Unique.Unique
234 mkUnique = Trans.lift $ mkUnique'
236 mkUnique' :: TransformSession Unique.Unique
238 us <- getA tsUniqSupply
239 let (us', us'') = UniqSupply.splitUniqSupply us
240 putA tsUniqSupply us'
241 return $ UniqSupply.uniqFromSupply us''
243 -- Replace each of the binders given with the coresponding expressions in the
245 substitute :: [(CoreBndr, CoreExpr)] -> CoreExpr -> CoreExpr
246 substitute [] expr = expr
247 -- Apply one substitution on the expression, but also on any remaining
248 -- substitutions. This seems to be the only way to handle substitutions like
249 -- [(b, c), (a, b)]. This means we reuse a substitution, which is not allowed
250 -- according to CoreSubst documentation (but it doesn't seem to be a problem).
251 -- TODO: Find out how this works, exactly.
252 substitute ((b, e):subss) expr = substitute subss' expr'
255 subs = (CoreSubst.extendSubst CoreSubst.emptySubst b e)
256 -- Apply this substitution to the main expression
257 expr' = CoreSubst.substExpr subs expr
258 -- Apply this substitution on all the expressions in the remaining
260 subss' = map (Arrow.second (CoreSubst.substExpr subs)) subss
262 -- Is the given expression representable at runtime, based on the type?
263 isRepr :: CoreSyn.CoreExpr -> TransformMonad Bool
264 isRepr (Type ty) = return False
265 isRepr expr = Trans.lift $ MonadState.lift tsType $ VHDLTools.isReprType (CoreUtils.exprType expr)
267 is_local_var :: CoreSyn.CoreExpr -> TransformSession Bool
268 is_local_var (CoreSyn.Var v) = do
269 bndrs <- getGlobalBinders
270 return $ not $ v `elem` bndrs
271 is_local_var _ = return False