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.Utils.Pretty
39 import CLasH.VHDL.VHDLTypes
40 import qualified CLasH.VHDL.VHDLTools as VHDLTools
42 -- Create a new internal var with the given name and type. A Unique is
43 -- appended to the given name, to ensure uniqueness (not strictly neccesary,
44 -- since the Unique is also stored in the name, but this ensures variable
45 -- names are unique in the output).
46 mkInternalVar :: String -> Type.Type -> TransformMonad Var.Var
47 mkInternalVar str ty = Trans.lift (mkInternalVar' str ty)
49 mkInternalVar' :: String -> Type.Type -> TransformSession Var.Var
50 mkInternalVar' str ty = do
52 let occname = OccName.mkVarOcc (str ++ show uniq)
53 let name = Name.mkInternalName uniq occname SrcLoc.noSrcSpan
54 return $ Var.mkLocalVar IdInfo.VanillaId name ty IdInfo.vanillaIdInfo
56 -- Create a new type variable with the given name and kind. A Unique is
57 -- appended to the given name, to ensure uniqueness (not strictly neccesary,
58 -- since the Unique is also stored in the name, but this ensures variable
59 -- names are unique in the output).
60 mkTypeVar :: String -> Type.Kind -> TransformMonad Var.Var
61 mkTypeVar str kind = Trans.lift (mkTypeVar' str kind)
63 mkTypeVar' :: String -> Type.Kind -> TransformSession Var.Var
64 mkTypeVar' str kind = do
66 let occname = OccName.mkVarOcc (str ++ show uniq)
67 let name = Name.mkInternalName uniq occname SrcLoc.noSrcSpan
68 return $ Var.mkTyVar name kind
70 -- Creates a binder for the given expression with the given name. This
71 -- works for both value and type level expressions, so it can return a Var or
72 -- TyVar (which is just an alias for Var).
73 mkBinderFor :: CoreExpr -> String -> TransformMonad Var.Var
74 mkBinderFor expr string = Trans.lift (mkBinderFor' expr string)
76 mkBinderFor' :: CoreExpr -> String -> TransformSession Var.Var
77 mkBinderFor' (Type ty) string = mkTypeVar' string (Type.typeKind ty)
78 mkBinderFor' expr string = mkInternalVar' string (CoreUtils.exprType expr)
80 -- Creates a reference to the given variable. This works for both a normal
81 -- variable as well as a type variable
82 mkReferenceTo :: Var.Var -> CoreExpr
83 mkReferenceTo var | Var.isTyVar var = (Type $ Type.mkTyVarTy var)
84 | otherwise = (Var var)
86 cloneVar :: Var.Var -> TransformMonad Var.Var
89 -- Swap out the unique, and reset the IdInfo (I'm not 100% sure what it
90 -- contains, but vannillaIdInfo is always correct, since it means "no info").
91 return $ Var.lazySetIdInfo (Var.setVarUnique v uniq) IdInfo.vanillaIdInfo
93 -- Creates a new function with the same name as the given binder (but with a
94 -- new unique) and with the given function body. Returns the new binder for
96 mkFunction :: CoreBndr -> CoreExpr -> TransformMonad CoreBndr
97 mkFunction bndr body = do
98 let ty = CoreUtils.exprType body
100 let newid = Var.setVarType id ty
101 Trans.lift $ addGlobalBind newid body
104 -- Apply the given transformation to all expressions in the given expression,
105 -- including the expression itself.
106 everywhere :: (String, Transform) -> Transform
107 everywhere trans = applyboth (subeverywhere (everywhere trans)) trans
109 -- Apply the first transformation, followed by the second transformation, and
110 -- keep applying both for as long as expression still changes.
111 applyboth :: Transform -> (String, Transform) -> Transform
112 applyboth first (name, second) expr = do
116 (expr'', changed) <- Writer.listen $ second expr'
118 -- trace ("Trying to apply transform " ++ name ++ " to:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n") $
121 -- trace ("Applying transform " ++ name ++ " to:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n") $
122 -- trace ("Result of applying " ++ name ++ ":\n" ++ showSDoc (nest 4 $ ppr expr'') ++ "\n" ++ "Type: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr'') ++ "\n" ) $
123 applyboth first (name, second) $
126 -- trace ("No changes") $
129 -- Apply the given transformation to all direct subexpressions (only), not the
130 -- expression itself.
131 subeverywhere :: Transform -> Transform
132 subeverywhere trans (App a b) = do
137 subeverywhere trans (Let (NonRec b bexpr) expr) = do
138 bexpr' <- trans bexpr
140 return $ Let (NonRec b bexpr') expr'
142 subeverywhere trans (Let (Rec binds) expr) = do
144 binds' <- mapM transbind binds
145 return $ Let (Rec binds') expr'
147 transbind :: (CoreBndr, CoreExpr) -> TransformMonad (CoreBndr, CoreExpr)
148 transbind (b, e) = do
152 subeverywhere trans (Lam x expr) = do
156 subeverywhere trans (Case scrut b t alts) = do
157 scrut' <- trans scrut
158 alts' <- mapM transalt alts
159 return $ Case scrut' b t alts'
161 transalt :: CoreAlt -> TransformMonad CoreAlt
162 transalt (con, binders, expr) = do
164 return (con, binders, expr')
166 subeverywhere trans (Var x) = return $ Var x
167 subeverywhere trans (Lit x) = return $ Lit x
168 subeverywhere trans (Type x) = return $ Type x
170 subeverywhere trans (Cast expr ty) = do
172 return $ Cast expr' ty
174 subeverywhere trans expr = error $ "\nNormalizeTools.subeverywhere: Unsupported expression: " ++ show expr
176 -- Apply the given transformation to all expressions, except for direct
177 -- arguments of an application
178 notappargs :: (String, Transform) -> Transform
179 notappargs trans = applyboth (subnotappargs trans) trans
181 -- Apply the given transformation to all (direct and indirect) subexpressions
182 -- (but not the expression itself), except for direct arguments of an
184 subnotappargs :: (String, Transform) -> Transform
185 subnotappargs trans (App a b) = do
186 a' <- subnotappargs trans a
187 b' <- subnotappargs trans b
190 -- Let subeverywhere handle all other expressions
191 subnotappargs trans expr = subeverywhere (notappargs trans) expr
193 -- Runs each of the transforms repeatedly inside the State monad.
194 dotransforms :: [Transform] -> CoreExpr -> TransformSession CoreExpr
195 dotransforms transs expr = do
196 (expr', changed) <- Writer.runWriterT $ Monad.foldM (flip ($)) expr transs
197 if Monoid.getAny changed then dotransforms transs expr' else return expr'
199 -- Inline all let bindings that satisfy the given condition
200 inlinebind :: ((CoreBndr, CoreExpr) -> TransformMonad Bool) -> Transform
201 inlinebind condition expr@(Let (Rec binds) res) = do
202 -- Find all bindings that adhere to the condition
203 res_eithers <- mapM docond binds
204 case Either.partitionEithers res_eithers of
205 -- No replaces? No change
206 ([], _) -> return expr
207 (replace, others) -> do
208 -- Substitute the to be replaced binders with their expression
209 let newexpr = substitute replace (Let (Rec others) res)
212 docond :: (CoreBndr, CoreExpr) -> TransformMonad (Either (CoreBndr, CoreExpr) (CoreBndr, CoreExpr))
215 return $ case res of True -> Left b; False -> Right b
217 -- Leave all other expressions unchanged
218 inlinebind _ expr = return expr
220 -- Sets the changed flag in the TransformMonad, to signify that some
221 -- transform has changed the result
222 setChanged :: TransformMonad ()
223 setChanged = Writer.tell (Monoid.Any True)
225 -- Sets the changed flag and returns the given value.
226 change :: a -> TransformMonad a
231 -- Create a new Unique
232 mkUnique :: TransformMonad Unique.Unique
233 mkUnique = Trans.lift $ mkUnique'
235 mkUnique' :: TransformSession Unique.Unique
237 us <- getA tsUniqSupply
238 let (us', us'') = UniqSupply.splitUniqSupply us
239 putA tsUniqSupply us'
240 return $ UniqSupply.uniqFromSupply us''
242 -- Replace each of the binders given with the coresponding expressions in the
244 substitute :: [(CoreBndr, CoreExpr)] -> CoreExpr -> CoreExpr
245 substitute [] expr = expr
246 -- Apply one substitution on the expression, but also on any remaining
247 -- substitutions. This seems to be the only way to handle substitutions like
248 -- [(b, c), (a, b)]. This means we reuse a substitution, which is not allowed
249 -- according to CoreSubst documentation (but it doesn't seem to be a problem).
250 -- TODO: Find out how this works, exactly.
251 substitute ((b, e):subss) expr = substitute subss' expr'
254 subs = (CoreSubst.extendSubst CoreSubst.emptySubst b e)
255 -- Apply this substitution to the main expression
256 expr' = CoreSubst.substExpr subs expr
257 -- Apply this substitution on all the expressions in the remaining
259 subss' = map (Arrow.second (CoreSubst.substExpr subs)) subss
261 -- Run a given TransformSession. Used mostly to setup the right calls and
263 runTransformSession :: HscTypes.HscEnv -> UniqSupply.UniqSupply -> TransformSession a -> a
264 runTransformSession env uniqSupply session = State.evalState session emptyTransformState
266 emptyTypeState = TypeState Map.empty [] Map.empty Map.empty env
267 emptyTransformState = TransformState uniqSupply Map.empty VarSet.emptyVarSet emptyTypeState
269 -- Is the given expression representable at runtime, based on the type?
270 isRepr :: CoreSyn.CoreExpr -> TransformMonad Bool
271 isRepr (Type ty) = return False
272 isRepr expr = Trans.lift $ MonadState.lift tsType $ VHDLTools.isReprType (CoreUtils.exprType expr)
274 is_local_var :: CoreSyn.CoreExpr -> TransformSession Bool
275 is_local_var (CoreSyn.Var v) = do
276 bndrs <- getGlobalBinders
277 return $ not $ v `elem` bndrs
278 is_local_var _ = return False