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 = do
49 let occname = OccName.mkVarOcc (str ++ show uniq)
50 let name = Name.mkInternalName uniq occname SrcLoc.noSrcSpan
51 return $ Var.mkLocalVar IdInfo.VanillaId name ty IdInfo.vanillaIdInfo
53 -- Create a new type variable with the given name and kind. A Unique is
54 -- appended to the given name, to ensure uniqueness (not strictly neccesary,
55 -- since the Unique is also stored in the name, but this ensures variable
56 -- names are unique in the output).
57 mkTypeVar :: String -> Type.Kind -> TransformMonad Var.Var
58 mkTypeVar str kind = do
60 let occname = OccName.mkVarOcc (str ++ show uniq)
61 let name = Name.mkInternalName uniq occname SrcLoc.noSrcSpan
62 return $ Var.mkTyVar name kind
64 -- Creates a binder for the given expression with the given name. This
65 -- works for both value and type level expressions, so it can return a Var or
66 -- TyVar (which is just an alias for Var).
67 mkBinderFor :: CoreExpr -> String -> TransformMonad Var.Var
68 mkBinderFor (Type ty) string = mkTypeVar string (Type.typeKind ty)
69 mkBinderFor expr string = mkInternalVar string (CoreUtils.exprType expr)
71 -- Creates a reference to the given variable. This works for both a normal
72 -- variable as well as a type variable
73 mkReferenceTo :: Var.Var -> CoreExpr
74 mkReferenceTo var | Var.isTyVar var = (Type $ Type.mkTyVarTy var)
75 | otherwise = (Var var)
77 cloneVar :: Var.Var -> TransformMonad Var.Var
80 -- Swap out the unique, and reset the IdInfo (I'm not 100% sure what it
81 -- contains, but vannillaIdInfo is always correct, since it means "no info").
82 return $ Var.lazySetIdInfo (Var.setVarUnique v uniq) IdInfo.vanillaIdInfo
84 -- Creates a new function with the same name as the given binder (but with a
85 -- new unique) and with the given function body. Returns the new binder for
87 mkFunction :: CoreBndr -> CoreExpr -> TransformMonad CoreBndr
88 mkFunction bndr body = do
89 let ty = CoreUtils.exprType body
91 let newid = Var.setVarType id ty
92 Trans.lift $ addGlobalBind newid body
95 -- Apply the given transformation to all expressions in the given expression,
96 -- including the expression itself.
97 everywhere :: (String, Transform) -> Transform
98 everywhere trans = applyboth (subeverywhere (everywhere trans)) trans
100 -- Apply the first transformation, followed by the second transformation, and
101 -- keep applying both for as long as expression still changes.
102 applyboth :: Transform -> (String, Transform) -> Transform
103 applyboth first (name, second) expr = do
107 (expr'', changed) <- Writer.listen $ second expr'
109 -- trace ("Trying to apply transform " ++ name ++ " to:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n") $
112 -- trace ("Applying transform " ++ name ++ " to:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n") $
113 -- trace ("Result of applying " ++ name ++ ":\n" ++ showSDoc (nest 4 $ ppr expr'') ++ "\n" ++ "Type: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr'') ++ "\n" ) $
114 applyboth first (name, second) $
117 -- trace ("No changes") $
120 -- Apply the given transformation to all direct subexpressions (only), not the
121 -- expression itself.
122 subeverywhere :: Transform -> Transform
123 subeverywhere trans (App a b) = do
128 subeverywhere trans (Let (NonRec b bexpr) expr) = do
129 bexpr' <- trans bexpr
131 return $ Let (NonRec b bexpr') expr'
133 subeverywhere trans (Let (Rec binds) expr) = do
135 binds' <- mapM transbind binds
136 return $ Let (Rec binds') expr'
138 transbind :: (CoreBndr, CoreExpr) -> TransformMonad (CoreBndr, CoreExpr)
139 transbind (b, e) = do
143 subeverywhere trans (Lam x expr) = do
147 subeverywhere trans (Case scrut b t alts) = do
148 scrut' <- trans scrut
149 alts' <- mapM transalt alts
150 return $ Case scrut' b t alts'
152 transalt :: CoreAlt -> TransformMonad CoreAlt
153 transalt (con, binders, expr) = do
155 return (con, binders, expr')
157 subeverywhere trans (Var x) = return $ Var x
158 subeverywhere trans (Lit x) = return $ Lit x
159 subeverywhere trans (Type x) = return $ Type x
161 subeverywhere trans (Cast expr ty) = do
163 return $ Cast expr' ty
165 subeverywhere trans expr = error $ "\nNormalizeTools.subeverywhere: Unsupported expression: " ++ show expr
167 -- Apply the given transformation to all expressions, except for direct
168 -- arguments of an application
169 notappargs :: (String, Transform) -> Transform
170 notappargs trans = applyboth (subnotappargs trans) trans
172 -- Apply the given transformation to all (direct and indirect) subexpressions
173 -- (but not the expression itself), except for direct arguments of an
175 subnotappargs :: (String, Transform) -> Transform
176 subnotappargs trans (App a b) = do
177 a' <- subnotappargs trans a
178 b' <- subnotappargs trans b
181 -- Let subeverywhere handle all other expressions
182 subnotappargs trans expr = subeverywhere (notappargs trans) expr
184 -- Runs each of the transforms repeatedly inside the State monad.
185 dotransforms :: [Transform] -> CoreExpr -> TransformSession CoreExpr
186 dotransforms transs expr = do
187 (expr', changed) <- Writer.runWriterT $ Monad.foldM (flip ($)) expr transs
188 if Monoid.getAny changed then dotransforms transs expr' else return expr'
190 -- Inline all let bindings that satisfy the given condition
191 inlinebind :: ((CoreBndr, CoreExpr) -> TransformMonad Bool) -> Transform
192 inlinebind condition expr@(Let (Rec binds) res) = do
193 -- Find all bindings that adhere to the condition
194 res_eithers <- mapM docond binds
195 case Either.partitionEithers res_eithers of
196 -- No replaces? No change
197 ([], _) -> return expr
198 (replace, others) -> do
199 -- Substitute the to be replaced binders with their expression
200 let newexpr = substitute replace (Let (Rec others) res)
203 docond :: (CoreBndr, CoreExpr) -> TransformMonad (Either (CoreBndr, CoreExpr) (CoreBndr, CoreExpr))
206 return $ case res of True -> Left b; False -> Right b
208 -- Leave all other expressions unchanged
209 inlinebind _ expr = return expr
211 -- Sets the changed flag in the TransformMonad, to signify that some
212 -- transform has changed the result
213 setChanged :: TransformMonad ()
214 setChanged = Writer.tell (Monoid.Any True)
216 -- Sets the changed flag and returns the given value.
217 change :: a -> TransformMonad a
222 -- Create a new Unique
223 mkUnique :: TransformMonad Unique.Unique
224 mkUnique = Trans.lift $ do
225 us <- getA tsUniqSupply
226 let (us', us'') = UniqSupply.splitUniqSupply us
227 putA tsUniqSupply us'
228 return $ UniqSupply.uniqFromSupply us''
230 -- Replace each of the binders given with the coresponding expressions in the
232 substitute :: [(CoreBndr, CoreExpr)] -> CoreExpr -> CoreExpr
233 substitute [] expr = expr
234 -- Apply one substitution on the expression, but also on any remaining
235 -- substitutions. This seems to be the only way to handle substitutions like
236 -- [(b, c), (a, b)]. This means we reuse a substitution, which is not allowed
237 -- according to CoreSubst documentation (but it doesn't seem to be a problem).
238 -- TODO: Find out how this works, exactly.
239 substitute ((b, e):subss) expr = substitute subss' expr'
242 subs = (CoreSubst.extendSubst CoreSubst.emptySubst b e)
243 -- Apply this substitution to the main expression
244 expr' = CoreSubst.substExpr subs expr
245 -- Apply this substitution on all the expressions in the remaining
247 subss' = map (Arrow.second (CoreSubst.substExpr subs)) subss
249 -- Run a given TransformSession. Used mostly to setup the right calls and
251 runTransformSession :: HscTypes.HscEnv -> UniqSupply.UniqSupply -> TransformSession a -> a
252 runTransformSession env uniqSupply session = State.evalState session emptyTransformState
254 emptyTypeState = TypeState Map.empty [] Map.empty Map.empty env
255 emptyTransformState = TransformState uniqSupply Map.empty VarSet.emptyVarSet emptyTypeState
257 -- Is the given expression representable at runtime, based on the type?
258 isRepr :: CoreSyn.CoreExpr -> TransformMonad Bool
259 isRepr (Type ty) = return False
260 isRepr expr = Trans.lift $ MonadState.lift tsType $ VHDLTools.isReprType (CoreUtils.exprType expr)
262 is_local_var :: CoreSyn.CoreExpr -> TransformSession Bool
263 is_local_var (CoreSyn.Var v) = do
264 bndrs <- getGlobalBinders
265 return $ not $ v `elem` bndrs
266 is_local_var _ = return False