{-# LANGUAGE RecordWildCards #-}

-- | This modules manages the two ways we remove nodes from a graph; collapsing and hiding.
--
-- Collapsing means absorbing a node's descendants into itself, including all edges.
--
-- Hiding means removing a node (and its descendants), moving the edges to the node's parent, if a parent exist.
--
-- There's also the special option --collapse-modules.
-- It's undeniably a little hacky, but for now this is the best home for that functionality.
-- Functionality-wise, it's still essentially just collapsing nodes into one another.
-- The thing that makes it hacky is that it then uses a value node to represent a module.
-- This is not actually a huge deal, because no other module actually cares about the node type, but it's something to watch out for.
-- There's more design discussion on https://github.com/jonascarpay/calligraphy/pull/5
module Calligraphy.Phases.NodeFilter (filterNodes, NodeFilterConfig (..), pNodeFilterConfig) where

import Calligraphy.Util.Types
import Control.Monad.State
import Data.EnumMap (EnumMap)
import qualified Data.EnumMap as EnumMap
import Data.Maybe (catMaybes)
import Data.Tree
import Options.Applicative

data Mode = Show | Collapse | Hide
  deriving (Mode -> Mode -> Bool
(Mode -> Mode -> Bool) -> (Mode -> Mode -> Bool) -> Eq Mode
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Mode -> Mode -> Bool
$c/= :: Mode -> Mode -> Bool
== :: Mode -> Mode -> Bool
$c== :: Mode -> Mode -> Bool
Eq, Int -> Mode -> ShowS
[Mode] -> ShowS
Mode -> String
(Int -> Mode -> ShowS)
-> (Mode -> String) -> ([Mode] -> ShowS) -> Show Mode
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Mode] -> ShowS
$cshowList :: [Mode] -> ShowS
show :: Mode -> String
$cshow :: Mode -> String
showsPrec :: Int -> Mode -> ShowS
$cshowsPrec :: Int -> Mode -> ShowS
Show)

data NodeFilterConfig = NodeFilterConfig
  { NodeFilterConfig -> Bool
hideLocals :: Bool,
    NodeFilterConfig -> Bool
collapseModules :: Bool,
    NodeFilterConfig -> Mode
collapseClasses :: Mode,
    NodeFilterConfig -> Mode
collapseData :: Mode,
    NodeFilterConfig -> Mode
collapseValues :: Mode,
    NodeFilterConfig -> Mode
collapseConstructors :: Mode,
    NodeFilterConfig -> Bool
hideRecords :: Bool
  }

pNodeFilterConfig :: Parser NodeFilterConfig
pNodeFilterConfig :: Parser NodeFilterConfig
pNodeFilterConfig =
  Bool
-> Bool -> Mode -> Mode -> Mode -> Mode -> Bool -> NodeFilterConfig
NodeFilterConfig
    (Bool
 -> Bool
 -> Mode
 -> Mode
 -> Mode
 -> Mode
 -> Bool
 -> NodeFilterConfig)
-> Parser Bool
-> Parser
     (Bool -> Mode -> Mode -> Mode -> Mode -> Bool -> NodeFilterConfig)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mod FlagFields Bool -> Parser Bool
switch (String -> Mod FlagFields Bool
forall (f :: * -> *) a. HasName f => String -> Mod f a
long String
"exports-only" Mod FlagFields Bool -> Mod FlagFields Bool -> Mod FlagFields Bool
forall a. Semigroup a => a -> a -> a
<> String -> Mod FlagFields Bool
forall (f :: * -> *) a. HasName f => String -> Mod f a
long String
"hide-local-bindings" Mod FlagFields Bool -> Mod FlagFields Bool -> Mod FlagFields Bool
forall a. Semigroup a => a -> a -> a
<> String -> Mod FlagFields Bool
forall (f :: * -> *) a. String -> Mod f a
help String
"Remove all non-exported bindings, merging all edges into its parent, if one exist.")
    Parser
  (Bool -> Mode -> Mode -> Mode -> Mode -> Bool -> NodeFilterConfig)
-> Parser Bool
-> Parser
     (Mode -> Mode -> Mode -> Mode -> Bool -> NodeFilterConfig)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mod FlagFields Bool -> Parser Bool
switch (String -> Mod FlagFields Bool
forall (f :: * -> *) a. HasName f => String -> Mod f a
long String
"collapse-modules" Mod FlagFields Bool -> Mod FlagFields Bool -> Mod FlagFields Bool
forall a. Semigroup a => a -> a -> a
<> String -> Mod FlagFields Bool
forall (f :: * -> *) a. String -> Mod f a
help String
"Collapse all nodes into a single node per module.")
    Parser (Mode -> Mode -> Mode -> Mode -> Bool -> NodeFilterConfig)
-> Parser Mode
-> Parser (Mode -> Mode -> Mode -> Bool -> NodeFilterConfig)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> String -> String -> Parser Mode
pMode String
"classes" String
"class"
    Parser (Mode -> Mode -> Mode -> Bool -> NodeFilterConfig)
-> Parser Mode -> Parser (Mode -> Mode -> Bool -> NodeFilterConfig)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> String -> String -> Parser Mode
pMode String
"data" String
"data"
    Parser (Mode -> Mode -> Bool -> NodeFilterConfig)
-> Parser Mode -> Parser (Mode -> Bool -> NodeFilterConfig)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> String -> String -> Parser Mode
pMode String
"values" String
"value"
    Parser (Mode -> Bool -> NodeFilterConfig)
-> Parser Mode -> Parser (Bool -> NodeFilterConfig)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> String -> String -> Parser Mode
pMode String
"constructors" String
"constructor"
    Parser (Bool -> NodeFilterConfig)
-> Parser Bool -> Parser NodeFilterConfig
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Bool -> Bool -> Mod FlagFields Bool -> Parser Bool
forall a. a -> a -> Mod FlagFields a -> Parser a
flag Bool
False Bool
True (String -> Mod FlagFields Bool
forall (f :: * -> *) a. HasName f => String -> Mod f a
long String
"hide-records" Mod FlagFields Bool -> Mod FlagFields Bool -> Mod FlagFields Bool
forall a. Semigroup a => a -> a -> a
<> String -> Mod FlagFields Bool
forall (f :: * -> *) a. String -> Mod f a
help String
"Remove all record nodes.")
  where
    pMode :: String -> String -> Parser Mode
    pMode :: String -> String -> Parser Mode
pMode String
flagName String
helpName =
      Mode -> Mod FlagFields Mode -> Parser Mode
forall a. a -> Mod FlagFields a -> Parser a
flag' Mode
Collapse (String -> Mod FlagFields Mode
forall (f :: * -> *) a. HasName f => String -> Mod f a
long (String
"collapse-" String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
flagName) Mod FlagFields Mode -> Mod FlagFields Mode -> Mod FlagFields Mode
forall a. Semigroup a => a -> a -> a
<> String -> Mod FlagFields Mode
forall (f :: * -> *) a. String -> Mod f a
help String
collapseHelp)
        Parser Mode -> Parser Mode -> Parser Mode
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Mode -> Mod FlagFields Mode -> Parser Mode
forall a. a -> Mod FlagFields a -> Parser a
flag' Mode
Hide (String -> Mod FlagFields Mode
forall (f :: * -> *) a. HasName f => String -> Mod f a
long (String
"hide-" String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
flagName) Mod FlagFields Mode -> Mod FlagFields Mode -> Mod FlagFields Mode
forall a. Semigroup a => a -> a -> a
<> String -> Mod FlagFields Mode
forall (f :: * -> *) a. String -> Mod f a
help String
hideHelp)
        Parser Mode -> Parser Mode -> Parser Mode
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Mode -> Parser Mode
forall (f :: * -> *) a. Applicative f => a -> f a
pure Mode
Show
      where
        collapseHelp :: String
collapseHelp = String
"Remove all " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
helpName String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" nodes's children, merging the children's edges into itself."
        hideHelp :: String
hideHelp = String
"Remove all " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
helpName String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" nodes and their children."

filterNodes :: NodeFilterConfig -> CallGraph -> CallGraph
filterNodes :: NodeFilterConfig -> CallGraph -> CallGraph
filterNodes NodeFilterConfig {Bool
Mode
hideRecords :: Bool
collapseConstructors :: Mode
collapseValues :: Mode
collapseData :: Mode
collapseClasses :: Mode
collapseModules :: Bool
hideLocals :: Bool
hideRecords :: NodeFilterConfig -> Bool
collapseConstructors :: NodeFilterConfig -> Mode
collapseValues :: NodeFilterConfig -> Mode
collapseData :: NodeFilterConfig -> Mode
collapseClasses :: NodeFilterConfig -> Mode
collapseModules :: NodeFilterConfig -> Bool
hideLocals :: NodeFilterConfig -> Bool
..} (CallGraph [Module]
modules Set (Key, Key)
calls Set (Key, Key)
types) =
  let ([Module]
modules', EnumMap Key Key
reps) =
        (State (EnumMap Key Key) [Module]
 -> EnumMap Key Key -> ([Module], EnumMap Key Key))
-> EnumMap Key Key
-> State (EnumMap Key Key) [Module]
-> ([Module], EnumMap Key Key)
forall a b c. (a -> b -> c) -> b -> a -> c
flip State (EnumMap Key Key) [Module]
-> EnumMap Key Key -> ([Module], EnumMap Key Key)
forall s a. State s a -> s -> (a, s)
runState EnumMap Key Key
forall a. Monoid a => a
mempty (State (EnumMap Key Key) [Module] -> ([Module], EnumMap Key Key))
-> State (EnumMap Key Key) [Module] -> ([Module], EnumMap Key Key)
forall a b. (a -> b) -> a -> b
$
          [Module]
-> (Module -> StateT (EnumMap Key Key) Identity Module)
-> State (EnumMap Key Key) [Module]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Module]
modules ((Module -> StateT (EnumMap Key Key) Identity Module)
 -> State (EnumMap Key Key) [Module])
-> (Module -> StateT (EnumMap Key Key) Identity Module)
-> State (EnumMap Key Key) [Module]
forall a b. (a -> b) -> a -> b
$
            if Bool
collapseModules
              then Module -> StateT (EnumMap Key Key) Identity Module
collapseModule
              else (Forest Decl -> StateT (EnumMap Key Key) Identity (Forest Decl))
-> Module -> StateT (EnumMap Key Key) Identity Module
Traversal' Module (Forest Decl)
modForest (([Maybe (Tree Decl)] -> Forest Decl)
-> StateT (EnumMap Key Key) Identity [Maybe (Tree Decl)]
-> StateT (EnumMap Key Key) Identity (Forest Decl)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Maybe (Tree Decl)] -> Forest Decl
forall a. [Maybe a] -> [a]
catMaybes (StateT (EnumMap Key Key) Identity [Maybe (Tree Decl)]
 -> StateT (EnumMap Key Key) Identity (Forest Decl))
-> (Forest Decl
    -> StateT (EnumMap Key Key) Identity [Maybe (Tree Decl)])
-> Forest Decl
-> StateT (EnumMap Key Key) Identity (Forest Decl)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Tree Decl
 -> StateT (EnumMap Key Key) Identity (Maybe (Tree Decl)))
-> Forest Decl
-> StateT (EnumMap Key Key) Identity [Maybe (Tree Decl)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Maybe Decl
-> Tree Decl
-> StateT (EnumMap Key Key) Identity (Maybe (Tree Decl))
go Maybe Decl
forall a. Maybe a
Nothing))
   in [Module] -> Set (Key, Key) -> Set (Key, Key) -> CallGraph
CallGraph [Module]
modules' (EnumMap Key Key -> Set (Key, Key) -> Set (Key, Key)
forall a b.
(Enum a, Ord b) =>
EnumMap a b -> Set (a, a) -> Set (b, b)
rekeyCalls EnumMap Key Key
reps Set (Key, Key)
calls) (EnumMap Key Key -> Set (Key, Key) -> Set (Key, Key)
forall a b.
(Enum a, Ord b) =>
EnumMap a b -> Set (a, a) -> Set (b, b)
rekeyCalls EnumMap Key Key
reps Set (Key, Key)
types)
  where
    collapseModule :: Module -> State (EnumMap Key Key) Module
    collapseModule :: Module -> StateT (EnumMap Key Key) Identity Module
collapseModule (Module String
modname String
path []) = Module -> StateT (EnumMap Key Key) Identity Module
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Module -> StateT (EnumMap Key Key) Identity Module)
-> Module -> StateT (EnumMap Key Key) Identity Module
forall a b. (a -> b) -> a -> b
$ String -> String -> Forest Decl -> Module
Module String
modname String
path []
    collapseModule (Module String
modname String
path forest :: Forest Decl
forest@(Node Decl
rep Forest Decl
_ : Forest Decl
_)) = do
      let repKey :: Key
repKey = Decl -> Key
declKey Decl
rep
      Traversal (Forest Decl) (Forest Any) Decl Any
-> Forest Decl
-> (Decl -> StateT (EnumMap Key Key) Identity ())
-> StateT (EnumMap Key Key) Identity ()
forall (m :: * -> *) s t a b.
Applicative m =>
Traversal s t a b -> s -> (a -> m ()) -> m ()
forT_ forall a b. Traversal (Forest a) (Forest b) a b
Traversal (Forest Decl) (Forest Any) Decl Any
forestT Forest Decl
forest ((Decl -> StateT (EnumMap Key Key) Identity ())
 -> StateT (EnumMap Key Key) Identity ())
-> (Decl -> StateT (EnumMap Key Key) Identity ())
-> StateT (EnumMap Key Key) Identity ()
forall a b. (a -> b) -> a -> b
$ \Decl
decl -> Key -> Key -> StateT (EnumMap Key Key) Identity ()
assoc (Decl -> Key
declKey Decl
decl) Key
repKey
      Module -> StateT (EnumMap Key Key) Identity Module
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Module -> StateT (EnumMap Key Key) Identity Module)
-> Module -> StateT (EnumMap Key Key) Identity Module
forall a b. (a -> b) -> a -> b
$ String -> String -> Forest Decl -> Module
Module String
modname String
path [Decl -> Forest Decl -> Tree Decl
forall a. a -> Forest a -> Tree a
Node (String -> Key -> EnumSet GHCKey -> Bool -> DeclType -> Loc -> Decl
Decl String
modname Key
repKey EnumSet GHCKey
forall a. Monoid a => a
mempty Bool
True DeclType
ValueDecl (Int -> Int -> Loc
Loc Int
1 Int
1)) []]

    shouldCollapse :: Decl -> Bool
    shouldCollapse :: Decl -> Bool
shouldCollapse Decl
decl = case Decl -> DeclType
declType Decl
decl of
      DeclType
ValueDecl -> Mode
collapseValues Mode -> Mode -> Bool
forall a. Eq a => a -> a -> Bool
== Mode
Collapse
      DeclType
ClassDecl -> Mode
collapseClasses Mode -> Mode -> Bool
forall a. Eq a => a -> a -> Bool
== Mode
Collapse
      DeclType
ConDecl -> Mode
collapseConstructors Mode -> Mode -> Bool
forall a. Eq a => a -> a -> Bool
== Mode
Collapse
      DeclType
DataDecl -> Mode
collapseData Mode -> Mode -> Bool
forall a. Eq a => a -> a -> Bool
== Mode
Collapse
      DeclType
_ -> Bool
False

    shouldHide :: Decl -> Bool
    shouldHide :: Decl -> Bool
shouldHide Decl
decl = DeclType -> Bool
typ (Decl -> DeclType
declType Decl
decl) Bool -> Bool -> Bool
|| (Bool
hideLocals Bool -> Bool -> Bool
&& Bool -> Bool
not (Decl -> Bool
declExported Decl
decl))
      where
        typ :: DeclType -> Bool
typ DeclType
ClassDecl = Mode
collapseClasses Mode -> Mode -> Bool
forall a. Eq a => a -> a -> Bool
== Mode
Hide
        typ DeclType
DataDecl = Mode
collapseData Mode -> Mode -> Bool
forall a. Eq a => a -> a -> Bool
== Mode
Hide
        typ DeclType
ValueDecl = Mode
collapseValues Mode -> Mode -> Bool
forall a. Eq a => a -> a -> Bool
== Mode
Hide
        typ DeclType
ConDecl = Mode
collapseConstructors Mode -> Mode -> Bool
forall a. Eq a => a -> a -> Bool
== Mode
Hide
        typ DeclType
RecDecl = Bool
hideRecords

    go :: Maybe Decl -> Tree Decl -> State (EnumMap Key Key) (Maybe (Tree Decl))
    go :: Maybe Decl
-> Tree Decl
-> StateT (EnumMap Key Key) Identity (Maybe (Tree Decl))
go Maybe Decl
mparent node :: Tree Decl
node@(Node Decl
decl Forest Decl
children)
      | Decl -> Bool
shouldHide Decl
decl = do
          Maybe Decl
-> (Decl -> StateT (EnumMap Key Key) Identity ())
-> StateT (EnumMap Key Key) Identity ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe Decl
mparent ((Decl -> StateT (EnumMap Key Key) Identity ())
 -> StateT (EnumMap Key Key) Identity ())
-> (Decl -> StateT (EnumMap Key Key) Identity ())
-> StateT (EnumMap Key Key) Identity ()
forall a b. (a -> b) -> a -> b
$ \Decl
parent ->
            Tree Decl
-> (Decl -> StateT (EnumMap Key Key) Identity ())
-> StateT (EnumMap Key Key) Identity ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Tree Decl
node ((Decl -> StateT (EnumMap Key Key) Identity ())
 -> StateT (EnumMap Key Key) Identity ())
-> (Decl -> StateT (EnumMap Key Key) Identity ())
-> StateT (EnumMap Key Key) Identity ()
forall a b. (a -> b) -> a -> b
$ \Decl
child -> Key -> Key -> StateT (EnumMap Key Key) Identity ()
assoc (Decl -> Key
declKey Decl
child) (Decl -> Key
declKey Decl
parent)
          Maybe (Tree Decl)
-> StateT (EnumMap Key Key) Identity (Maybe (Tree Decl))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Tree Decl)
forall a. Maybe a
Nothing
      | Decl -> Bool
shouldCollapse Decl
decl = do
          Tree Decl
-> (Decl -> StateT (EnumMap Key Key) Identity ())
-> StateT (EnumMap Key Key) Identity ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Tree Decl
node ((Decl -> StateT (EnumMap Key Key) Identity ())
 -> StateT (EnumMap Key Key) Identity ())
-> (Decl -> StateT (EnumMap Key Key) Identity ())
-> StateT (EnumMap Key Key) Identity ()
forall a b. (a -> b) -> a -> b
$ \Decl
child ->
            Key -> Key -> StateT (EnumMap Key Key) Identity ()
assoc (Decl -> Key
declKey Decl
child) (Decl -> Key
declKey Decl
decl)
          Maybe (Tree Decl)
-> StateT (EnumMap Key Key) Identity (Maybe (Tree Decl))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Tree Decl)
 -> StateT (EnumMap Key Key) Identity (Maybe (Tree Decl)))
-> Maybe (Tree Decl)
-> StateT (EnumMap Key Key) Identity (Maybe (Tree Decl))
forall a b. (a -> b) -> a -> b
$ Tree Decl -> Maybe (Tree Decl)
forall a. a -> Maybe a
Just (Tree Decl -> Maybe (Tree Decl)) -> Tree Decl -> Maybe (Tree Decl)
forall a b. (a -> b) -> a -> b
$ Decl -> Forest Decl -> Tree Decl
forall a. a -> Forest a -> Tree a
Node Decl
decl []
      | Bool
otherwise = do
          Key -> Key -> StateT (EnumMap Key Key) Identity ()
assoc (Decl -> Key
declKey Decl
decl) (Decl -> Key
declKey Decl
decl)
          Forest Decl
children' <- [Maybe (Tree Decl)] -> Forest Decl
forall a. [Maybe a] -> [a]
catMaybes ([Maybe (Tree Decl)] -> Forest Decl)
-> StateT (EnumMap Key Key) Identity [Maybe (Tree Decl)]
-> StateT (EnumMap Key Key) Identity (Forest Decl)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Tree Decl
 -> StateT (EnumMap Key Key) Identity (Maybe (Tree Decl)))
-> Forest Decl
-> StateT (EnumMap Key Key) Identity [Maybe (Tree Decl)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Maybe Decl
-> Tree Decl
-> StateT (EnumMap Key Key) Identity (Maybe (Tree Decl))
go (Decl -> Maybe Decl
forall a. a -> Maybe a
Just Decl
decl)) Forest Decl
children
          Maybe (Tree Decl)
-> StateT (EnumMap Key Key) Identity (Maybe (Tree Decl))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Tree Decl)
 -> StateT (EnumMap Key Key) Identity (Maybe (Tree Decl)))
-> Maybe (Tree Decl)
-> StateT (EnumMap Key Key) Identity (Maybe (Tree Decl))
forall a b. (a -> b) -> a -> b
$ Tree Decl -> Maybe (Tree Decl)
forall a. a -> Maybe a
Just (Tree Decl -> Maybe (Tree Decl)) -> Tree Decl -> Maybe (Tree Decl)
forall a b. (a -> b) -> a -> b
$ Decl -> Forest Decl -> Tree Decl
forall a. a -> Forest a -> Tree a
Node Decl
decl Forest Decl
children'

    assoc :: Key -> Key -> State (EnumMap Key Key) ()
    assoc :: Key -> Key -> StateT (EnumMap Key Key) Identity ()
assoc Key
key Key
rep = (EnumMap Key Key -> EnumMap Key Key)
-> StateT (EnumMap Key Key) Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (Key -> Key -> EnumMap Key Key -> EnumMap Key Key
forall k a. Enum k => k -> a -> EnumMap k a -> EnumMap k a
EnumMap.insert Key
key Key
rep)