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