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