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