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 Control.Arrow as Arrow
  11 import qualified Control.Monad as Monad
  12 import qualified Control.Monad.Trans.State as State
  13 import qualified Control.Monad.Trans.Writer as Writer
  14 import qualified "transformers" Control.Monad.Trans as Trans
  15 import qualified Data.Map as Map
  16 import Data.Accessor
  17
  18 -- GHC API
  19 import CoreSyn
  20 import qualified UniqSupply
  21 import qualified Unique
  22 import qualified OccName
  23 import qualified Name
  24 import qualified Var
  25 import qualified SrcLoc
  26 import qualified Type
  27 import qualified IdInfo
  28 import qualified CoreUtils
  29 import qualified CoreSubst
  30 import qualified VarSet
  31 import Outputable ( showSDoc, ppr, nest )
  32
  33 -- Local imports
  34 import NormalizeTypes
  35
  36 -- Create a new internal var with the given name and type. A Unique is
  37 -- appended to the given name, to ensure uniqueness (not strictly neccesary,
  38 -- since the Unique is also stored in the name, but this ensures variable
  39 -- names are unique in the output).
  40 mkInternalVar :: String -> Type.Type -> TransformMonad Var.Var
  41 mkInternalVar str ty = do
  42   uniq <- mkUnique
  43   let occname = OccName.mkVarOcc (str ++ show uniq)
  44   let name = Name.mkInternalName uniq occname SrcLoc.noSrcSpan
  45   return $ Var.mkLocalIdVar name ty IdInfo.vanillaIdInfo
  46
  47 -- Create a new type variable with the given name and kind. A Unique is
  48 -- appended to the given name, to ensure uniqueness (not strictly neccesary,
  49 -- since the Unique is also stored in the name, but this ensures variable
  50 -- names are unique in the output).
  51 mkTypeVar :: String -> Type.Kind -> TransformMonad Var.Var
  52 mkTypeVar str kind = do
  53   uniq <- mkUnique
  54   let occname = OccName.mkVarOcc (str ++ show uniq)
  55   let name = Name.mkInternalName uniq occname SrcLoc.noSrcSpan
  56   return $ Var.mkTyVar name kind
  57
  58 -- Creates a binder for the given expression with the given name. This
  59 -- works for both value and type level expressions, so it can return a Var or
  60 -- TyVar (which is just an alias for Var).
  61 mkBinderFor :: CoreExpr -> String -> TransformMonad Var.Var
  62 mkBinderFor (Type ty) string = mkTypeVar string (Type.typeKind ty)
  63 mkBinderFor expr string = mkInternalVar string (CoreUtils.exprType expr)
  64
  65 -- Creates a reference to the given variable. This works for both a normal
  66 -- variable as well as a type variable
  67 mkReferenceTo :: Var.Var -> CoreExpr
  68 mkReferenceTo var | Var.isTyVar var = (Type $ Type.mkTyVarTy var)
  69                   | otherwise       = (Var var)
  70
  71 cloneVar :: Var.Var -> TransformMonad Var.Var
  72 cloneVar v = do
  73   uniq <- mkUnique
  74   -- Swap out the unique, and reset the IdInfo (I'm not 100% sure what it
  75   -- contains, but vannillaIdInfo is always correct, since it means "no info").
  76   return $ Var.lazySetVarIdInfo (Var.setVarUnique v uniq) IdInfo.vanillaIdInfo
  77
  78 -- Creates a new function with the same name as the given binder (but with a
  79 -- new unique) and with the given function body. Returns the new binder for
  80 -- this function.
  81 mkFunction :: CoreBndr -> CoreExpr -> TransformMonad CoreBndr
  82 mkFunction bndr body = do
  83   let ty = CoreUtils.exprType body
  84   id <- cloneVar bndr
  85   let newid = Var.setVarType id ty
  86   Trans.lift $ addGlobalBind newid body
  87   return newid
  88
  89 -- Apply the given transformation to all expressions in the given expression,
  90 -- including the expression itself.
  91 everywhere :: (String, Transform) -> Transform
  92 everywhere trans = applyboth (subeverywhere (everywhere trans)) trans
  93
  94 -- Apply the first transformation, followed by the second transformation, and
  95 -- keep applying both for as long as expression still changes.
  96 applyboth :: Transform -> (String, Transform) -> Transform
  97 applyboth first (name, second) expr  = do
  98   -- Apply the first
  99   expr' <- first expr
 100   -- Apply the second
 101   (expr'', changed) <- Writer.listen $ second expr'
 102   if Monoid.getAny $
 103   --      trace ("Trying to apply transform " ++ name ++ " to:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n") $
 104         changed
 105     then
 106 --      trace ("Applying transform " ++ name ++ " to:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n") $
 107  --     trace ("Result of applying " ++ name ++ ":\n" ++ showSDoc (nest 4 $ ppr expr'') ++ "\n" ++ "Type: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr'') ++ "\n" ) $
 108       applyboth first (name, second) $
 109         expr''
 110     else
 111     --  trace ("No changes") $
 112       return expr''
 113
 114 -- Apply the given transformation to all direct subexpressions (only), not the
 115 -- expression itself.
 116 subeverywhere :: Transform -> Transform
 117 subeverywhere trans (App a b) = do
 118   a' <- trans a
 119   b' <- trans b
 120   return $ App a' b'
 121
 122 subeverywhere trans (Let (NonRec b bexpr) expr) = do
 123   bexpr' <- trans bexpr
 124   expr' <- trans expr
 125   return $ Let (NonRec b bexpr') expr'
 126
 127 subeverywhere trans (Let (Rec binds) expr) = do
 128   expr' <- trans expr
 129   binds' <- mapM transbind binds
 130   return $ Let (Rec binds') expr'
 131   where
 132     transbind :: (CoreBndr, CoreExpr) -> TransformMonad (CoreBndr, CoreExpr)
 133     transbind (b, e) = do
 134       e' <- trans e
 135       return (b, e')
 136
 137 subeverywhere trans (Lam x expr) = do
 138   expr' <- trans expr
 139   return $ Lam x expr'
 140
 141 subeverywhere trans (Case scrut b t alts) = do
 142   scrut' <- trans scrut
 143   alts' <- mapM transalt alts
 144   return $ Case scrut' b t alts'
 145   where
 146     transalt :: CoreAlt -> TransformMonad CoreAlt
 147     transalt (con, binders, expr) = do
 148       expr' <- trans expr
 149       return (con, binders, expr')
 150
 151 subeverywhere trans (Var x) = return $ Var x
 152 subeverywhere trans (Lit x) = return $ Lit x
 153 subeverywhere trans (Type x) = return $ Type x
 154
 155 subeverywhere trans (Cast expr ty) = do
 156   expr' <- trans expr
 157   return $ Cast expr' ty
 158
 159 subeverywhere trans expr = error $ "\nNormalizeTools.subeverywhere: Unsupported expression: " ++ show expr
 160
 161 -- Apply the given transformation to all expressions, except for direct
 162 -- arguments of an application
 163 notappargs :: (String, Transform) -> Transform
 164 notappargs trans = applyboth (subnotappargs trans) trans
 165
 166 -- Apply the given transformation to all (direct and indirect) subexpressions
 167 -- (but not the expression itself), except for direct arguments of an
 168 -- application
 169 subnotappargs :: (String, Transform) -> Transform
 170 subnotappargs trans (App a b) = do
 171   a' <- subnotappargs trans a
 172   b' <- subnotappargs trans b
 173   return $ App a' b'
 174
 175 -- Let subeverywhere handle all other expressions
 176 subnotappargs trans expr = subeverywhere (notappargs trans) expr
 177
 178 -- Runs each of the transforms repeatedly inside the State monad.
 179 dotransforms :: [Transform] -> CoreExpr -> TransformSession CoreExpr
 180 dotransforms transs expr = do
 181   (expr', changed) <- Writer.runWriterT $ Monad.foldM (flip ($)) expr transs
 182   if Monoid.getAny changed then dotransforms transs expr' else return expr'
 183
 184 -- Inline all let bindings that satisfy the given condition
 185 inlinebind :: ((CoreBndr, CoreExpr) -> Bool) -> Transform
 186 inlinebind condition (Let (Rec binds) expr) | not $ null replace =
 187     change newexpr
 188   where
 189     -- Find all simple bindings
 190     (replace, others) = List.partition condition binds
 191     -- Substitute the to be replaced binders with their expression
 192     newexpr = substitute replace (Let (Rec others) expr)
 193 -- Leave all other expressions unchanged
 194 inlinebind _ expr = return expr
 195
 196 -- Sets the changed flag in the TransformMonad, to signify that some
 197 -- transform has changed the result
 198 setChanged :: TransformMonad ()
 199 setChanged = Writer.tell (Monoid.Any True)
 200
 201 -- Sets the changed flag and returns the given value.
 202 change :: a -> TransformMonad a
 203 change val = do
 204   setChanged
 205   return val
 206
 207 -- Create a new Unique
 208 mkUnique :: TransformMonad Unique.Unique
 209 mkUnique = Trans.lift $ do
 210     us <- getA tsUniqSupply
 211     let (us', us'') = UniqSupply.splitUniqSupply us
 212     putA tsUniqSupply us'
 213     return $ UniqSupply.uniqFromSupply us''
 214
 215 -- Replace each of the binders given with the coresponding expressions in the
 216 -- given expression.
 217 substitute :: [(CoreBndr, CoreExpr)] -> CoreExpr -> CoreExpr
 218 substitute [] expr = expr
 219 -- Apply one substitution on the expression, but also on any remaining
 220 -- substitutions. This seems to be the only way to handle substitutions like
 221 -- [(b, c), (a, b)]. This means we reuse a substitution, which is not allowed
 222 -- according to CoreSubst documentation (but it doesn't seem to be a problem).
 223 -- TODO: Find out how this works, exactly.
 224 substitute ((b, e):subss) expr = substitute subss' expr'
 225   where
 226     -- Create the Subst
 227     subs = (CoreSubst.extendSubst CoreSubst.emptySubst b e)
 228     -- Apply this substitution to the main expression
 229     expr' = CoreSubst.substExpr subs expr
 230     -- Apply this substitution on all the expressions in the remaining
 231     -- substitutions
 232     subss' = map (Arrow.second (CoreSubst.substExpr subs)) subss
 233
 234 -- Run a given TransformSession. Used mostly to setup the right calls and
 235 -- an initial state.
 236 runTransformSession :: UniqSupply.UniqSupply -> TransformSession a -> a
 237 runTransformSession uniqSupply session = State.evalState session (emptyTransformState uniqSupply)