{-# LANGUAGE OverloadedStrings #-}

-- | This module exports functionality for generating a call graph of
-- an Futhark program.
module Futhark.Analysis.CallGraph
  ( CallGraph,
    buildCallGraph,
    isFunInCallGraph,
    calls,
    calledByConsts,
    allCalledBy,
    numOccurences,
    findNoninlined,
  )
where

import Control.Monad.Writer.Strict
import Data.List (foldl')
import qualified Data.Map.Strict as M
import Data.Maybe (isJust)
import qualified Data.Set as S
import Futhark.IR.SOACS
import Futhark.Util.Pretty

type FunctionTable = M.Map Name (FunDef SOACS)

buildFunctionTable :: Prog SOACS -> FunctionTable
buildFunctionTable :: Prog SOACS -> FunctionTable
buildFunctionTable = (FunctionTable -> FunDef SOACS -> FunctionTable)
-> FunctionTable -> [FunDef SOACS] -> FunctionTable
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl FunctionTable -> FunDef SOACS -> FunctionTable
forall rep.
Map Name (FunDef rep) -> FunDef rep -> Map Name (FunDef rep)
expand FunctionTable
forall k a. Map k a
M.empty ([FunDef SOACS] -> FunctionTable)
-> (Prog SOACS -> [FunDef SOACS]) -> Prog SOACS -> FunctionTable
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog SOACS -> [FunDef SOACS]
forall rep. Prog rep -> [FunDef rep]
progFuns
  where
    expand :: Map Name (FunDef rep) -> FunDef rep -> Map Name (FunDef rep)
expand Map Name (FunDef rep)
ftab FunDef rep
f = Name
-> FunDef rep -> Map Name (FunDef rep) -> Map Name (FunDef rep)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (FunDef rep -> Name
forall rep. FunDef rep -> Name
funDefName FunDef rep
f) FunDef rep
f Map Name (FunDef rep)
ftab

-- | A unique (at least within a function) name identifying a function
-- call.  In practice the first element of the corresponding pattern.
type CallId = VName

data FunCalls = FunCalls
  { FunCalls -> Map CallId (Attrs, Name)
fcMap :: M.Map CallId (Attrs, Name),
    FunCalls -> Set Name
fcAllCalled :: S.Set Name
  }
  deriving (FunCalls -> FunCalls -> Bool
(FunCalls -> FunCalls -> Bool)
-> (FunCalls -> FunCalls -> Bool) -> Eq FunCalls
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: FunCalls -> FunCalls -> Bool
$c/= :: FunCalls -> FunCalls -> Bool
== :: FunCalls -> FunCalls -> Bool
$c== :: FunCalls -> FunCalls -> Bool
Eq, Eq FunCalls
Eq FunCalls
-> (FunCalls -> FunCalls -> Ordering)
-> (FunCalls -> FunCalls -> Bool)
-> (FunCalls -> FunCalls -> Bool)
-> (FunCalls -> FunCalls -> Bool)
-> (FunCalls -> FunCalls -> Bool)
-> (FunCalls -> FunCalls -> FunCalls)
-> (FunCalls -> FunCalls -> FunCalls)
-> Ord FunCalls
FunCalls -> FunCalls -> Bool
FunCalls -> FunCalls -> Ordering
FunCalls -> FunCalls -> FunCalls
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: FunCalls -> FunCalls -> FunCalls
$cmin :: FunCalls -> FunCalls -> FunCalls
max :: FunCalls -> FunCalls -> FunCalls
$cmax :: FunCalls -> FunCalls -> FunCalls
>= :: FunCalls -> FunCalls -> Bool
$c>= :: FunCalls -> FunCalls -> Bool
> :: FunCalls -> FunCalls -> Bool
$c> :: FunCalls -> FunCalls -> Bool
<= :: FunCalls -> FunCalls -> Bool
$c<= :: FunCalls -> FunCalls -> Bool
< :: FunCalls -> FunCalls -> Bool
$c< :: FunCalls -> FunCalls -> Bool
compare :: FunCalls -> FunCalls -> Ordering
$ccompare :: FunCalls -> FunCalls -> Ordering
$cp1Ord :: Eq FunCalls
Ord, Int -> FunCalls -> ShowS
[FunCalls] -> ShowS
FunCalls -> String
(Int -> FunCalls -> ShowS)
-> (FunCalls -> String) -> ([FunCalls] -> ShowS) -> Show FunCalls
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [FunCalls] -> ShowS
$cshowList :: [FunCalls] -> ShowS
show :: FunCalls -> String
$cshow :: FunCalls -> String
showsPrec :: Int -> FunCalls -> ShowS
$cshowsPrec :: Int -> FunCalls -> ShowS
Show)

instance Monoid FunCalls where
  mempty :: FunCalls
mempty = Map CallId (Attrs, Name) -> Set Name -> FunCalls
FunCalls Map CallId (Attrs, Name)
forall a. Monoid a => a
mempty Set Name
forall a. Monoid a => a
mempty

instance Semigroup FunCalls where
  FunCalls Map CallId (Attrs, Name)
x1 Set Name
y1 <> :: FunCalls -> FunCalls -> FunCalls
<> FunCalls Map CallId (Attrs, Name)
x2 Set Name
y2 = Map CallId (Attrs, Name) -> Set Name -> FunCalls
FunCalls (Map CallId (Attrs, Name)
x1 Map CallId (Attrs, Name)
-> Map CallId (Attrs, Name) -> Map CallId (Attrs, Name)
forall a. Semigroup a => a -> a -> a
<> Map CallId (Attrs, Name)
x2) (Set Name
y1 Set Name -> Set Name -> Set Name
forall a. Semigroup a => a -> a -> a
<> Set Name
y2)

fcCalled :: Name -> FunCalls -> Bool
fcCalled :: Name -> FunCalls -> Bool
fcCalled Name
f FunCalls
fcs = Name
f Name -> Set Name -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` FunCalls -> Set Name
fcAllCalled FunCalls
fcs

type FunGraph = M.Map Name FunCalls

-- | The call graph is a mapping from a function name, i.e., the
-- caller, to a record of the names of functions called *directly* (not
-- transitively!) by the function.
--
-- We keep track separately of the functions called by constants.
data CallGraph = CallGraph
  { CallGraph -> FunGraph
cgCalledByFuns :: FunGraph,
    CallGraph -> FunCalls
cgCalledByConsts :: FunCalls
  }
  deriving (CallGraph -> CallGraph -> Bool
(CallGraph -> CallGraph -> Bool)
-> (CallGraph -> CallGraph -> Bool) -> Eq CallGraph
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: CallGraph -> CallGraph -> Bool
$c/= :: CallGraph -> CallGraph -> Bool
== :: CallGraph -> CallGraph -> Bool
$c== :: CallGraph -> CallGraph -> Bool
Eq, Eq CallGraph
Eq CallGraph
-> (CallGraph -> CallGraph -> Ordering)
-> (CallGraph -> CallGraph -> Bool)
-> (CallGraph -> CallGraph -> Bool)
-> (CallGraph -> CallGraph -> Bool)
-> (CallGraph -> CallGraph -> Bool)
-> (CallGraph -> CallGraph -> CallGraph)
-> (CallGraph -> CallGraph -> CallGraph)
-> Ord CallGraph
CallGraph -> CallGraph -> Bool
CallGraph -> CallGraph -> Ordering
CallGraph -> CallGraph -> CallGraph
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: CallGraph -> CallGraph -> CallGraph
$cmin :: CallGraph -> CallGraph -> CallGraph
max :: CallGraph -> CallGraph -> CallGraph
$cmax :: CallGraph -> CallGraph -> CallGraph
>= :: CallGraph -> CallGraph -> Bool
$c>= :: CallGraph -> CallGraph -> Bool
> :: CallGraph -> CallGraph -> Bool
$c> :: CallGraph -> CallGraph -> Bool
<= :: CallGraph -> CallGraph -> Bool
$c<= :: CallGraph -> CallGraph -> Bool
< :: CallGraph -> CallGraph -> Bool
$c< :: CallGraph -> CallGraph -> Bool
compare :: CallGraph -> CallGraph -> Ordering
$ccompare :: CallGraph -> CallGraph -> Ordering
$cp1Ord :: Eq CallGraph
Ord, Int -> CallGraph -> ShowS
[CallGraph] -> ShowS
CallGraph -> String
(Int -> CallGraph -> ShowS)
-> (CallGraph -> String)
-> ([CallGraph] -> ShowS)
-> Show CallGraph
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [CallGraph] -> ShowS
$cshowList :: [CallGraph] -> ShowS
show :: CallGraph -> String
$cshow :: CallGraph -> String
showsPrec :: Int -> CallGraph -> ShowS
$cshowsPrec :: Int -> CallGraph -> ShowS
Show)

-- | Is the given function known to the call graph?
isFunInCallGraph :: Name -> CallGraph -> Bool
isFunInCallGraph :: Name -> CallGraph -> Bool
isFunInCallGraph Name
f = Name -> FunGraph -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.member Name
f (FunGraph -> Bool) -> (CallGraph -> FunGraph) -> CallGraph -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CallGraph -> FunGraph
cgCalledByFuns

-- | Does the first function call the second?
calls :: Name -> Name -> CallGraph -> Bool
calls :: Name -> Name -> CallGraph -> Bool
calls Name
caller Name
callee =
  Bool -> (FunCalls -> Bool) -> Maybe FunCalls -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (Name -> FunCalls -> Bool
fcCalled Name
callee) (Maybe FunCalls -> Bool)
-> (CallGraph -> Maybe FunCalls) -> CallGraph -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> FunGraph -> Maybe FunCalls
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
caller (FunGraph -> Maybe FunCalls)
-> (CallGraph -> FunGraph) -> CallGraph -> Maybe FunCalls
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CallGraph -> FunGraph
cgCalledByFuns

-- | Is the function called in any of the constants?
calledByConsts :: Name -> CallGraph -> Bool
calledByConsts :: Name -> CallGraph -> Bool
calledByConsts Name
callee = Name -> FunCalls -> Bool
fcCalled Name
callee (FunCalls -> Bool) -> (CallGraph -> FunCalls) -> CallGraph -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CallGraph -> FunCalls
cgCalledByConsts

-- | All functions called by this function.
allCalledBy :: Name -> CallGraph -> S.Set Name
allCalledBy :: Name -> CallGraph -> Set Name
allCalledBy Name
f = Set Name -> (FunCalls -> Set Name) -> Maybe FunCalls -> Set Name
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set Name
forall a. Monoid a => a
mempty FunCalls -> Set Name
fcAllCalled (Maybe FunCalls -> Set Name)
-> (CallGraph -> Maybe FunCalls) -> CallGraph -> Set Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> FunGraph -> Maybe FunCalls
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f (FunGraph -> Maybe FunCalls)
-> (CallGraph -> FunGraph) -> CallGraph -> Maybe FunCalls
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CallGraph -> FunGraph
cgCalledByFuns

-- | @buildCallGraph prog@ build the program's call graph.
buildCallGraph :: Prog SOACS -> CallGraph
buildCallGraph :: Prog SOACS -> CallGraph
buildCallGraph Prog SOACS
prog =
  FunGraph -> FunCalls -> CallGraph
CallGraph FunGraph
fg FunCalls
cg
  where
    fg :: FunGraph
fg = (FunGraph -> Name -> FunGraph) -> FunGraph -> Set Name -> FunGraph
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (FunctionTable -> FunGraph -> Name -> FunGraph
buildFGfun FunctionTable
ftable) FunGraph
forall k a. Map k a
M.empty Set Name
entry_points
    cg :: FunCalls
cg = Stms SOACS -> FunCalls
buildFGStms (Stms SOACS -> FunCalls) -> Stms SOACS -> FunCalls
forall a b. (a -> b) -> a -> b
$ Prog SOACS -> Stms SOACS
forall rep. Prog rep -> Stms rep
progConsts Prog SOACS
prog

    entry_points :: Set Name
entry_points =
      [Name] -> Set Name
forall a. Ord a => [a] -> Set a
S.fromList ((FunDef SOACS -> Name) -> [FunDef SOACS] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map FunDef SOACS -> Name
forall rep. FunDef rep -> Name
funDefName ((FunDef SOACS -> Bool) -> [FunDef SOACS] -> [FunDef SOACS]
forall a. (a -> Bool) -> [a] -> [a]
filter (Maybe EntryPoint -> Bool
forall a. Maybe a -> Bool
isJust (Maybe EntryPoint -> Bool)
-> (FunDef SOACS -> Maybe EntryPoint) -> FunDef SOACS -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FunDef SOACS -> Maybe EntryPoint
forall rep. FunDef rep -> Maybe EntryPoint
funDefEntryPoint) ([FunDef SOACS] -> [FunDef SOACS])
-> [FunDef SOACS] -> [FunDef SOACS]
forall a b. (a -> b) -> a -> b
$ Prog SOACS -> [FunDef SOACS]
forall rep. Prog rep -> [FunDef rep]
progFuns Prog SOACS
prog))
        Set Name -> Set Name -> Set Name
forall a. Semigroup a => a -> a -> a
<> FunCalls -> Set Name
fcAllCalled FunCalls
cg
    ftable :: FunctionTable
ftable = Prog SOACS -> FunctionTable
buildFunctionTable Prog SOACS
prog

count :: Ord k => [k] -> M.Map k Int
count :: [k] -> Map k Int
count [k]
ks = (Int -> Int -> Int) -> [(k, Int)] -> Map k Int
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ([(k, Int)] -> Map k Int) -> [(k, Int)] -> Map k Int
forall a b. (a -> b) -> a -> b
$ [k] -> [Int] -> [(k, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [k]
ks ([Int] -> [(k, Int)]) -> [Int] -> [(k, Int)]
forall a b. (a -> b) -> a -> b
$ Int -> [Int]
forall a. a -> [a]
repeat Int
1

-- | Produce a mapping of the number of occurences in the call graph
-- of each function.  Only counts functions that are called at least
-- once.
numOccurences :: CallGraph -> M.Map Name Int
numOccurences :: CallGraph -> Map Name Int
numOccurences (CallGraph FunGraph
funs FunCalls
consts) =
  [Name] -> Map Name Int
forall k. Ord k => [k] -> Map k Int
count ([Name] -> Map Name Int) -> [Name] -> Map Name Int
forall a b. (a -> b) -> a -> b
$ ((Attrs, Name) -> Name) -> [(Attrs, Name)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Attrs, Name) -> Name
forall a b. (a, b) -> b
snd ([(Attrs, Name)] -> [Name]) -> [(Attrs, Name)] -> [Name]
forall a b. (a -> b) -> a -> b
$ Map CallId (Attrs, Name) -> [(Attrs, Name)]
forall k a. Map k a -> [a]
M.elems (FunCalls -> Map CallId (Attrs, Name)
fcMap FunCalls
consts Map CallId (Attrs, Name)
-> Map CallId (Attrs, Name) -> Map CallId (Attrs, Name)
forall a. Semigroup a => a -> a -> a
<> (FunCalls -> Map CallId (Attrs, Name))
-> [FunCalls] -> Map CallId (Attrs, Name)
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap FunCalls -> Map CallId (Attrs, Name)
fcMap (FunGraph -> [FunCalls]
forall k a. Map k a -> [a]
M.elems FunGraph
funs))

-- | @buildCallGraph ftable fg fname@ updates @fg@ with the
-- contributions of function @fname@.
buildFGfun :: FunctionTable -> FunGraph -> Name -> FunGraph
buildFGfun :: FunctionTable -> FunGraph -> Name -> FunGraph
buildFGfun FunctionTable
ftable FunGraph
fg Name
fname =
  -- Check if function is a non-builtin that we have not already
  -- processed.
  case Name -> FunctionTable -> Maybe (FunDef SOACS)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
fname FunctionTable
ftable of
    Just FunDef SOACS
f | Maybe FunCalls
Nothing <- Name -> FunGraph -> Maybe FunCalls
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
fname FunGraph
fg -> do
      let callees :: FunCalls
callees = Body -> FunCalls
buildFGBody (Body -> FunCalls) -> Body -> FunCalls
forall a b. (a -> b) -> a -> b
$ FunDef SOACS -> Body
forall rep. FunDef rep -> BodyT rep
funDefBody FunDef SOACS
f
          fg' :: FunGraph
fg' = Name -> FunCalls -> FunGraph -> FunGraph
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Name
fname FunCalls
callees FunGraph
fg
      -- recursively build the callees
      (FunGraph -> Name -> FunGraph) -> FunGraph -> Set Name -> FunGraph
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (FunctionTable -> FunGraph -> Name -> FunGraph
buildFGfun FunctionTable
ftable) FunGraph
fg' (Set Name -> FunGraph) -> Set Name -> FunGraph
forall a b. (a -> b) -> a -> b
$ FunCalls -> Set Name
fcAllCalled FunCalls
callees
    Maybe (FunDef SOACS)
_ -> FunGraph
fg

buildFGStms :: Stms SOACS -> FunCalls
buildFGStms :: Stms SOACS -> FunCalls
buildFGStms = [FunCalls] -> FunCalls
forall a. Monoid a => [a] -> a
mconcat ([FunCalls] -> FunCalls)
-> (Stms SOACS -> [FunCalls]) -> Stms SOACS -> FunCalls
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Stm -> FunCalls) -> [Stm] -> [FunCalls]
forall a b. (a -> b) -> [a] -> [b]
map Stm -> FunCalls
buildFGstm ([Stm] -> [FunCalls])
-> (Stms SOACS -> [Stm]) -> Stms SOACS -> [FunCalls]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stms SOACS -> [Stm]
forall rep. Stms rep -> [Stm rep]
stmsToList

buildFGBody :: Body -> FunCalls
buildFGBody :: Body -> FunCalls
buildFGBody = Stms SOACS -> FunCalls
buildFGStms (Stms SOACS -> FunCalls)
-> (Body -> Stms SOACS) -> Body -> FunCalls
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body -> Stms SOACS
forall rep. BodyT rep -> Stms rep
bodyStms

buildFGstm :: Stm -> FunCalls
buildFGstm :: Stm -> FunCalls
buildFGstm (Let (Pat (PatElemT (LetDec SOACS)
p : [PatElemT (LetDec SOACS)]
_)) StmAux (ExpDec SOACS)
aux (Apply Name
fname [(SubExp, Diet)]
_ [RetType SOACS]
_ (Safety, SrcLoc, [SrcLoc])
_)) =
  Map CallId (Attrs, Name) -> Set Name -> FunCalls
FunCalls (CallId -> (Attrs, Name) -> Map CallId (Attrs, Name)
forall k a. k -> a -> Map k a
M.singleton (PatElemT Type -> CallId
forall dec. PatElemT dec -> CallId
patElemName PatElemT Type
PatElemT (LetDec SOACS)
p) (StmAux () -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux ()
StmAux (ExpDec SOACS)
aux, Name
fname)) (Name -> Set Name
forall a. a -> Set a
S.singleton Name
fname)
buildFGstm (Let PatT (LetDec SOACS)
_ StmAux (ExpDec SOACS)
_ (Op Op SOACS
op)) = Writer FunCalls (SOAC SOACS) -> FunCalls
forall w a. Writer w a -> w
execWriter (Writer FunCalls (SOAC SOACS) -> FunCalls)
-> Writer FunCalls (SOAC SOACS) -> FunCalls
forall a b. (a -> b) -> a -> b
$ SOACMapper SOACS SOACS (WriterT FunCalls Identity)
-> SOAC SOACS -> Writer FunCalls (SOAC SOACS)
forall (m :: * -> *) frep trep.
(Applicative m, Monad m) =>
SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
mapSOACM SOACMapper SOACS SOACS (WriterT FunCalls Identity)
folder Op SOACS
SOAC SOACS
op
  where
    folder :: SOACMapper SOACS SOACS (WriterT FunCalls Identity)
folder =
      SOACMapper Any Any (WriterT FunCalls Identity)
forall (m :: * -> *) rep. Monad m => SOACMapper rep rep m
identitySOACMapper
        { mapOnSOACLambda :: Lambda SOACS -> WriterT FunCalls Identity (Lambda SOACS)
mapOnSOACLambda = \Lambda SOACS
lam -> do
            FunCalls -> WriterT FunCalls Identity ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (FunCalls -> WriterT FunCalls Identity ())
-> FunCalls -> WriterT FunCalls Identity ()
forall a b. (a -> b) -> a -> b
$ Body -> FunCalls
buildFGBody (Body -> FunCalls) -> Body -> FunCalls
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body
forall rep. LambdaT rep -> BodyT rep
lambdaBody Lambda SOACS
lam
            Lambda SOACS -> WriterT FunCalls Identity (Lambda SOACS)
forall (m :: * -> *) a. Monad m => a -> m a
return Lambda SOACS
lam
        }
buildFGstm (Let PatT (LetDec SOACS)
_ StmAux (ExpDec SOACS)
_ ExpT SOACS
e) = Writer FunCalls (ExpT SOACS) -> FunCalls
forall w a. Writer w a -> w
execWriter (Writer FunCalls (ExpT SOACS) -> FunCalls)
-> Writer FunCalls (ExpT SOACS) -> FunCalls
forall a b. (a -> b) -> a -> b
$ Mapper SOACS SOACS (WriterT FunCalls Identity)
-> ExpT SOACS -> Writer FunCalls (ExpT SOACS)
forall (m :: * -> *) frep trep.
(Applicative m, Monad m) =>
Mapper frep trep m -> Exp frep -> m (Exp trep)
mapExpM Mapper SOACS SOACS (WriterT FunCalls Identity)
folder ExpT SOACS
e
  where
    folder :: Mapper SOACS SOACS (WriterT FunCalls Identity)
folder =
      Mapper SOACS SOACS (WriterT FunCalls Identity)
forall (m :: * -> *) rep. Monad m => Mapper rep rep m
identityMapper
        { mapOnBody :: Scope SOACS -> Body -> WriterT FunCalls Identity Body
mapOnBody = \Scope SOACS
_ Body
body -> do
            FunCalls -> WriterT FunCalls Identity ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (FunCalls -> WriterT FunCalls Identity ())
-> FunCalls -> WriterT FunCalls Identity ()
forall a b. (a -> b) -> a -> b
$ Body -> FunCalls
buildFGBody Body
body
            Body -> WriterT FunCalls Identity Body
forall (m :: * -> *) a. Monad m => a -> m a
return Body
body
        }

-- | The set of all functions that are called noinline somewhere, or
-- have a noinline attribute on their definition.
findNoninlined :: Prog SOACS -> S.Set Name
findNoninlined :: Prog SOACS -> Set Name
findNoninlined Prog SOACS
prog =
  (FunDef SOACS -> Set Name) -> [FunDef SOACS] -> Set Name
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap FunDef SOACS -> Set Name
forall rep. FunDef rep -> Set Name
noinlineDef (Prog SOACS -> [FunDef SOACS]
forall rep. Prog rep -> [FunDef rep]
progFuns Prog SOACS
prog)
    Set Name -> Set Name -> Set Name
forall a. Semigroup a => a -> a -> a
<> (Stm -> Set Name) -> Stms SOACS -> Set Name
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Stm -> Set Name
onStm ((FunDef SOACS -> Stms SOACS) -> [FunDef SOACS] -> Stms SOACS
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Body -> Stms SOACS
forall rep. BodyT rep -> Stms rep
bodyStms (Body -> Stms SOACS)
-> (FunDef SOACS -> Body) -> FunDef SOACS -> Stms SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FunDef SOACS -> Body
forall rep. FunDef rep -> BodyT rep
funDefBody) (Prog SOACS -> [FunDef SOACS]
forall rep. Prog rep -> [FunDef rep]
progFuns Prog SOACS
prog) Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> Prog SOACS -> Stms SOACS
forall rep. Prog rep -> Stms rep
progConsts Prog SOACS
prog)
  where
    onStm :: Stm -> S.Set Name
    onStm :: Stm -> Set Name
onStm (Let PatT (LetDec SOACS)
_ StmAux (ExpDec SOACS)
aux (Apply Name
fname [(SubExp, Diet)]
_ [RetType SOACS]
_ (Safety, SrcLoc, [SrcLoc])
_))
      | Attr
"noinline" Attr -> Attrs -> Bool
`inAttrs` StmAux () -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux ()
StmAux (ExpDec SOACS)
aux =
        Name -> Set Name
forall a. a -> Set a
S.singleton Name
fname
    onStm (Let PatT (LetDec SOACS)
_ StmAux (ExpDec SOACS)
_ ExpT SOACS
e) = Writer (Set Name) (ExpT SOACS) -> Set Name
forall w a. Writer w a -> w
execWriter (Writer (Set Name) (ExpT SOACS) -> Set Name)
-> Writer (Set Name) (ExpT SOACS) -> Set Name
forall a b. (a -> b) -> a -> b
$ Mapper SOACS SOACS (WriterT (Set Name) Identity)
-> ExpT SOACS -> Writer (Set Name) (ExpT SOACS)
forall (m :: * -> *) frep trep.
(Applicative m, Monad m) =>
Mapper frep trep m -> Exp frep -> m (Exp trep)
mapExpM Mapper SOACS SOACS (WriterT (Set Name) Identity)
folder ExpT SOACS
e
      where
        folder :: Mapper SOACS SOACS (WriterT (Set Name) Identity)
folder =
          Mapper SOACS SOACS (WriterT (Set Name) Identity)
forall (m :: * -> *) rep. Monad m => Mapper rep rep m
identityMapper
            { mapOnBody :: Scope SOACS -> Body -> WriterT (Set Name) Identity Body
mapOnBody = \Scope SOACS
_ Body
body -> do
                Set Name -> WriterT (Set Name) Identity ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (Set Name -> WriterT (Set Name) Identity ())
-> Set Name -> WriterT (Set Name) Identity ()
forall a b. (a -> b) -> a -> b
$ (Stm -> Set Name) -> Stms SOACS -> Set Name
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Stm -> Set Name
onStm (Stms SOACS -> Set Name) -> Stms SOACS -> Set Name
forall a b. (a -> b) -> a -> b
$ Body -> Stms SOACS
forall rep. BodyT rep -> Stms rep
bodyStms Body
body
                Body -> WriterT (Set Name) Identity Body
forall (m :: * -> *) a. Monad m => a -> m a
return Body
body,
              mapOnOp :: Op SOACS -> WriterT (Set Name) Identity (Op SOACS)
mapOnOp =
                SOACMapper SOACS SOACS (WriterT (Set Name) Identity)
-> SOAC SOACS -> WriterT (Set Name) Identity (SOAC SOACS)
forall (m :: * -> *) frep trep.
(Applicative m, Monad m) =>
SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
mapSOACM
                  SOACMapper Any Any (WriterT (Set Name) Identity)
forall (m :: * -> *) rep. Monad m => SOACMapper rep rep m
identitySOACMapper
                    { mapOnSOACLambda :: Lambda SOACS -> WriterT (Set Name) Identity (Lambda SOACS)
mapOnSOACLambda = \Lambda SOACS
lam -> do
                        Set Name -> WriterT (Set Name) Identity ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (Set Name -> WriterT (Set Name) Identity ())
-> Set Name -> WriterT (Set Name) Identity ()
forall a b. (a -> b) -> a -> b
$ (Stm -> Set Name) -> Stms SOACS -> Set Name
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Stm -> Set Name
onStm (Stms SOACS -> Set Name) -> Stms SOACS -> Set Name
forall a b. (a -> b) -> a -> b
$ Body -> Stms SOACS
forall rep. BodyT rep -> Stms rep
bodyStms (Body -> Stms SOACS) -> Body -> Stms SOACS
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body
forall rep. LambdaT rep -> BodyT rep
lambdaBody Lambda SOACS
lam
                        Lambda SOACS -> WriterT (Set Name) Identity (Lambda SOACS)
forall (m :: * -> *) a. Monad m => a -> m a
return Lambda SOACS
lam
                    }
            }

    noinlineDef :: FunDef rep -> Set Name
noinlineDef FunDef rep
fd
      | Attr
"noinline" Attr -> Attrs -> Bool
`inAttrs` FunDef rep -> Attrs
forall rep. FunDef rep -> Attrs
funDefAttrs FunDef rep
fd =
        Name -> Set Name
forall a. a -> Set a
S.singleton (Name -> Set Name) -> Name -> Set Name
forall a b. (a -> b) -> a -> b
$ FunDef rep -> Name
forall rep. FunDef rep -> Name
funDefName FunDef rep
fd
      | Bool
otherwise =
        Set Name
forall a. Monoid a => a
mempty

instance Pretty FunCalls where
  ppr :: FunCalls -> Doc
ppr = [Doc] -> Doc
stack ([Doc] -> Doc) -> (FunCalls -> [Doc]) -> FunCalls -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((CallId, (Attrs, Name)) -> Doc)
-> [(CallId, (Attrs, Name))] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map (CallId, (Attrs, Name)) -> Doc
forall a a a. (Pretty a, Pretty a, Pretty a) => (a, (a, a)) -> Doc
f ([(CallId, (Attrs, Name))] -> [Doc])
-> (FunCalls -> [(CallId, (Attrs, Name))]) -> FunCalls -> [Doc]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map CallId (Attrs, Name) -> [(CallId, (Attrs, Name))]
forall k a. Map k a -> [(k, a)]
M.toList (Map CallId (Attrs, Name) -> [(CallId, (Attrs, Name))])
-> (FunCalls -> Map CallId (Attrs, Name))
-> FunCalls
-> [(CallId, (Attrs, Name))]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FunCalls -> Map CallId (Attrs, Name)
fcMap
    where
      f :: (a, (a, a)) -> Doc
f (a
x, (a
attrs, a
y)) = Doc
"=>" Doc -> Doc -> Doc
<+> a -> Doc
forall a. Pretty a => a -> Doc
ppr a
y Doc -> Doc -> Doc
<+> Doc -> Doc
parens (Doc
"at" Doc -> Doc -> Doc
<+> a -> Doc
forall a. Pretty a => a -> Doc
ppr a
x Doc -> Doc -> Doc
<+> a -> Doc
forall a. Pretty a => a -> Doc
ppr a
attrs)

instance Pretty CallGraph where
  ppr :: CallGraph -> Doc
ppr (CallGraph FunGraph
fg FunCalls
cg) =
    [Doc] -> Doc
stack ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$
      Doc -> [Doc] -> [Doc]
punctuate Doc
line ([Doc] -> [Doc]) -> [Doc] -> [Doc]
forall a b. (a -> b) -> a -> b
$
        (Name, FunCalls) -> Doc
forall a. Pretty a => (Name, a) -> Doc
ppFunCalls (Name
"called at top level", FunCalls
cg) Doc -> [Doc] -> [Doc]
forall a. a -> [a] -> [a]
: ((Name, FunCalls) -> Doc) -> [(Name, FunCalls)] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map (Name, FunCalls) -> Doc
forall a. Pretty a => (Name, a) -> Doc
ppFunCalls (FunGraph -> [(Name, FunCalls)]
forall k a. Map k a -> [(k, a)]
M.toList FunGraph
fg)
    where
      ppFunCalls :: (Name, a) -> Doc
ppFunCalls (Name
f, a
fcalls) =
        Name -> Doc
forall a. Pretty a => a -> Doc
ppr Name
f Doc -> Doc -> Doc
</> String -> Doc
text ((Char -> Char) -> ShowS
forall a b. (a -> b) -> [a] -> [b]
map (Char -> Char -> Char
forall a b. a -> b -> a
const Char
'=') (Name -> String
nameToString Name
f))
          Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 (a -> Doc
forall a. Pretty a => a -> Doc
ppr a
fcalls)