ext-core library: Add dead code eliminator for Core

Added code for dead code elimination to the ext-core library. This can be used in concert with Language.Core.Merge to produce a single self-contained module without unnecessary cruft. 

3 files changed, 128 insertions, 2 deletions

diff --git a/utils/ext-core/Language/Core/ElimDeadCode.hs b/utils/ext-core/Language/Core/ElimDeadCode.hs
new file mode 100644
index 0000000000..e32568e5ef
--- /dev/null
+++ b/utils/ext-core/Language/Core/ElimDeadCode.hs
@@ -0,0 +1,126 @@
+{- 
+   This module eliminates unused top-level bindings, under the
+   assumption that all top-level bindings with qualified names
+   should be retained.
+-}
+module Language.Core.ElimDeadCode(elimDeadCode) where
+
+import Language.Core.Core
+import Language.Core.Printer()
+import Language.Core.Utils
+
+import Control.Monad.Reader
+import Data.Generics
+import Data.List
+import Data.Maybe
+import qualified Data.Map as M
+import qualified Data.Set as S
+
+elimDeadCode :: Module -> Module
+elimDeadCode (Module mn tdefs vdefgs) = runReader (do
+  (usedVars, usedDcons, usedTcons) <- findUsed emptySet 
+     (mkStartSet mn vdefgs) 
+  let isUsed (Vdef (v,_,_)) = v `S.member` usedVars
+  let newVdefgs = [Rec $ filter isUsed (flattenBinds vdefgs)]
+  let newTdefs  = filter (tdefIsUsed usedTcons usedDcons) tdefs in
+    return $ Module mn newTdefs newVdefgs) ((mkVarEnv vdefgs), mkTyEnv tdefs)
+
+tdefIsUsed :: S.Set (Qual Tcon) -> S.Set (Qual Dcon) -> Tdef -> Bool
+tdefIsUsed tcs dcs (Data qtc _ cdefs) = 
+  (qtc `S.member` tcs || any (\ (Constr qdc _ _) -> qdc `S.member` dcs) cdefs)
+tdefIsUsed tcs _ (Newtype qtc qtc_co _ _) = 
+  qtc `S.member` tcs || qtc_co `S.member` tcs
+
+mkVarEnv :: [Vdefg] -> M.Map (Qual Var) Exp
+mkVarEnv vgs =
+  let vdefs = flattenBinds vgs in
+    M.fromList [(v, e) | (Vdef (v, _, e)) <- vdefs]
+
+-- if there is a Newtype qtc qtc_co ty,
+-- generate: qtc |-> ty and qtc_co |-> ty
+-- roughly the same for rhs's of Data decls
+mkTyEnv :: [Tdef] -> M.Map (Qual Tcon) [Ty]
+mkTyEnv tdefs = 
+  M.fromList ([(qtc, [ty]) | (Newtype qtc _ _ ty) <- tdefs]
+         ++   [(qtc, [ty]) | (Newtype _ qtc _ ty) <- tdefs]
+         ++   concatMap (\ td -> case td of
+                                   Data qtc _ cdefs -> [(qtc, concatMap 
+                                     (\ (Constr _ _ ts) -> ts) cdefs)]
+                                   _ -> []) tdefs)
+
+findUsed :: DeadSet -> DeadSet -> DeadM DeadSet
+findUsed _old@(oldVars,oldDcs,oldTcs) _new@(newVars,newDcs,newTcs) = do
+  let (todoVars, todoTcs) = ((S.\\) newVars oldVars, (S.\\) newTcs oldTcs)
+  let nextOld = (oldVars `S.union` todoVars, oldDcs `S.union` newDcs,
+                 oldTcs `S.union` todoTcs)
+  nextStuff <- getVarsAndConsIn (todoVars, todoTcs)
+  if (S.null todoVars && S.null todoTcs)
+      then return nextOld
+      else findUsed nextOld nextStuff
+
+getVarsAndConsIn :: (S.Set (Qual Var), S.Set (Qual Tcon)) -> DeadM DeadSet
+getVarsAndConsIn (vs, tcs) = do
+  vs1 <- mapM varsAndConsInOne (S.toList vs)
+  ts1 <- mapM varsAndConsInOne' (S.toList tcs)
+  let (vs'::[S.Set (Qual Var)], dcs'::[S.Set (Qual Dcon)],
+        ts'::[S.Set (Qual Tcon)]) = unzip3 (vs1 ++ ts1)
+  return (foldl' S.union S.empty vs', foldl' S.union S.empty dcs',
+          foldl' S.union S.empty ts')
+
+varsAndConsInOne :: Qual Var -> DeadM DeadSet
+varsAndConsInOne vr = do
+  def <- findDefn vr
+  return $ maybe emptySet 
+    (noNames emptySet unionThree (mkQ emptySet usedNamesAll)) def
+
+varsAndConsInOne' :: Qual Tcon -> DeadM DeadSet
+varsAndConsInOne' tc = do
+  ty <- findRepTy tc
+  return $ maybe emptySet
+    (noNames emptySet unionThree 
+              (mkQ emptySet usedStuffTys)) ty
+
+emptySet :: DeadSet
+emptySet = (S.empty, S.empty, S.empty)
+mkStartSet :: AnMname -> [Vdefg] -> DeadSet
+-- Initially, we assume the definitions of any exported functions are not
+-- dead, and work backwards from there.
+mkStartSet mn vds = 
+  (S.fromList (filter ((== Just mn) . getModule) (exportedNames vds)), 
+   S.empty, S.empty)
+
+exportedNames :: [Vdefg] -> [Qual Var]
+exportedNames vdefgs = 
+  let vds = flattenBinds vdefgs in
+    filter isQual (vdefNames vds)
+      where isQual    = isJust . fst
+            vdefNames = map (\ (Vdef (n,_,_)) -> n)
+
+
+type DeadSet = (S.Set (Qual Var), S.Set (Qual Dcon), S.Set (Qual Tcon))
+type DeadM = Reader (M.Map (Qual Var) Exp, M.Map (Qual Tcon) [Ty])
+
+findDefn :: Qual Var -> DeadM (Maybe Exp)
+findDefn vr = asks ((M.lookup vr) . fst)
+findRepTy :: Qual Tcon -> DeadM (Maybe [Ty])
+findRepTy tc = asks ((M.lookup tc) . snd)
+
+unionThree :: DeadSet -> DeadSet -> DeadSet
+unionThree (a,b,c) (d,e,f) = (a `S.union` d, b `S.union` e, c `S.union` f)
+
+usedNamesAll :: Exp -> DeadSet
+usedNamesAll = (noNames emptySet unionThree 
+  ((mkQ emptySet usedStuff) `extQ` usedStuffTys `extQ` usedStuffAlts))
+            
+usedStuff :: Exp -> DeadSet
+usedStuff (Var qv)  = (S.singleton qv, S.empty, S.empty)
+usedStuff (Dcon dc) = (S.empty, S.singleton dc, S.empty)
+usedStuff _         = emptySet
+
+usedStuffAlts :: Alt -> DeadSet
+usedStuffAlts (Acon qdc _ _ _) = (S.empty, S.singleton qdc, S.empty)
+usedStuffAlts _ = emptySet
+
+usedStuffTys :: Ty -> DeadSet
+usedStuffTys (Tcon qtc) = (S.empty, S.empty, S.singleton qtc)
+usedStuffTys _          = emptySet
diff --git a/utils/ext-core/Language/Core/Prims.hs b/utils/ext-core/Language/Core/Prims.hs
index 181977b93c..13254d350c 100644
--- a/utils/ext-core/Language/Core/Prims.hs
+++ b/utils/ext-core/Language/Core/Prims.hs
@@ -7,7 +7,7 @@
 module Language.Core.Prims(initialEnv, primEnv, primId, bv,
              tIntzh, tInt64zh, tCharzh, tFloatzh, tAddrzh, tDoublezh, tcStatezh,
              tWordzh, tWord64zh, tByteArrayzh,
-             tcStablePtrzh, tcIO, mkInitialEnv, mkTypeEnv, tRWS, tBool,
+             tcStablePtrzh, tcIO, mkInitialEnv, mkTypeEnv, tRWS, tBool, tcBool,
              ioBaseMname) where
 
 import Control.Monad
diff --git a/utils/ext-core/extcore.cabal b/utils/ext-core/extcore.cabal
index 51f4b22867..ea3b13282c 100644
--- a/utils/ext-core/extcore.cabal
+++ b/utils/ext-core/extcore.cabal
@@ -10,7 +10,7 @@ maintainer:          chevalier@alum.wellesley.edu
 stability:           alpha
 build-depends:       base, containers, directory, filepath, mtl, parsec, pretty
 exposed-modules:     Language.Core.Check, Language.Core.Dependencies, Language.Core.Core, Language.Core.Interp, Language.Core.Overrides, Language.Core.ParsecParser, Language.Core.Prep, Language.Core.Prims, Language.Core.Printer, Language.Core.Merge, Language.Core.ElimDeadCode, Language.Core.Encoding, Language.Core.Env
-other-modules:       Language.Core.PrimCoercions, Language.Core.PrimEnv, Language.Core.Utils, Language.Core.CoreUtils
+other-modules:       Language.Core.PrimCoercions, Language.Core.PrimEnv, Language.Core.Utils, Language.Core.CoreUtils, Language.Core.Environments
 extensions:          DeriveDataTypeable PatternGuards PatternSignatures
 ghc-options:         -Wall -O2
 tested-with:         GHC ==6.8.2