[app] properly filter callGraph

app/Main.hs

@@ -19,6 +19,8 @@ import MiniJuvix.Syntax.Concrete.Scoped.Pretty.Html
 import MiniJuvix.Syntax.Concrete.Scoped.Pretty.Base (defaultOptions)
 import qualified MiniJuvix.Syntax.Abstract.Pretty.Ansi as A
 import qualified MiniJuvix.Termination.CallGraph as T
+import qualified Data.Text as Text
+import qualified MiniJuvix.Termination.CallGraph as T
 
 data Command
   = Scope ScopeOptions
@@ -48,13 +50,13 @@ data HtmlOptions = HtmlOptions
 data CallsOptions = CallsOptions
   { _callsInputFile :: FilePath,
     _callsShowIds :: Bool,
-    _callsFunctionNameFilter :: Maybe Text,
+    _callsFunctionNameFilter :: Maybe [Text],
     _callsShowDecreasingArgs :: A.ShowDecrArgs
   }
 
 data CallGraphOptions = CallGraphOptions
   { _graphInputFile :: FilePath,
-    _graphFunctionNameFilter :: Maybe Text
+    _graphFunctionNameFilter :: Maybe [Text]
   }
 
 parseHtml :: Parser HtmlOptions
@@ -89,10 +91,11 @@ parseCalls = do
           <> help "Show the unique number of each identifier"
       )
   _callsFunctionNameFilter <-
-    optional $ option str
+    optional $ Text.words <$> option str
       ( long "function"
           <> short 'f'
-          <> help "Only shows the specified function"
+          <> metavar "fun1 fun2 ..."
+          <> help "Only shows the specified functions"
       )
   _callsShowDecreasingArgs <-
     option decrArgsParser
@@ -116,7 +119,7 @@ parseCallGraph :: Parser CallGraphOptions
 parseCallGraph = do
   _graphInputFile <- parseInputFile
   _graphFunctionNameFilter <-
-    optional $ option str
+    optional $ Text.words <$> option str
       ( long "function"
           <> short 'f'
           <> help "Only shows the specified function"
@@ -296,14 +299,12 @@ go c = do
       m <- parseModuleIO _graphInputFile
       s <- fromRightIO' printErrorAnsi $ M.scopeCheck1IO root m
       a <- fromRightIO' putStrLn (return $ A.translateModule s)
-      let callMap0 = T.buildCallMap a
-          callMap = case _graphFunctionNameFilter of
-            Nothing -> callMap0
-            Just f -> T.filterCallMap f callMap0
+      let callMap = T.buildCallMap a
           opts' = A.defaultOptions
           completeGraph = T.completeCallGraph callMap
-          recBehav = map T.recursiveBehaviour (T.reflexiveEdges completeGraph)
-      A.printPrettyCode opts' completeGraph
+          filteredGraph = maybe completeGraph (`T.unsafeFilterGraph` completeGraph) _graphFunctionNameFilter
+          recBehav = map T.recursiveBehaviour (T.reflexiveEdges filteredGraph)
+      A.printPrettyCode opts' filteredGraph
       putStrLn ""
       forM_ recBehav $ \r -> do
         let n = M.renderPrettyCode M.defaultOptions $ A._recBehaviourFunction r
@@ -311,10 +312,8 @@ go c = do
         putStrLn ""
         case T.findOrder r of
           Nothing -> putStrLn (n <> " Fails the termination checking")
-          Just (T.LexOrder k) -> putStrLn (n<> " Terminates with order " <> show (toList k))
+          Just (T.LexOrder k) -> putStrLn (n <> " Terminates with order " <> show (toList k))
         putStrLn ""
 
-
-
 main :: IO ()
 main = execParser descr >>= go

src/MiniJuvix/Syntax/Concrete/Scoped/Name.hs

@@ -107,6 +107,9 @@ isConstructor Name' {..} = case _nameKind of
 fromQualifiedName :: C.QualifiedName -> C.Symbol
 fromQualifiedName (C.QualifiedName _ s) = s
 
+symbolText :: Symbol -> Text
+symbolText = C._symbolText . _nameConcrete
+
 fromName :: Name -> Symbol
 fromName Name' {..} = Name' {_nameConcrete = unqual, ..}
   where

src/MiniJuvix/Termination/CallGraph.hs

@@ -8,6 +8,7 @@ import MiniJuvix.Prelude
 import MiniJuvix.Syntax.Abstract.Language.Extra
 import qualified Data.HashMap.Strict as HashMap
 import MiniJuvix.Termination.Types
+import qualified MiniJuvix.Syntax.Concrete.Scoped.Name as S
 import Prettyprinter as PP
 import MiniJuvix.Syntax.Abstract.Pretty.Base
 import qualified Data.HashSet as HashSet
@@ -69,6 +70,12 @@ fromFunCall caller fc =
           _callMatrix = map fst (fc ^. callArgs)
        }
 
+-- | IMPORTANT: the resulting call graph is not complete. Use this function
+-- only to filter the pretty printed graph
+unsafeFilterGraph :: [Text] -> CompleteCallGraph -> CompleteCallGraph
+unsafeFilterGraph funNames (CompleteCallGraph g) =
+  CompleteCallGraph (HashMap.filterWithKey  (\(f , _) _ -> S.symbolText f `elem`funNames) g)
+
 completeCallGraph :: CallMap -> CompleteCallGraph
 completeCallGraph cm = CompleteCallGraph (go startingEdges)
   where

src/MiniJuvix/Termination/Types.hs

@@ -99,8 +99,5 @@ instance PrettyCode CallRow where
 instance PrettyCode CallMatrix where
   ppCode l = vsep <$> mapM ppCode l
 
-filterCallMap :: Text -> CallMap -> CallMap
-filterCallMap funName = over callMap (HashMap.filterWithKey (\k _ -> funStr k == funName))
-  where
-  funStr :: A.FunctionName -> Text
-  funStr = C._symbolText . S._nameConcrete
+filterCallMap :: [Text] -> CallMap -> CallMap
+filterCallMap funNames = over callMap (HashMap.filterWithKey (\k _ -> S.symbolText k `elem` funNames))