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.Utils.Pretty
  39 import CLasH.VHDL.VHDLTypes
  40 import qualified CLasH.VHDL.VHDLTools as 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 = Trans.lift (mkInternalVar' str ty)
  48
  49 mkInternalVar' :: String -> Type.Type -> TransformSession Var.Var
  50 mkInternalVar' str ty = do
  51   uniq <- mkUnique'
  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
  55
  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)
  62
  63 mkTypeVar' :: String -> Type.Kind -> TransformSession Var.Var
  64 mkTypeVar' str kind = do
  65   uniq <- mkUnique'
  66   let occname = OccName.mkVarOcc (str ++ show uniq)
  67   let name = Name.mkInternalName uniq occname SrcLoc.noSrcSpan
  68   return $ Var.mkTyVar name kind
  69
  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)
  75
  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)
  79
  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)
  85
  86 cloneVar :: Var.Var -> TransformMonad Var.Var
  87 cloneVar v = do
  88   uniq <- mkUnique
  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
  92
  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
  95 -- this function.
  96 mkFunction :: CoreBndr -> CoreExpr -> TransformMonad CoreBndr
  97 mkFunction bndr body = do
  98   let ty = CoreUtils.exprType body
  99   id <- cloneVar bndr
 100   let newid = Var.setVarType id ty
 101   Trans.lift $ addGlobalBind newid body
 102   return newid
 103
 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
 108
 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
 113   -- Apply the first
 114   expr' <- first expr
 115   -- Apply the second
 116   (expr'', changed) <- Writer.listen $ second expr'
 117   if Monoid.getAny $
 118 --        trace ("Trying to apply transform " ++ name ++ " to:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n") $
 119         changed
 120     then
 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) $
 124         expr''
 125     else
 126 --      trace ("No changes") $
 127       return expr''
 128
 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
 133   a' <- trans a
 134   b' <- trans b
 135   return $ App a' b'
 136
 137 subeverywhere trans (Let (NonRec b bexpr) expr) = do
 138   bexpr' <- trans bexpr
 139   expr' <- trans expr
 140   return $ Let (NonRec b bexpr') expr'
 141
 142 subeverywhere trans (Let (Rec binds) expr) = do
 143   expr' <- trans expr
 144   binds' <- mapM transbind binds
 145   return $ Let (Rec binds') expr'
 146   where
 147     transbind :: (CoreBndr, CoreExpr) -> TransformMonad (CoreBndr, CoreExpr)
 148     transbind (b, e) = do
 149       e' <- trans e
 150       return (b, e')
 151
 152 subeverywhere trans (Lam x expr) = do
 153   expr' <- trans expr
 154   return $ Lam x expr'
 155
 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'
 160   where
 161     transalt :: CoreAlt -> TransformMonad CoreAlt
 162     transalt (con, binders, expr) = do
 163       expr' <- trans expr
 164       return (con, binders, expr')
 165
 166 subeverywhere trans (Var x) = return $ Var x
 167 subeverywhere trans (Lit x) = return $ Lit x
 168 subeverywhere trans (Type x) = return $ Type x
 169
 170 subeverywhere trans (Cast expr ty) = do
 171   expr' <- trans expr
 172   return $ Cast expr' ty
 173
 174 subeverywhere trans expr = error $ "\nNormalizeTools.subeverywhere: Unsupported expression: " ++ show expr
 175
 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
 180
 181 -- Apply the given transformation to all (direct and indirect) subexpressions
 182 -- (but not the expression itself), except for direct arguments of an
 183 -- application
 184 subnotappargs :: (String, Transform) -> Transform
 185 subnotappargs trans (App a b) = do
 186   a' <- subnotappargs trans a
 187   b' <- subnotappargs trans b
 188   return $ App a' b'
 189
 190 -- Let subeverywhere handle all other expressions
 191 subnotappargs trans expr = subeverywhere (notappargs trans) expr
 192
 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'
 198
 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)
 210         change newexpr
 211   where
 212     docond :: (CoreBndr, CoreExpr) -> TransformMonad (Either (CoreBndr, CoreExpr) (CoreBndr, CoreExpr))
 213     docond b = do
 214       res <- condition b
 215       return $ case res of True -> Left b; False -> Right b
 216
 217 -- Leave all other expressions unchanged
 218 inlinebind _ expr = return expr
 219
 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)
 224
 225 -- Sets the changed flag and returns the given value.
 226 change :: a -> TransformMonad a
 227 change val = do
 228   setChanged
 229   return val
 230
 231 -- Create a new Unique
 232 mkUnique :: TransformMonad Unique.Unique
 233 mkUnique = Trans.lift $ mkUnique'
 234
 235 mkUnique' :: TransformSession Unique.Unique
 236 mkUnique' = do
 237   us <- getA tsUniqSupply
 238   let (us', us'') = UniqSupply.splitUniqSupply us
 239   putA tsUniqSupply us'
 240   return $ UniqSupply.uniqFromSupply us''
 241
 242 -- Replace each of the binders given with the coresponding expressions in the
 243 -- given expression.
 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'
 252   where
 253     -- Create the Subst
 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
 258     -- substitutions
 259     subss' = map (Arrow.second (CoreSubst.substExpr subs)) subss
 260
 261 -- Run a given TransformSession. Used mostly to setup the right calls and
 262 -- an initial state.
 263 runTransformSession :: HscTypes.HscEnv -> UniqSupply.UniqSupply -> TransformSession a -> a
 264 runTransformSession env uniqSupply session = State.evalState session emptyTransformState
 265   where
 266     emptyTypeState = TypeState Map.empty [] Map.empty Map.empty env
 267     emptyTransformState = TransformState uniqSupply Map.empty VarSet.emptyVarSet emptyTypeState
 268
 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)
 273
 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