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 Data.Accessor
  15
  16 -- GHC API
  17 import CoreSyn
  18 import qualified UniqSupply
  19 import qualified Unique
  20 import qualified OccName
  21 import qualified Name
  22 import qualified Var
  23 import qualified SrcLoc
  24 import qualified Type
  25 import qualified IdInfo
  26 import qualified CoreUtils
  27 import qualified CoreSubst
  28 import Outputable ( showSDoc, ppr, nest )
  29
  30 -- Local imports
  31 import NormalizeTypes
  32
  33 -- Create a new internal var with the given name and type. A Unique is
  34 -- appended to the given name, to ensure uniqueness (not strictly neccesary,
  35 -- since the Unique is also stored in the name, but this ensures variable
  36 -- names are unique in the output).
  37 mkInternalVar :: String -> Type.Type -> TransformMonad Var.Var
  38 mkInternalVar str ty = do
  39   uniq <- mkUnique
  40   let occname = OccName.mkVarOcc (str ++ show uniq)
  41   let name = Name.mkInternalName uniq occname SrcLoc.noSrcSpan
  42   return $ Var.mkLocalIdVar name ty IdInfo.vanillaIdInfo
  43
  44 -- Apply the given transformation to all expressions in the given expression,
  45 -- including the expression itself.
  46 everywhere :: (String, Transform) -> Transform
  47 everywhere trans = applyboth (subeverywhere (everywhere trans)) trans
  48
  49 -- Apply the first transformation, followed by the second transformation, and
  50 -- keep applying both for as long as expression still changes.
  51 applyboth :: Transform -> (String, Transform) -> Transform
  52 applyboth first (name, second) expr  = do
  53   -- Apply the first
  54   expr' <- first expr
  55   -- Apply the second
  56   (expr'', changed) <- Writer.listen $ second expr'
  57   if Monoid.getAny changed
  58     then
  59       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" ) $
  60       applyboth first (name, second) expr''
  61     else
  62       return expr''
  63
  64 -- Apply the given transformation to all direct subexpressions (only), not the
  65 -- expression itself.
  66 subeverywhere :: Transform -> Transform
  67 subeverywhere trans (App a b) = do
  68   a' <- trans a
  69   b' <- trans b
  70   return $ App a' b'
  71
  72 subeverywhere trans (Let (Rec binds) expr) = do
  73   expr' <- trans expr
  74   binds' <- mapM transbind binds
  75   return $ Let (Rec binds') expr'
  76   where
  77     transbind :: (CoreBndr, CoreExpr) -> TransformMonad (CoreBndr, CoreExpr)
  78     transbind (b, e) = do
  79       e' <- trans e
  80       return (b, e')
  81
  82 subeverywhere trans (Lam x expr) = do
  83   expr' <- trans expr
  84   return $ Lam x expr'
  85
  86 subeverywhere trans (Case scrut b t alts) = do
  87   scrut' <- trans scrut
  88   alts' <- mapM transalt alts
  89   return $ Case scrut' b t alts'
  90   where
  91     transalt :: CoreAlt -> TransformMonad CoreAlt
  92     transalt (con, binders, expr) = do
  93       expr' <- trans expr
  94       return (con, binders, expr')
  95
  96
  97 subeverywhere trans expr = return expr
  98
  99 -- Apply the given transformation to all expressions, except for every first
 100 -- argument of an application.
 101 notapplied :: (String, Transform) -> Transform
 102 notapplied trans = applyboth (subnotapplied trans) trans
 103
 104 -- Apply the given transformation to all (direct and indirect) subexpressions
 105 -- (but not the expression itself), except for the first argument of an
 106 -- applicfirst argument of an application
 107 subnotapplied :: (String, Transform) -> Transform
 108 subnotapplied trans (App a b) = do
 109   a' <- subnotapplied trans a
 110   b' <- notapplied trans b
 111   return $ App a' b'
 112
 113 -- Let subeverywhere handle all other expressions
 114 subnotapplied trans expr = subeverywhere (notapplied trans) expr
 115
 116 -- Run the given transforms over the given expression
 117 dotransforms :: [Transform] -> UniqSupply.UniqSupply -> CoreExpr -> CoreExpr
 118 dotransforms transs uniqSupply = (flip State.evalState initState) . (dotransforms' transs)
 119                        where initState = TransformState uniqSupply
 120
 121 -- Runs each of the transforms repeatedly inside the State monad.
 122 dotransforms' :: [Transform] -> CoreExpr -> State.State TransformState CoreExpr
 123 dotransforms' transs expr = do
 124   (expr', changed) <- Writer.runWriterT $ Monad.foldM (flip ($)) expr transs
 125   if Monoid.getAny changed then dotransforms' transs expr' else return expr'
 126
 127 -- Inline all let bindings that satisfy the given condition
 128 inlinebind :: ((CoreBndr, CoreExpr) -> Bool) -> Transform
 129 inlinebind condition (Let (Rec binds) expr) | not $ null replace =
 130     change newexpr
 131   where
 132     -- Find all simple bindings
 133     (replace, others) = List.partition condition binds
 134     -- Substitute the to be replaced binders with their expression
 135     newexpr = substitute replace (Let (Rec others) expr)
 136 -- Leave all other expressions unchanged
 137 inlinebind _ expr = return expr
 138
 139 -- Sets the changed flag in the TransformMonad, to signify that some
 140 -- transform has changed the result
 141 setChanged :: TransformMonad ()
 142 setChanged = Writer.tell (Monoid.Any True)
 143
 144 -- Sets the changed flag and returns the given value.
 145 change :: a -> TransformMonad a
 146 change val = do
 147   setChanged
 148   return val
 149
 150 -- Create a new Unique
 151 mkUnique :: TransformMonad Unique.Unique
 152 mkUnique = Trans.lift $ do
 153     us <- getA tsUniqSupply
 154     let (us', us'') = UniqSupply.splitUniqSupply us
 155     putA tsUniqSupply us'
 156     return $ UniqSupply.uniqFromSupply us''
 157
 158 -- Replace each of the binders given with the coresponding expressions in the
 159 -- given expression.
 160 substitute :: [(CoreBndr, CoreExpr)] -> CoreExpr -> CoreExpr
 161 substitute replace expr = CoreSubst.substExpr subs expr
 162     where subs = foldl (\s (b, e) -> CoreSubst.extendIdSubst s b e) CoreSubst.emptySubst replace