{-# 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
  , findNoninlined
  )
  where

import Control.Monad.Writer.Strict
import qualified Data.Map.Strict as M
import qualified Data.Set as S
import Data.Maybe (fromMaybe)
import Data.List (foldl')

import Futhark.IR.SOACS

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 lore.
Map Name (FunDef lore) -> FunDef lore -> Map Name (FunDef lore)
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 lore. Prog lore -> [FunDef lore]
progFuns
  where expand :: Map Name (FunDef lore) -> FunDef lore -> Map Name (FunDef lore)
expand Map Name (FunDef lore)
ftab FunDef lore
f = Name
-> FunDef lore -> Map Name (FunDef lore) -> Map Name (FunDef lore)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (FunDef lore -> Name
forall lore. FunDef lore -> Name
funDefName FunDef lore
f) FunDef lore
f Map Name (FunDef lore)
ftab

type FunGraph = M.Map Name (S.Set Name)

-- | The call graph is a mapping from a function name, i.e., the
-- caller, to a set 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 -> Map Name (Set Name)
calledByFuns :: M.Map Name (S.Set Name)
                           , CallGraph -> Set Name
calledInConsts :: S.Set Name
                           }

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

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

-- | Is the function called in any of the constants?
calledByConsts :: Name -> CallGraph -> Bool
calledByConsts :: Name -> CallGraph -> Bool
calledByConsts Name
f = Name -> Set Name -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Name
f (Set Name -> Bool) -> (CallGraph -> Set Name) -> CallGraph -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CallGraph -> Set Name
calledInConsts

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

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

        entry_points :: [Name]
entry_points = (FunDef SOACS -> Name) -> [FunDef SOACS] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map FunDef SOACS -> Name
forall lore. FunDef lore -> Name
funDefName ([FunDef SOACS] -> [Name]) -> [FunDef SOACS] -> [Name]
forall a b. (a -> b) -> a -> b
$ Prog SOACS -> [FunDef SOACS]
forall lore. Prog lore -> [FunDef lore]
progFuns Prog SOACS
prog
        ftable :: FunctionTable
ftable = Prog SOACS -> FunctionTable
buildFunctionTable Prog SOACS
prog

-- | @buildCallGraph ftable fg fname@ updates @fg@ with the
-- contributions of function @fname@.
buildFGfun :: FunctionTable -> FunGraph -> Name -> FunGraph
buildFGfun :: FunctionTable -> Map Name (Set Name) -> Name -> Map Name (Set Name)
buildFGfun FunctionTable
ftable Map Name (Set Name)
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 (Set Name)
Nothing <- Name -> Map Name (Set Name) -> Maybe (Set Name)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
fname Map Name (Set Name)
fg -> do
               let callees :: Set Name
callees = Body -> Set Name
buildFGBody (Body -> Set Name) -> Body -> Set Name
forall a b. (a -> b) -> a -> b
$ FunDef SOACS -> Body
forall lore. FunDef lore -> BodyT lore
funDefBody FunDef SOACS
f
                   fg' :: Map Name (Set Name)
fg' = Name -> Set Name -> Map Name (Set Name) -> Map Name (Set Name)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Name
fname Set Name
callees Map Name (Set Name)
fg
               -- recursively build the callees
               (Map Name (Set Name) -> Name -> Map Name (Set Name))
-> Map Name (Set Name) -> Set Name -> Map Name (Set Name)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (FunctionTable -> Map Name (Set Name) -> Name -> Map Name (Set Name)
buildFGfun FunctionTable
ftable) Map Name (Set Name)
fg' Set Name
callees
    Maybe (FunDef SOACS)
_ -> Map Name (Set Name)
fg

buildFGStms :: Stms SOACS -> S.Set Name
buildFGStms :: Stms SOACS -> Set Name
buildFGStms = [Set Name] -> Set Name
forall a. Monoid a => [a] -> a
mconcat ([Set Name] -> Set Name)
-> (Stms SOACS -> [Set Name]) -> Stms SOACS -> Set Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Stm SOACS -> Set Name) -> [Stm SOACS] -> [Set Name]
forall a b. (a -> b) -> [a] -> [b]
map (Exp -> Set Name
buildFGexp (Exp -> Set Name) -> (Stm SOACS -> Exp) -> Stm SOACS -> Set Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm SOACS -> Exp
forall lore. Stm lore -> Exp lore
stmExp) ([Stm SOACS] -> [Set Name])
-> (Stms SOACS -> [Stm SOACS]) -> Stms SOACS -> [Set Name]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stms SOACS -> [Stm SOACS]
forall lore. Stms lore -> [Stm lore]
stmsToList

buildFGBody :: Body -> S.Set Name
buildFGBody :: Body -> Set Name
buildFGBody = Stms SOACS -> Set Name
buildFGStms (Stms SOACS -> Set Name)
-> (Body -> Stms SOACS) -> Body -> Set Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body -> Stms SOACS
forall lore. BodyT lore -> Stms lore
bodyStms

buildFGexp :: Exp -> S.Set Name
buildFGexp :: Exp -> Set Name
buildFGexp (Apply Name
fname [(SubExp, Diet)]
_ [RetType SOACS]
_ (Safety, SrcLoc, [SrcLoc])
_) = Name -> Set Name
forall a. a -> Set a
S.singleton Name
fname
buildFGexp (Op Op SOACS
op) = Writer (Set Name) (SOAC SOACS) -> Set Name
forall w a. Writer w a -> w
execWriter (Writer (Set Name) (SOAC SOACS) -> Set Name)
-> Writer (Set Name) (SOAC SOACS) -> Set Name
forall a b. (a -> b) -> a -> b
$ SOACMapper SOACS SOACS (WriterT (Set Name) Identity)
-> SOAC SOACS -> Writer (Set Name) (SOAC SOACS)
forall (m :: * -> *) flore tlore.
(Applicative m, Monad m) =>
SOACMapper flore tlore m -> SOAC flore -> m (SOAC tlore)
mapSOACM SOACMapper SOACS SOACS (WriterT (Set Name) Identity)
folder Op SOACS
SOAC SOACS
op
  where folder :: SOACMapper SOACS SOACS (WriterT (Set Name) Identity)
folder = SOACMapper Any Any (WriterT (Set Name) Identity)
forall (m :: * -> *) lore. Monad m => SOACMapper lore lore 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
$ Body -> Set Name
buildFGBody (Body -> Set Name) -> Body -> Set Name
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda SOACS
lam
                                       Lambda SOACS -> WriterT (Set Name) Identity (Lambda SOACS)
forall (m :: * -> *) a. Monad m => a -> m a
return Lambda SOACS
lam
          }
buildFGexp Exp
e = Writer (Set Name) Exp -> Set Name
forall w a. Writer w a -> w
execWriter (Writer (Set Name) Exp -> Set Name)
-> Writer (Set Name) Exp -> Set Name
forall a b. (a -> b) -> a -> b
$ Mapper SOACS SOACS (WriterT (Set Name) Identity)
-> Exp -> Writer (Set Name) Exp
forall (m :: * -> *) flore tlore.
(Applicative m, Monad m) =>
Mapper flore tlore m -> Exp flore -> m (Exp tlore)
mapExpM Mapper SOACS SOACS (WriterT (Set Name) Identity)
folder Exp
e
  where folder :: Mapper SOACS SOACS (WriterT (Set Name) Identity)
folder = Mapper SOACS SOACS (WriterT (Set Name) Identity)
forall (m :: * -> *) lore. Monad m => Mapper lore lore 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
$ Body -> Set Name
buildFGBody Body
body
                                    Body -> WriterT (Set Name) Identity Body
forall (m :: * -> *) a. Monad m => a -> m a
return Body
body
          }

-- | The set of all functions that are called noinline somewhere.
findNoninlined :: Prog SOACS -> S.Set Name
findNoninlined :: Prog SOACS -> Set Name
findNoninlined Prog SOACS
prog =
  (Stm SOACS -> Set Name) -> Stms SOACS -> Set Name
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Stm SOACS -> Set Name
onStm (Stms SOACS -> Set Name) -> Stms SOACS -> Set Name
forall a b. (a -> b) -> a -> b
$ (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 lore. BodyT lore -> Stms lore
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 lore. FunDef lore -> BodyT lore
funDefBody) (Prog SOACS -> [FunDef SOACS]
forall lore. Prog lore -> [FunDef lore]
progFuns Prog SOACS
prog) Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> Prog SOACS -> Stms SOACS
forall lore. Prog lore -> Stms lore
progConsts Prog SOACS
prog
  where onStm :: Stm -> S.Set Name
        onStm :: Stm SOACS -> Set Name
onStm (Let Pattern 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 Pattern SOACS
_ StmAux (ExpDec SOACS)
_ Exp
e) = Writer (Set Name) Exp -> Set Name
forall w a. Writer w a -> w
execWriter (Writer (Set Name) Exp -> Set Name)
-> Writer (Set Name) Exp -> Set Name
forall a b. (a -> b) -> a -> b
$ Mapper SOACS SOACS (WriterT (Set Name) Identity)
-> Exp -> Writer (Set Name) Exp
forall (m :: * -> *) flore tlore.
(Applicative m, Monad m) =>
Mapper flore tlore m -> Exp flore -> m (Exp tlore)
mapExpM Mapper SOACS SOACS (WriterT (Set Name) Identity)
folder Exp
e
          where folder :: Mapper SOACS SOACS (WriterT (Set Name) Identity)
folder =
                  Mapper SOACS SOACS (WriterT (Set Name) Identity)
forall (m :: * -> *) lore. Monad m => Mapper lore lore 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 SOACS -> Set Name) -> Stms SOACS -> Set Name
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Stm SOACS -> Set Name
onStm (Stms SOACS -> Set Name) -> Stms SOACS -> Set Name
forall a b. (a -> b) -> a -> b
$ Body -> Stms SOACS
forall lore. BodyT lore -> Stms lore
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 -> Writer (Set Name) (SOAC SOACS)
forall (m :: * -> *) flore tlore.
(Applicative m, Monad m) =>
SOACMapper flore tlore m -> SOAC flore -> m (SOAC tlore)
mapSOACM
                              SOACMapper Any Any (WriterT (Set Name) Identity)
forall (m :: * -> *) lore. Monad m => SOACMapper lore lore 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 SOACS -> Set Name) -> Stms SOACS -> Set Name
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Stm SOACS -> Set Name
onStm (Stms SOACS -> Set Name) -> Stms SOACS -> Set Name
forall a b. (a -> b) -> a -> b
$ Body -> Stms SOACS
forall lore. BodyT lore -> Stms lore
bodyStms (Body -> Stms SOACS) -> Body -> Stms SOACS
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda SOACS
lam
                                  Lambda SOACS -> WriterT (Set Name) Identity (Lambda SOACS)
forall (m :: * -> *) a. Monad m => a -> m a
return Lambda SOACS
lam
                              }
                  }