Merge pull request #69 from grin-compiler/andorp/dead-parameter-elim

Improve Simple Dead Parameter Elimination.

grin/src/Transformations/Optimising/SimpleDeadParameterElimination.hs

@@ -7,23 +7,81 @@ import Data.Map (Map)
 import qualified Data.Map as Map
 import Data.Maybe (mapMaybe)
 
+import Control.Monad (msum)
+import Data.Maybe
 import Data.Functor.Foldable as Foldable
 import qualified Data.Foldable
 import Grin.Grin
 import Transformations.Util
 
-collectUsedNames :: Exp -> Set Name
-collectUsedNames = cata folder where
-  folder exp = foldNameUseExpF Set.singleton exp `mappend` Data.Foldable.fold exp
+{-
+Check if the parameter name is used in value position.
+ * pure val1
+ * store val, fetch name, update name val
+ * otherfun val1 ... valn
+ * recfun val1 ... valn -- ignore if the name is used in its original pos
+ * case val of
+-}
+
+-- Checks if the given name are present in the value and and return
+-- Just name if is there, otherwise Nothing.
+nameInVal :: Val -> Name -> Maybe Name
+nameInVal val name = case val of
+  Var name0 | name == name0 -> Just name
+            | otherwise     -> Nothing
+  ConstTagNode  tag  vals   -> msum $ map (`nameInVal` name) vals
+  VarTagNode    _    vals   -> msum $ map (`nameInVal` name) vals
+  ValTag        tag         -> Nothing
+  Unit                      -> Nothing
+  Lit           lit         -> Nothing
+  Undefined     typ         -> Nothing
+
+collectUsedArguments :: Name -> [(Int, Name)] -> Exp -> Set Name
+collectUsedArguments fun args = cata collect where
+
+  -- Collect all the arguments that are referred in the given val
+  argsInVal :: Val -> Set Name
+  argsInVal val = Set.fromList $ mapMaybe (nameInVal val . snd) args
+
+  -- Collect the name if is an argument
+  isArg :: Name -> Set Name
+  isArg n = if n `elem` (snd <$> args) then Set.singleton n else Set.empty
+
+  -- Collect the name when the recursive argument is out of its calling index.
+  recursiveArg :: (Int, Val) -> Set Name
+  recursiveArg (i, Var n)
+    = if (i,n) `elem` args
+        then Set.empty -- recursively used
+        else Set.singleton n
+  recursiveArg (_, val)
+    = argsInVal val
+
+  -- Collect all the args that are not recursively (in their original place) used
+  collect :: ExpF (Set Name) -> Set Name
+  collect = \case
+    SReturnF val      -> argsInVal val
+    SStoreF val       -> argsInVal val
+    SFetchF name      -> isArg name
+    SUpdateF name val -> isArg name <> argsInVal val
+    SAppF name vals
+      | name == fun -> mconcat $ map recursiveArg $ zip [0..] vals
+      | otherwise   -> mconcat $ map argsInVal vals
+    ECaseF val names -> argsInVal val <> mconcat names
+    exp -> Data.Foldable.fold exp
+
+type DeadArgMap
+  = Map
+    Name      -- Name  of the function
+    (Set Int) -- Index of dead argument
 
 simpleDeadParameterElimination :: Program -> Program
 simpleDeadParameterElimination prog@(Program exts defs) = ana builder prog where
-  deadArgMap :: Map Name (Set Int)
+  deadArgMap :: DeadArgMap
   deadArgMap = mconcat $ mapMaybe deadArgsInDef defs
 
-  deadArgsInDef :: Def -> Maybe (Map Name (Set Int))
-  deadArgsInDef def@(Def name args _)
-    | usedNames       <- collectUsedNames def
+  deadArgsInDef :: Def -> Maybe DeadArgMap
+  deadArgsInDef def@(Def name args body)
+    | usedNames       <- collectUsedArguments name (zip [0..] args) body
     , deadArgIndices  <- Set.fromList . map fst . filter (flip Set.notMember usedNames . snd) $ zip [0..] args
     = if null deadArgIndices
         then Nothing
@@ -70,4 +128,4 @@ simpleDeadParameterElimination prog@(Program exts defs) = ana builder prog where
   isPFtag = \case
     F{} -> True
     P{} -> True
-    _   -> False
+    _   -> False

grin/test/Transformations/Optimising/SimpleDeadParameterEliminationSpec.hs

@@ -25,6 +25,30 @@ spec = do
       |]
     simpleDeadParameterElimination before `sameAs` after
 
+  it "recursive non-used parameter" $ do
+    let before = [prog|
+          fun f1 f2 =
+            f3 <- fun2 f1
+            fun f1 f2
+      |]
+    let after = [prog|
+          fun f1 =
+            f3 <- fun2 f1
+            fun f1
+      |]
+    simpleDeadParameterElimination before `sameAs` after
+
+  it "recursive switched parameter" $ do
+    let before = [prog|
+        fun f1 f2 =
+          fun f2 f1
+      |]
+    let after = [prog|
+        fun f1 f2 =
+          fun f2 f1
+      |]
+    simpleDeadParameterElimination before `sameAs` after
+
   it "Pnode + Fnode ; val - lpat - cpat" $ do
     let before = [prog|
           funA a b = pure b
@@ -69,3 +93,4 @@ spec = do
                 pure (P0funA b6)
       |]
     simpleDeadParameterElimination before `sameAs` after
+