{-# 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 Data.List (foldl')
import qualified Data.Map.Strict as M
import Data.Maybe (fromMaybe)
import qualified Data.Set as S
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, 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 {lore}. FunDef lore -> Set Name
noinlineDef (Prog SOACS -> [FunDef SOACS]
forall lore. Prog lore -> [FunDef lore]
progFuns Prog SOACS
prog)
    Set Name -> Set Name -> Set Name
forall a. Semigroup a => a -> a -> a
<> (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 ((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
                    }
            }

    noinlineDef :: FunDef lore -> Set Name
noinlineDef FunDef lore
fd
      | Attr
"noinline" Attr -> Attrs -> Bool
`inAttrs` FunDef lore -> Attrs
forall lore. FunDef lore -> Attrs
funDefAttrs FunDef lore
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 lore -> Name
forall lore. FunDef lore -> Name
funDefName FunDef lore
fd
      | Bool
otherwise =
        Set Name
forall a. Monoid a => a
mempty