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 cloneVar :: Var.Var -> TransformMonad Var.Var
  47 cloneVar v = do
  48   uniq <- mkUnique
  49   -- Swap out the unique, and reset the IdInfo (I'm not 100% sure what it
  50   -- contains, but vannillaIdInfo is always correct, since it means "no info").
  51   return $ Var.lazySetVarIdInfo (Var.setVarUnique v uniq) IdInfo.vanillaIdInfo
  52
  53 -- Apply the given transformation to all expressions in the given expression,
  54 -- including the expression itself.
  55 everywhere :: (String, Transform) -> Transform
  56 everywhere trans = applyboth (subeverywhere (everywhere trans)) trans
  57
  58 -- Apply the first transformation, followed by the second transformation, and
  59 -- keep applying both for as long as expression still changes.
  60 applyboth :: Transform -> (String, Transform) -> Transform
  61 applyboth first (name, second) expr  = do
  62   -- Apply the first
  63   expr' <- first expr
  64   -- Apply the second
  65   (expr'', changed) <- Writer.listen $ second expr'
  66   if Monoid.getAny changed
  67     then
  68       trace ("Transform " ++ name ++ " changed from:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n" ++ "\nTo:\n" ++ showSDoc (nest 4 $ ppr expr'') ++ "\n" ++ "Type: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr'') ++ "\n" ) $
  69       applyboth first (name, second) expr''
  70     else
  71       return expr''
  72
  73 -- Apply the given transformation to all direct subexpressions (only), not the
  74 -- expression itself.
  75 subeverywhere :: Transform -> Transform
  76 subeverywhere trans (App a b) = do
  77   a' <- trans a
  78   b' <- trans b
  79   return $ App a' b'
  80
  81 subeverywhere trans (Let (Rec binds) expr) = do
  82   expr' <- trans expr
  83   binds' <- mapM transbind binds
  84   return $ Let (Rec binds') expr'
  85   where
  86     transbind :: (CoreBndr, CoreExpr) -> TransformMonad (CoreBndr, CoreExpr)
  87     transbind (b, e) = do
  88       e' <- trans e
  89       return (b, e')
  90
  91 subeverywhere trans (Lam x expr) = do
  92   expr' <- trans expr
  93   return $ Lam x expr'
  94
  95 subeverywhere trans (Case scrut b t alts) = do
  96   scrut' <- trans scrut
  97   alts' <- mapM transalt alts
  98   return $ Case scrut' b t alts'
  99   where
 100     transalt :: CoreAlt -> TransformMonad CoreAlt
 101     transalt (con, binders, expr) = do
 102       expr' <- trans expr
 103       return (con, binders, expr')
 104
 105
 106 subeverywhere trans expr = return expr
 107
 108 -- Apply the given transformation to all expressions, except for every first
 109 -- argument of an application.
 110 notapplied :: (String, Transform) -> Transform
 111 notapplied trans = applyboth (subnotapplied trans) trans
 112
 113 -- Apply the given transformation to all (direct and indirect) subexpressions
 114 -- (but not the expression itself), except for the first argument of an
 115 -- applicfirst argument of an application
 116 subnotapplied :: (String, Transform) -> Transform
 117 subnotapplied trans (App a b) = do
 118   a' <- subnotapplied trans a
 119   b' <- notapplied trans b
 120   return $ App a' b'
 121
 122 -- Let subeverywhere handle all other expressions
 123 subnotapplied trans expr = subeverywhere (notapplied trans) expr
 124
 125 -- Runs each of the transforms repeatedly inside the State monad.
 126 dotransforms :: [Transform] -> CoreExpr -> TransformSession CoreExpr
 127 dotransforms transs expr = do
 128   (expr', changed) <- Writer.runWriterT $ Monad.foldM (flip ($)) expr transs
 129   if Monoid.getAny changed then dotransforms transs expr' else return expr'
 130
 131 -- Inline all let bindings that satisfy the given condition
 132 inlinebind :: ((CoreBndr, CoreExpr) -> Bool) -> Transform
 133 inlinebind condition (Let (Rec binds) expr) | not $ null replace =
 134     change newexpr
 135   where
 136     -- Find all simple bindings
 137     (replace, others) = List.partition condition binds
 138     -- Substitute the to be replaced binders with their expression
 139     newexpr = substitute replace (Let (Rec others) expr)
 140 -- Leave all other expressions unchanged
 141 inlinebind _ expr = return expr
 142
 143 -- Sets the changed flag in the TransformMonad, to signify that some
 144 -- transform has changed the result
 145 setChanged :: TransformMonad ()
 146 setChanged = Writer.tell (Monoid.Any True)
 147
 148 -- Sets the changed flag and returns the given value.
 149 change :: a -> TransformMonad a
 150 change val = do
 151   setChanged
 152   return val
 153
 154 -- Create a new Unique
 155 mkUnique :: TransformMonad Unique.Unique
 156 mkUnique = Trans.lift $ do
 157     us <- getA tsUniqSupply
 158     let (us', us'') = UniqSupply.splitUniqSupply us
 159     putA tsUniqSupply us'
 160     return $ UniqSupply.uniqFromSupply us''
 161
 162 -- Replace each of the binders given with the coresponding expressions in the
 163 -- given expression.
 164 substitute :: [(CoreBndr, CoreExpr)] -> CoreExpr -> CoreExpr
 165 substitute replace expr = CoreSubst.substExpr subs expr
 166     where subs = foldl (\s (b, e) -> CoreSubst.extendIdSubst s b e) CoreSubst.emptySubst replace
 167
 168 -- Run a given TransformSession. Used mostly to setup the right calls and
 169 -- an initial state.
 170 runTransformSession :: UniqSupply.UniqSupply -> TransformSession a -> a
 171 runTransformSession uniqSupply session = State.evalState session initState
 172                        where initState = TransformState uniqSupply Map.empty VarSet.emptyVarSet