cλash/CLasH/Normalize/NormalizeTools.hs

   1 {-# LANGUAGE PackageImports #-}
   2 --
   3 -- This module provides functions for program transformations.
   4 --
   5 module 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 NormalizeTypes
  38 import Pretty
  39 import VHDLTypes
  40 import qualified VHDLTools
  41
  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
  48   uniq <- mkUnique
  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
  52
  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
  59   uniq <- mkUnique
  60   let occname = OccName.mkVarOcc (str ++ show uniq)
  61   let name = Name.mkInternalName uniq occname SrcLoc.noSrcSpan
  62   return $ Var.mkTyVar name kind
  63
  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)
  70
  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)
  76
  77 cloneVar :: Var.Var -> TransformMonad Var.Var
  78 cloneVar v = do
  79   uniq <- mkUnique
  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
  83
  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
  86 -- this function.
  87 mkFunction :: CoreBndr -> CoreExpr -> TransformMonad CoreBndr
  88 mkFunction bndr body = do
  89   let ty = CoreUtils.exprType body
  90   id <- cloneVar bndr
  91   let newid = Var.setVarType id ty
  92   Trans.lift $ addGlobalBind newid body
  93   return newid
  94
  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
  99
 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
 104   -- Apply the first
 105   expr' <- first expr
 106   -- Apply the second
 107   (expr'', changed) <- Writer.listen $ second expr'
 108   if Monoid.getAny $
 109 --        trace ("Trying to apply transform " ++ name ++ " to:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n") $
 110         changed
 111     then
 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) $
 115         expr''
 116     else
 117 --      trace ("No changes") $
 118       return expr''
 119
 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
 124   a' <- trans a
 125   b' <- trans b
 126   return $ App a' b'
 127
 128 subeverywhere trans (Let (NonRec b bexpr) expr) = do
 129   bexpr' <- trans bexpr
 130   expr' <- trans expr
 131   return $ Let (NonRec b bexpr') expr'
 132
 133 subeverywhere trans (Let (Rec binds) expr) = do
 134   expr' <- trans expr
 135   binds' <- mapM transbind binds
 136   return $ Let (Rec binds') expr'
 137   where
 138     transbind :: (CoreBndr, CoreExpr) -> TransformMonad (CoreBndr, CoreExpr)
 139     transbind (b, e) = do
 140       e' <- trans e
 141       return (b, e')
 142
 143 subeverywhere trans (Lam x expr) = do
 144   expr' <- trans expr
 145   return $ Lam x expr'
 146
 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'
 151   where
 152     transalt :: CoreAlt -> TransformMonad CoreAlt
 153     transalt (con, binders, expr) = do
 154       expr' <- trans expr
 155       return (con, binders, expr')
 156
 157 subeverywhere trans (Var x) = return $ Var x
 158 subeverywhere trans (Lit x) = return $ Lit x
 159 subeverywhere trans (Type x) = return $ Type x
 160
 161 subeverywhere trans (Cast expr ty) = do
 162   expr' <- trans expr
 163   return $ Cast expr' ty
 164
 165 subeverywhere trans expr = error $ "\nNormalizeTools.subeverywhere: Unsupported expression: " ++ show expr
 166
 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
 171
 172 -- Apply the given transformation to all (direct and indirect) subexpressions
 173 -- (but not the expression itself), except for direct arguments of an
 174 -- application
 175 subnotappargs :: (String, Transform) -> Transform
 176 subnotappargs trans (App a b) = do
 177   a' <- subnotappargs trans a
 178   b' <- subnotappargs trans b
 179   return $ App a' b'
 180
 181 -- Let subeverywhere handle all other expressions
 182 subnotappargs trans expr = subeverywhere (notappargs trans) expr
 183
 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'
 189
 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)
 201         change newexpr
 202   where
 203     docond :: (CoreBndr, CoreExpr) -> TransformMonad (Either (CoreBndr, CoreExpr) (CoreBndr, CoreExpr))
 204     docond b = do
 205       res <- condition b
 206       return $ case res of True -> Left b; False -> Right b
 207
 208 -- Leave all other expressions unchanged
 209 inlinebind _ expr = return expr
 210
 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)
 215
 216 -- Sets the changed flag and returns the given value.
 217 change :: a -> TransformMonad a
 218 change val = do
 219   setChanged
 220   return val
 221
 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''
 229
 230 -- Replace each of the binders given with the coresponding expressions in the
 231 -- given expression.
 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'
 240   where
 241     -- Create the Subst
 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
 246     -- substitutions
 247     subss' = map (Arrow.second (CoreSubst.substExpr subs)) subss
 248
 249 -- Run a given TransformSession. Used mostly to setup the right calls and
 250 -- an initial state.
 251 runTransformSession :: HscTypes.HscEnv -> UniqSupply.UniqSupply -> TransformSession a -> a
 252 runTransformSession env uniqSupply session = State.evalState session emptyTransformState
 253   where
 254     emptyTypeState = TypeState Map.empty [] Map.empty Map.empty env
 255     emptyTransformState = TransformState uniqSupply Map.empty VarSet.emptyVarSet emptyTypeState
 256
 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)
 261
 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