utils/ext-core/Language/Core/ElimDeadCode.hs

   1 {-
   2    This module eliminates unused top-level bindings, under the
   3    assumption that all top-level bindings with qualified names
   4    should be retained.
   5 -}
   6 module Language.Core.ElimDeadCode(elimDeadCode) where
   7
   8 import Language.Core.Core
   9 import Language.Core.Printer()
  10 import Language.Core.CoreUtils
  11 import Language.Core.Utils
  12
  13 import Control.Monad.Reader
  14 import Data.Generics
  15 import Data.List
  16 import Data.Maybe
  17 import qualified Data.Map as M
  18 import qualified Data.Set as S
  19
  20 elimDeadCode :: Bool -> Module -> Module
  21 -- exports = true <=> it's assumed we want to keep exported functions;
  22 -- otherwise, we assume the module is "closed" and eliminate everything
  23 -- not reachable from Main
  24 elimDeadCode exports (Module mn tdefs vdefgs) = runReader (do
  25   (usedVars, usedDcons, usedTcons) <- findUsed emptySet
  26      (mkStartSet exports mn vdefgs)
  27   let isUsed (Vdef (v,_,_)) = v `S.member` usedVars
  28   let newVdefgs = filterVdefgs isUsed vdefgs
  29   let newTdefs  = filter (tdefIsUsed usedTcons usedDcons) tdefs in
  30     return $ Module mn newTdefs newVdefgs) ((mkVarEnv vdefgs), mkTyEnv tdefs)
  31
  32 tdefIsUsed :: S.Set (Qual Tcon) -> S.Set (Qual Dcon) -> Tdef -> Bool
  33 tdefIsUsed tcs dcs (Data qtc _ cdefs) =
  34   (qtc `S.member` tcs || any (\ (Constr qdc _ _) -> qdc `S.member` dcs) cdefs)
  35 tdefIsUsed tcs _ (Newtype qtc qtc_co _ _) =
  36   qtc `S.member` tcs || qtc_co `S.member` tcs
  37
  38 mkVarEnv :: [Vdefg] -> M.Map (Qual Var) Exp
  39 mkVarEnv vgs =
  40   let vdefs = flattenBinds vgs in
  41     M.fromList [(v, e) | (Vdef (v, _, e)) <- vdefs]
  42
  43 -- if there is a Newtype qtc qtc_co ty,
  44 -- generate: qtc |-> ty and qtc_co |-> ty
  45 -- roughly the same for rhs's of Data decls
  46 mkTyEnv :: [Tdef] -> M.Map (Qual Tcon) [Ty]
  47 mkTyEnv tdefs =
  48   M.fromList ([(qtc, [ty]) | (Newtype qtc _ _ ty) <- tdefs]
  49          ++   [(qtc, [ty]) | (Newtype _ qtc _ ty) <- tdefs]
  50          ++   concatMap (\ td -> case td of
  51                                    Data qtc _ cdefs -> [(qtc, concatMap
  52                                      (\ (Constr _ _ ts) -> ts) cdefs)]
  53                                    _ -> []) tdefs)
  54
  55 findUsed :: DeadSet -> DeadSet -> DeadM DeadSet
  56 findUsed _old@(oldVars,oldDcs,oldTcs) _new@(newVars,newDcs,newTcs) = do
  57   let (todoVars, todoTcs) = ((S.\\) newVars oldVars, (S.\\) newTcs oldTcs)
  58   let nextOld = (oldVars `S.union` todoVars, oldDcs `S.union` newDcs,
  59                  oldTcs `S.union` todoTcs)
  60   nextStuff <- getVarsAndConsIn (todoVars, todoTcs)
  61   if (S.null todoVars && S.null todoTcs)
  62       then return nextOld
  63       else findUsed nextOld nextStuff
  64
  65 getVarsAndConsIn :: (S.Set (Qual Var), S.Set (Qual Tcon)) -> DeadM DeadSet
  66 getVarsAndConsIn (vs, tcs) = do
  67   vs1 <- mapM varsAndConsInOne (S.toList vs)
  68   ts1 <- mapM varsAndConsInOne' (S.toList tcs)
  69   let (vs'::[S.Set (Qual Var)], dcs'::[S.Set (Qual Dcon)],
  70         ts'::[S.Set (Qual Tcon)]) = unzip3 (vs1 ++ ts1)
  71   return (foldl' S.union S.empty vs', foldl' S.union S.empty dcs',
  72           foldl' S.union S.empty ts')
  73
  74 varsAndConsInOne :: Qual Var -> DeadM DeadSet
  75 varsAndConsInOne vr = do
  76   def <- findDefn vr
  77   return $ maybe emptySet
  78     (noNames emptySet unionThree (mkQ emptySet usedNamesAll)) def
  79
  80 varsAndConsInOne' :: Qual Tcon -> DeadM DeadSet
  81 varsAndConsInOne' tc = do
  82   ty <- findRepTy tc
  83   return $ maybe emptySet
  84     (noNames emptySet unionThree
  85               (mkQ emptySet usedStuffTys)) ty
  86
  87 emptySet :: DeadSet
  88 emptySet = (S.empty, S.empty, S.empty)
  89 mkStartSet :: Bool -> AnMname -> [Vdefg] -> DeadSet
  90 -- Initially, we assume the definitions of any exported functions are not
  91 -- dead, and work backwards from there.
  92 mkStartSet exports mn vds =
  93   (S.fromList (filter ((== Just mn) . getModule) (if exports then exportedNames vds else [mainVar])),
  94    S.empty, S.empty)
  95
  96 exportedNames :: [Vdefg] -> [Qual Var]
  97 exportedNames vdefgs =
  98   let vds = flattenBinds vdefgs in
  99     filter isQual (ns vds)
 100       where isQual    = isJust . fst
 101             ns = map (\ (Vdef (n,_,_)) -> n)
 102
 103 type DeadSet = (S.Set (Qual Var), S.Set (Qual Dcon), S.Set (Qual Tcon))
 104 type DeadM = Reader (M.Map (Qual Var) Exp, M.Map (Qual Tcon) [Ty])
 105
 106 findDefn :: Qual Var -> DeadM (Maybe Exp)
 107 findDefn vr = asks ((M.lookup vr) . fst)
 108 findRepTy :: Qual Tcon -> DeadM (Maybe [Ty])
 109 findRepTy tc = asks ((M.lookup tc) . snd)
 110
 111 unionThree :: DeadSet -> DeadSet -> DeadSet
 112 unionThree (a,b,c) (d,e,f) = (a `S.union` d, b `S.union` e, c `S.union` f)
 113
 114 usedNamesAll :: Exp -> DeadSet
 115 usedNamesAll = (noNames emptySet unionThree
 116   ((mkQ emptySet usedStuff) `extQ` usedStuffTys `extQ` usedStuffAlts))
 117
 118 usedStuff :: Exp -> DeadSet
 119 usedStuff (Var qv)  = (S.singleton qv, S.empty, S.empty)
 120 usedStuff (Dcon dc) = (S.empty, S.singleton dc, S.empty)
 121 usedStuff _         = emptySet
 122
 123 usedStuffAlts :: Alt -> DeadSet
 124 usedStuffAlts (Acon qdc _ _ _) = (S.empty, S.singleton qdc, S.empty)
 125 usedStuffAlts _ = emptySet
 126
 127 usedStuffTys :: Ty -> DeadSet
 128 usedStuffTys (Tcon qtc) = (S.empty, S.empty, S.singleton qtc)
 129 usedStuffTys _          = emptySet