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