module NormalizeTools where
-- Standard modules
import Debug.Trace
+import qualified List
import qualified Data.Monoid as Monoid
+import qualified Control.Arrow as Arrow
import qualified Control.Monad as Monad
import qualified Control.Monad.Trans.State as State
import qualified Control.Monad.Trans.Writer as Writer
import qualified "transformers" Control.Monad.Trans as Trans
+import qualified Data.Map as Map
import Data.Accessor
-- GHC API
import qualified SrcLoc
import qualified Type
import qualified IdInfo
+import qualified CoreUtils
+import qualified CoreSubst
+import qualified VarSet
import Outputable ( showSDoc, ppr, nest )
-- Local imports
let name = Name.mkInternalName uniq occname SrcLoc.noSrcSpan
return $ Var.mkLocalIdVar name ty IdInfo.vanillaIdInfo
+-- Create a new type variable with the given name and kind. A Unique is
+-- appended to the given name, to ensure uniqueness (not strictly neccesary,
+-- since the Unique is also stored in the name, but this ensures variable
+-- names are unique in the output).
+mkTypeVar :: String -> Type.Kind -> TransformMonad Var.Var
+mkTypeVar str kind = do
+ uniq <- mkUnique
+ let occname = OccName.mkVarOcc (str ++ show uniq)
+ let name = Name.mkInternalName uniq occname SrcLoc.noSrcSpan
+ return $ Var.mkTyVar name kind
+
+-- Creates a binder for the given expression with the given name. This
+-- works for both value and type level expressions, so it can return a Var or
+-- TyVar (which is just an alias for Var).
+mkBinderFor :: CoreExpr -> String -> TransformMonad Var.Var
+mkBinderFor (Type ty) string = mkTypeVar string (Type.typeKind ty)
+mkBinderFor expr string = mkInternalVar string (CoreUtils.exprType expr)
+
+-- Creates a reference to the given variable. This works for both a normal
+-- variable as well as a type variable
+mkReferenceTo :: Var.Var -> CoreExpr
+mkReferenceTo var | Var.isTyVar var = (Type $ Type.mkTyVarTy var)
+ | otherwise = (Var var)
+
+cloneVar :: Var.Var -> TransformMonad Var.Var
+cloneVar v = do
+ uniq <- mkUnique
+ -- Swap out the unique, and reset the IdInfo (I'm not 100% sure what it
+ -- contains, but vannillaIdInfo is always correct, since it means "no info").
+ return $ Var.lazySetVarIdInfo (Var.setVarUnique v uniq) IdInfo.vanillaIdInfo
+
+-- Creates a new function with the same name as the given binder (but with a
+-- new unique) and with the given function body. Returns the new binder for
+-- this function.
+mkFunction :: CoreBndr -> CoreExpr -> TransformMonad CoreBndr
+mkFunction bndr body = do
+ let ty = CoreUtils.exprType body
+ id <- cloneVar bndr
+ let newid = Var.setVarType id ty
+ Trans.lift $ addGlobalBind newid body
+ return newid
+
-- Apply the given transformation to all expressions in the given expression,
-- including the expression itself.
everywhere :: (String, Transform) -> Transform
expr' <- first expr
-- Apply the second
(expr'', changed) <- Writer.listen $ second expr'
- if Monoid.getAny changed
+ if Monoid.getAny $
+ -- trace ("Trying to apply transform " ++ name ++ " to:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n") $
+ changed
then
- trace ("Transform " ++ name ++ " changed from:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nTo:\n" ++ showSDoc (nest 4 $ ppr expr'') ++ "\n") $
- applyboth first (name, second) expr''
+-- trace ("Applying transform " ++ name ++ " to:\n" ++ showSDoc (nest 4 $ ppr expr') ++ "\nType: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr') ++ "\n") $
+ -- trace ("Result of applying " ++ name ++ ":\n" ++ showSDoc (nest 4 $ ppr expr'') ++ "\n" ++ "Type: \n" ++ (showSDoc $ nest 4 $ ppr $ CoreUtils.exprType expr'') ++ "\n" ) $
+ applyboth first (name, second) $
+ expr''
else
+ -- trace ("No changes") $
return expr''
-- Apply the given transformation to all direct subexpressions (only), not the
-- Let subeverywhere handle all other expressions
subnotapplied trans expr = subeverywhere (notapplied trans) expr
--- Run the given transforms over the given expression
-dotransforms :: [Transform] -> UniqSupply.UniqSupply -> CoreExpr -> CoreExpr
-dotransforms transs uniqSupply = (flip State.evalState initState) . (dotransforms' transs)
- where initState = TransformState uniqSupply
-
-- Runs each of the transforms repeatedly inside the State monad.
-dotransforms' :: [Transform] -> CoreExpr -> State.State TransformState CoreExpr
-dotransforms' transs expr = do
+dotransforms :: [Transform] -> CoreExpr -> TransformSession CoreExpr
+dotransforms transs expr = do
(expr', changed) <- Writer.runWriterT $ Monad.foldM (flip ($)) expr transs
- if Monoid.getAny changed then dotransforms' transs expr' else return expr'
+ if Monoid.getAny changed then dotransforms transs expr' else return expr'
+
+-- Inline all let bindings that satisfy the given condition
+inlinebind :: ((CoreBndr, CoreExpr) -> Bool) -> Transform
+inlinebind condition (Let (Rec binds) expr) | not $ null replace =
+ change newexpr
+ where
+ -- Find all simple bindings
+ (replace, others) = List.partition condition binds
+ -- Substitute the to be replaced binders with their expression
+ newexpr = substitute replace (Let (Rec others) expr)
+-- Leave all other expressions unchanged
+inlinebind _ expr = return expr
-- Sets the changed flag in the TransformMonad, to signify that some
-- transform has changed the result
let (us', us'') = UniqSupply.splitUniqSupply us
putA tsUniqSupply us'
return $ UniqSupply.uniqFromSupply us''
+
+-- Replace each of the binders given with the coresponding expressions in the
+-- given expression.
+substitute :: [(CoreBndr, CoreExpr)] -> CoreExpr -> CoreExpr
+substitute [] expr = expr
+-- Apply one substitution on the expression, but also on any remaining
+-- substitutions. This seems to be the only way to handle substitutions like
+-- [(b, c), (a, b)]. This means we reuse a substitution, which is not allowed
+-- according to CoreSubst documentation (but it doesn't seem to be a problem).
+-- TODO: Find out how this works, exactly.
+substitute ((b, e):subss) expr = substitute subss' expr'
+ where
+ -- Create the Subst
+ subs = (CoreSubst.extendSubst CoreSubst.emptySubst b e)
+ -- Apply this substitution to the main expression
+ expr' = CoreSubst.substExpr subs expr
+ -- Apply this substitution on all the expressions in the remaining
+ -- substitutions
+ subss' = map (Arrow.second (CoreSubst.substExpr subs)) subss
+
+-- Run a given TransformSession. Used mostly to setup the right calls and
+-- an initial state.
+runTransformSession :: UniqSupply.UniqSupply -> TransformSession a -> a
+runTransformSession uniqSupply session = State.evalState session initState
+ where initState = TransformState uniqSupply Map.empty VarSet.emptyVarSet