cλash/CLasH/Normalize/NormalizeTools.hs

   1 {-# LANGUAGE PackageImports #-}
   2 --
   3 -- This module provides functions for program transformations.
   4 --
   5 module CLasH.Normalize.NormalizeTools where
   6 -- Standard modules
   7 import Debug.Trace
   8 import qualified List
   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
  17 import Data.Accessor
  18 import Data.Accessor.MonadState as MonadState
  19
  20 -- GHC API
  21 import CoreSyn
  22 import qualified UniqSupply
  23 import qualified Unique
  24 import qualified OccName
  25 import qualified Name
  26 import qualified Var
  27 import qualified SrcLoc
  28 import qualified Type
  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 )
  35
  36 -- Local imports
  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
  42
  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)
  49
  50 mkInternalVar' :: String -> Type.Type -> TranslatorSession Var.Var
  51 mkInternalVar' str ty = do
  52   uniq <- mkUnique'
  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
  56
  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)
  63
  64 mkTypeVar' :: String -> Type.Kind -> TranslatorSession Var.Var
  65 mkTypeVar' str kind = do
  66   uniq <- mkUnique'
  67   let occname = OccName.mkVarOcc (str ++ show uniq)
  68   let name = Name.mkInternalName uniq occname SrcLoc.noSrcSpan
  69   return $ Var.mkTyVar name kind
  70
  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)
  76
  77 mkBinderFor' :: CoreExpr -> String -> TranslatorSession Var.Var
  78 mkBinderFor' (Type ty) string = mkTypeVar' string (Type.typeKind ty)
  79 mkBinderFor' expr string = mkInternalVar' string (CoreUtils.exprType expr)
  80
  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)
  86
  87 cloneVar :: Var.Var -> TransformMonad Var.Var
  88 cloneVar v = do
  89   uniq <- mkUnique
  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
  93
  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
  96 -- this function.
  97 mkFunction :: CoreBndr -> CoreExpr -> TransformMonad CoreBndr
  98 mkFunction bndr body = do
  99   let ty = CoreUtils.exprType body
 100   id <- cloneVar bndr
 101   let newid = Var.setVarType id ty
 102   Trans.lift $ addGlobalBind newid body
 103   return newid
 104
 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
 109
 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
 114   -- Apply the first
 115   expr' <- first expr
 116   -- Apply the second
 117   (expr'', changed) <- Writer.listen $ second expr'
 118   if Monoid.getAny $
 119 --        trace ("Trying to apply transform " ++ name ++ " to:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n") $
 120         changed
 121     then
 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) $
 125         expr''
 126     else
 127 --      trace ("No changes") $
 128       return expr''
 129
 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
 134   a' <- trans a
 135   b' <- trans b
 136   return $ App a' b'
 137
 138 subeverywhere trans (Let (NonRec b bexpr) expr) = do
 139   bexpr' <- trans bexpr
 140   expr' <- trans expr
 141   return $ Let (NonRec b bexpr') expr'
 142
 143 subeverywhere trans (Let (Rec binds) expr) = do
 144   expr' <- trans expr
 145   binds' <- mapM transbind binds
 146   return $ Let (Rec binds') expr'
 147   where
 148     transbind :: (CoreBndr, CoreExpr) -> TransformMonad (CoreBndr, CoreExpr)
 149     transbind (b, e) = do
 150       e' <- trans e
 151       return (b, e')
 152
 153 subeverywhere trans (Lam x expr) = do
 154   expr' <- trans expr
 155   return $ Lam x expr'
 156
 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'
 161   where
 162     transalt :: CoreAlt -> TransformMonad CoreAlt
 163     transalt (con, binders, expr) = do
 164       expr' <- trans expr
 165       return (con, binders, expr')
 166
 167 subeverywhere trans (Var x) = return $ Var x
 168 subeverywhere trans (Lit x) = return $ Lit x
 169 subeverywhere trans (Type x) = return $ Type x
 170
 171 subeverywhere trans (Cast expr ty) = do
 172   expr' <- trans expr
 173   return $ Cast expr' ty
 174
 175 subeverywhere trans expr = error $ "\nNormalizeTools.subeverywhere: Unsupported expression: " ++ show expr
 176
 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
 181
 182 -- Apply the given transformation to all (direct and indirect) subexpressions
 183 -- (but not the expression itself), except for direct arguments of an
 184 -- application
 185 subnotappargs :: (String, Transform) -> Transform
 186 subnotappargs trans (App a b) = do
 187   a' <- subnotappargs trans a
 188   b' <- subnotappargs trans b
 189   return $ App a' b'
 190
 191 -- Let subeverywhere handle all other expressions
 192 subnotappargs trans expr = subeverywhere (notappargs trans) expr
 193
 194 -- Runs each of the transforms repeatedly inside the State monad.
 195 dotransforms :: [Transform] -> CoreExpr -> TranslatorSession 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'
 199
 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)
 211         change newexpr
 212   where
 213     docond :: (CoreBndr, CoreExpr) -> TransformMonad (Either (CoreBndr, CoreExpr) (CoreBndr, CoreExpr))
 214     docond b = do
 215       res <- condition b
 216       return $ case res of True -> Left b; False -> Right b
 217
 218 -- Leave all other expressions unchanged
 219 inlinebind _ expr = return expr
 220
 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)
 225
 226 -- Sets the changed flag and returns the given value.
 227 change :: a -> TransformMonad a
 228 change val = do
 229   setChanged
 230   return val
 231
 232 -- Create a new Unique
 233 mkUnique :: TransformMonad Unique.Unique
 234 mkUnique = Trans.lift $ mkUnique'
 235
 236 mkUnique' :: TranslatorSession Unique.Unique
 237 mkUnique' = do
 238   us <- getA tsUniqSupply
 239   let (us', us'') = UniqSupply.splitUniqSupply us
 240   putA tsUniqSupply us'
 241   return $ UniqSupply.uniqFromSupply us''
 242
 243 -- Replace each of the binders given with the coresponding expressions in the
 244 -- given expression.
 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'
 253   where
 254     -- Create the Subst
 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
 259     -- substitutions
 260     subss' = map (Arrow.second (CoreSubst.substExpr subs)) subss
 261
 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)
 266
 267 is_local_var :: CoreSyn.CoreExpr -> TranslatorSession Bool
 268 is_local_var (CoreSyn.Var v) = do
 269   bndrs <- getGlobalBinders
 270   return $ not $ v `elem` bndrs
 271 is_local_var _ = return False