{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}

-- | This module implements a compiler pass for inlining functions,
-- then removing those that have become dead.
module Futhark.Optimise.InliningDeadFun
  ( inlineAggressively,
    inlineConservatively,
    removeDeadFunctions,
  )
where

import Control.Monad.Identity
import Control.Monad.State
import Control.Parallel.Strategies
import Data.List (partition)
import qualified Data.Map.Strict as M
import qualified Data.Set as S
import Futhark.Analysis.CallGraph
import qualified Futhark.Analysis.SymbolTable as ST
import Futhark.Builder
import Futhark.IR.SOACS
import Futhark.IR.SOACS.Simplify
  ( simpleSOACS,
    simplifyConsts,
    simplifyFun,
  )
import Futhark.Optimise.CSE
import Futhark.Optimise.Simplify.Rep (addScopeWisdom, informStms)
import Futhark.Pass
import Futhark.Transform.CopyPropagate
  ( copyPropagateInFun,
    copyPropagateInProg,
  )
import Futhark.Transform.Rename

parMapM :: MonadFreshNames m => (a -> State VNameSource b) -> [a] -> m [b]
-- The special-casing of [] is quite important here!  If 'as' is
-- empty, then we might otherwise create an empty name source below,
-- which can wreak all kinds of havoc.
parMapM :: (a -> State VNameSource b) -> [a] -> m [b]
parMapM a -> State VNameSource b
_ [] = [b] -> m [b]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
parMapM a -> State VNameSource b
f [a]
as =
  (VNameSource -> ([b], VNameSource)) -> m [b]
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> ([b], VNameSource)) -> m [b])
-> (VNameSource -> ([b], VNameSource)) -> m [b]
forall a b. (a -> b) -> a -> b
$ \VNameSource
src ->
    let f' :: a -> (b, VNameSource)
f' a
a = State VNameSource b -> VNameSource -> (b, VNameSource)
forall s a. State s a -> s -> (a, s)
runState (a -> State VNameSource b
f a
a) VNameSource
src
        ([b]
bs, [VNameSource]
srcs) = [(b, VNameSource)] -> ([b], [VNameSource])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(b, VNameSource)] -> ([b], [VNameSource]))
-> [(b, VNameSource)] -> ([b], [VNameSource])
forall a b. (a -> b) -> a -> b
$ Strategy (b, VNameSource)
-> (a -> (b, VNameSource)) -> [a] -> [(b, VNameSource)]
forall b a. Strategy b -> (a -> b) -> [a] -> [b]
parMap Strategy (b, VNameSource)
forall a. Strategy a
rpar a -> (b, VNameSource)
f' [a]
as
     in ([b]
bs, [VNameSource] -> VNameSource
forall a. Monoid a => [a] -> a
mconcat [VNameSource]
srcs)

-- It is more efficient to shrink the program as soon as possible,
-- rather than wait until it has balooned after full inlining.  This
-- is the inverse rate at which we perform full simplification after
-- inlining.  For the other steps we just do copy propagation.  The
-- simplification rates used have been determined heuristically and
-- are probably not optimal for any given program.
inlineFunctions ::
  MonadFreshNames m =>
  Int ->
  CallGraph ->
  S.Set Name ->
  Prog SOACS ->
  m (Prog SOACS)
inlineFunctions :: Int -> CallGraph -> Set Name -> Prog SOACS -> m (Prog SOACS)
inlineFunctions Int
simplify_rate CallGraph
cg Set Name
what_should_be_inlined Prog SOACS
prog = do
  let Prog Stms SOACS
consts [FunDef SOACS]
funs = Prog SOACS
prog
      vtable :: SymbolTable (Wise SOACS)
vtable = Scope (Wise SOACS) -> SymbolTable (Wise SOACS)
forall rep. ASTRep rep => Scope rep -> SymbolTable rep
ST.fromScope (Scope SOACS -> Scope (Wise SOACS)
forall rep. Scope rep -> Scope (Wise rep)
addScopeWisdom (Stms SOACS -> Scope SOACS
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stms SOACS
consts))

  (Stms SOACS -> [FunDef SOACS] -> Prog SOACS)
-> (Stms SOACS, [FunDef SOACS]) -> Prog SOACS
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Stms SOACS -> [FunDef SOACS] -> Prog SOACS
forall rep. Stms rep -> [FunDef rep] -> Prog rep
Prog ((Stms SOACS, [FunDef SOACS]) -> Prog SOACS)
-> m (Stms SOACS, [FunDef SOACS]) -> m (Prog SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int, SymbolTable (Wise SOACS))
-> (Stms SOACS, [FunDef SOACS])
-> Set Name
-> m (Stms SOACS, [FunDef SOACS])
forall (f :: * -> *).
MonadFreshNames f =>
(Int, SymbolTable (Wise SOACS))
-> (Stms SOACS, [FunDef SOACS])
-> Set Name
-> f (Stms SOACS, [FunDef SOACS])
recurse (Int
1, SymbolTable (Wise SOACS)
vtable) (Stms SOACS
consts, [FunDef SOACS]
funs) Set Name
what_should_be_inlined
  where
    fdmap :: [FunDef rep] -> Map Name (FunDef rep)
fdmap [FunDef rep]
fds = [(Name, FunDef rep)] -> Map Name (FunDef rep)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, FunDef rep)] -> Map Name (FunDef rep))
-> [(Name, FunDef rep)] -> Map Name (FunDef rep)
forall a b. (a -> b) -> a -> b
$ [Name] -> [FunDef rep] -> [(Name, FunDef rep)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((FunDef rep -> Name) -> [FunDef rep] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map FunDef rep -> Name
forall rep. FunDef rep -> Name
funDefName [FunDef rep]
fds) [FunDef rep]
fds

    noCallsTo :: Set Name -> Name -> Bool
noCallsTo Set Name
which Name
from = Set Name -> Bool
forall a. Set a -> Bool
S.null (Set Name -> Bool) -> Set Name -> Bool
forall a b. (a -> b) -> a -> b
$ Name -> CallGraph -> Set Name
allCalledBy Name
from CallGraph
cg Set Name -> Set Name -> Set Name
forall a. Ord a => Set a -> Set a -> Set a
`S.intersection` Set Name
which

    recurse :: (Int, SymbolTable (Wise SOACS))
-> (Stms SOACS, [FunDef SOACS])
-> Set Name
-> f (Stms SOACS, [FunDef SOACS])
recurse (Int
i, SymbolTable (Wise SOACS)
vtable) (Stms SOACS
consts, [FunDef SOACS]
funs) Set Name
to_inline = do
      let (Set Name
to_inline_now, Set Name
to_inline_later) =
            (Name -> Bool) -> Set Name -> (Set Name, Set Name)
forall a. (a -> Bool) -> Set a -> (Set a, Set a)
S.partition (Set Name -> Name -> Bool
noCallsTo Set Name
to_inline) Set Name
to_inline
          ([FunDef SOACS]
dont_inline_in, [FunDef SOACS]
to_inline_in) =
            (FunDef SOACS -> Bool)
-> [FunDef SOACS] -> ([FunDef SOACS], [FunDef SOACS])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition (Set Name -> Name -> Bool
noCallsTo Set Name
to_inline_now (Name -> Bool) -> (FunDef SOACS -> Name) -> FunDef SOACS -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FunDef SOACS -> Name
forall rep. FunDef rep -> Name
funDefName) [FunDef SOACS]
funs

      if Set Name -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Set Name
to_inline_now
        then (Stms SOACS, [FunDef SOACS]) -> f (Stms SOACS, [FunDef SOACS])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms SOACS
consts, [FunDef SOACS]
funs)
        else do
          let inlinemap :: Map Name (FunDef SOACS)
inlinemap =
                [FunDef SOACS] -> Map Name (FunDef SOACS)
forall rep. [FunDef rep] -> Map Name (FunDef rep)
fdmap ([FunDef SOACS] -> Map Name (FunDef SOACS))
-> [FunDef SOACS] -> Map Name (FunDef SOACS)
forall a b. (a -> b) -> a -> b
$ (FunDef SOACS -> Bool) -> [FunDef SOACS] -> [FunDef SOACS]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Name -> Set Name -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set Name
to_inline_now) (Name -> Bool) -> (FunDef SOACS -> Name) -> FunDef SOACS -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FunDef SOACS -> Name
forall rep. FunDef rep -> Name
funDefName) [FunDef SOACS]
dont_inline_in
          (SymbolTable (Wise SOACS)
vtable', Stms SOACS
consts') <-
            if (Name -> Bool) -> Set Name -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Name -> CallGraph -> Bool
`calledByConsts` CallGraph
cg) Set Name
to_inline_now
              then do
                Stms SOACS
consts' <-
                  Stms SOACS -> f (Stms SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Stms SOACS -> m (Stms SOACS)
simplifyConsts (Stms SOACS -> f (Stms SOACS))
-> (Stms SOACS -> Stms SOACS) -> Stms SOACS -> f (Stms SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Stms SOACS -> Stms SOACS
forall rep.
(ASTRep rep, CanBeAliased (Op rep),
 CSEInOp (OpWithAliases (Op rep))) =>
Bool -> Stms rep -> Stms rep
performCSEOnStms Bool
True
                    (Stms SOACS -> f (Stms SOACS)) -> f (Stms SOACS) -> f (Stms SOACS)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Map Name (FunDef SOACS) -> Stms SOACS -> f (Stms SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Map Name (FunDef SOACS) -> Stms SOACS -> m (Stms SOACS)
inlineInStms Map Name (FunDef SOACS)
inlinemap Stms SOACS
consts
                (SymbolTable (Wise SOACS), Stms SOACS)
-> f (SymbolTable (Wise SOACS), Stms SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms (Wise SOACS)
-> SymbolTable (Wise SOACS) -> SymbolTable (Wise SOACS)
forall rep.
(ASTRep rep, IndexOp (Op rep), Aliased rep) =>
Stms rep -> SymbolTable rep -> SymbolTable rep
ST.insertStms (Stms SOACS -> Stms (Wise SOACS)
forall rep. Informing rep => Stms rep -> Stms (Wise rep)
informStms Stms SOACS
consts') SymbolTable (Wise SOACS)
forall a. Monoid a => a
mempty, Stms SOACS
consts')
              else (SymbolTable (Wise SOACS), Stms SOACS)
-> f (SymbolTable (Wise SOACS), Stms SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SymbolTable (Wise SOACS)
vtable, Stms SOACS
consts)

          let simplifyFun' :: FunDef SOACS -> m (FunDef SOACS)
simplifyFun' FunDef SOACS
fd
                | Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`rem` Int
simplify_rate Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 =
                    SimpleOps SOACS
-> SymbolTable (Wise SOACS) -> FunDef SOACS -> m (FunDef SOACS)
forall (m :: * -> *) rep.
(MonadFreshNames m, SimplifiableRep rep) =>
SimpleOps rep
-> SymbolTable (Wise rep) -> FunDef rep -> m (FunDef rep)
copyPropagateInFun SimpleOps SOACS
simpleSOACS SymbolTable (Wise SOACS)
vtable'
                      (FunDef SOACS -> m (FunDef SOACS))
-> (FunDef SOACS -> FunDef SOACS)
-> FunDef SOACS
-> m (FunDef SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> FunDef SOACS -> FunDef SOACS
forall rep.
(ASTRep rep, CanBeAliased (Op rep),
 CSEInOp (OpWithAliases (Op rep))) =>
Bool -> FunDef rep -> FunDef rep
performCSEOnFunDef Bool
True
                      (FunDef SOACS -> m (FunDef SOACS))
-> m (FunDef SOACS) -> m (FunDef SOACS)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SymbolTable (Wise SOACS) -> FunDef SOACS -> m (FunDef SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
SymbolTable (Wise SOACS) -> FunDef SOACS -> m (FunDef SOACS)
simplifyFun SymbolTable (Wise SOACS)
vtable' FunDef SOACS
fd
                | Bool
otherwise =
                    SimpleOps SOACS
-> SymbolTable (Wise SOACS) -> FunDef SOACS -> m (FunDef SOACS)
forall (m :: * -> *) rep.
(MonadFreshNames m, SimplifiableRep rep) =>
SimpleOps rep
-> SymbolTable (Wise rep) -> FunDef rep -> m (FunDef rep)
copyPropagateInFun SimpleOps SOACS
simpleSOACS SymbolTable (Wise SOACS)
vtable' FunDef SOACS
fd

              onFun :: FunDef SOACS -> StateT VNameSource Identity (FunDef SOACS)
onFun = FunDef SOACS -> StateT VNameSource Identity (FunDef SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
FunDef SOACS -> m (FunDef SOACS)
simplifyFun' (FunDef SOACS -> StateT VNameSource Identity (FunDef SOACS))
-> (FunDef SOACS -> StateT VNameSource Identity (FunDef SOACS))
-> FunDef SOACS
-> StateT VNameSource Identity (FunDef SOACS)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Map Name (FunDef SOACS)
-> FunDef SOACS -> StateT VNameSource Identity (FunDef SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Map Name (FunDef SOACS) -> FunDef SOACS -> m (FunDef SOACS)
inlineInFunDef Map Name (FunDef SOACS)
inlinemap

          [FunDef SOACS]
to_inline_in' <- (FunDef SOACS -> StateT VNameSource Identity (FunDef SOACS))
-> [FunDef SOACS] -> f [FunDef SOACS]
forall (m :: * -> *) a b.
MonadFreshNames m =>
(a -> State VNameSource b) -> [a] -> m [b]
parMapM FunDef SOACS -> StateT VNameSource Identity (FunDef SOACS)
onFun [FunDef SOACS]
to_inline_in

          (Int, SymbolTable (Wise SOACS))
-> (Stms SOACS, [FunDef SOACS])
-> Set Name
-> f (Stms SOACS, [FunDef SOACS])
recurse
            (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, SymbolTable (Wise SOACS)
vtable')
            (Stms SOACS
consts', [FunDef SOACS]
dont_inline_in [FunDef SOACS] -> [FunDef SOACS] -> [FunDef SOACS]
forall a. Semigroup a => a -> a -> a
<> [FunDef SOACS]
to_inline_in')
            Set Name
to_inline_later

calledOnce :: CallGraph -> S.Set Name
calledOnce :: CallGraph -> Set Name
calledOnce = [Name] -> Set Name
forall a. Ord a => [a] -> Set a
S.fromList ([Name] -> Set Name)
-> (CallGraph -> [Name]) -> CallGraph -> Set Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Name, Int) -> Name) -> [(Name, Int)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, Int) -> Name
forall a b. (a, b) -> a
fst ([(Name, Int)] -> [Name])
-> (CallGraph -> [(Name, Int)]) -> CallGraph -> [Name]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Name, Int) -> Bool) -> [(Name, Int)] -> [(Name, Int)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1) (Int -> Bool) -> ((Name, Int) -> Int) -> (Name, Int) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, Int) -> Int
forall a b. (a, b) -> b
snd) ([(Name, Int)] -> [(Name, Int)])
-> (CallGraph -> [(Name, Int)]) -> CallGraph -> [(Name, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name Int -> [(Name, Int)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Name Int -> [(Name, Int)])
-> (CallGraph -> Map Name Int) -> CallGraph -> [(Name, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CallGraph -> Map Name Int
numOccurences

inlineBecauseTiny :: Prog SOACS -> S.Set Name
inlineBecauseTiny :: Prog SOACS -> Set Name
inlineBecauseTiny = (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
onFunDef ([FunDef SOACS] -> Set Name)
-> (Prog SOACS -> [FunDef SOACS]) -> Prog SOACS -> Set Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog SOACS -> [FunDef SOACS]
forall rep. Prog rep -> [FunDef rep]
progFuns
  where
    onFunDef :: FunDef rep -> Set Name
onFunDef FunDef rep
fd
      | Seq (Stm rep) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Body rep -> Seq (Stm rep)
forall rep. Body rep -> Stms rep
bodyStms (FunDef rep -> Body rep
forall rep. FunDef rep -> Body rep
funDefBody FunDef rep
fd)) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2
          Bool -> Bool -> Bool
|| Attr
"inline" 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 (FunDef rep -> Name
forall rep. FunDef rep -> Name
funDefName FunDef rep
fd)
      | Bool
otherwise = Set Name
forall a. Monoid a => a
mempty

-- Conservative inlining of functions that are called just once, or
-- have #[inline] on them.
consInlineFunctions :: MonadFreshNames m => Prog SOACS -> m (Prog SOACS)
consInlineFunctions :: Prog SOACS -> m (Prog SOACS)
consInlineFunctions Prog SOACS
prog =
  Int -> CallGraph -> Set Name -> Prog SOACS -> m (Prog SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Int -> CallGraph -> Set Name -> Prog SOACS -> m (Prog SOACS)
inlineFunctions Int
4 CallGraph
cg (CallGraph -> Set Name
calledOnce CallGraph
cg Set Name -> Set Name -> Set Name
forall a. Semigroup a => a -> a -> a
<> Prog SOACS -> Set Name
inlineBecauseTiny Prog SOACS
prog) Prog SOACS
prog
  where
    cg :: CallGraph
cg = Prog SOACS -> CallGraph
buildCallGraph Prog SOACS
prog

aggInlineFunctions :: MonadFreshNames m => Prog SOACS -> m (Prog SOACS)
aggInlineFunctions :: Prog SOACS -> m (Prog SOACS)
aggInlineFunctions Prog SOACS
prog =
  Int -> CallGraph -> Set Name -> Prog SOACS -> m (Prog SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Int -> CallGraph -> Set Name -> Prog SOACS -> m (Prog SOACS)
inlineFunctions Int
3 CallGraph
cg ([Name] -> Set Name
forall a. Ord a => [a] -> Set a
S.fromList ([Name] -> Set Name) -> [Name] -> Set Name
forall a b. (a -> b) -> a -> b
$ (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] -> [Name]) -> [FunDef SOACS] -> [Name]
forall a b. (a -> b) -> a -> b
$ Prog SOACS -> [FunDef SOACS]
forall rep. Prog rep -> [FunDef rep]
progFuns Prog SOACS
prog) Prog SOACS
prog
  where
    cg :: CallGraph
cg = Prog SOACS -> CallGraph
buildCallGraph Prog SOACS
prog

-- | @inlineInFunDef constf fdmap caller@ inlines in @calleer@ the
-- functions in @fdmap@ that are called as @constf@. At this point the
-- preconditions are that if @fdmap@ is not empty, and, more
-- importantly, the functions in @fdmap@ do not call any other
-- functions.
inlineInFunDef ::
  MonadFreshNames m =>
  M.Map Name (FunDef SOACS) ->
  FunDef SOACS ->
  m (FunDef SOACS)
inlineInFunDef :: Map Name (FunDef SOACS) -> FunDef SOACS -> m (FunDef SOACS)
inlineInFunDef Map Name (FunDef SOACS)
fdmap (FunDef Maybe EntryPoint
entry Attrs
attrs Name
name [RetType SOACS]
rtp [FParam SOACS]
args Body SOACS
body) =
  Maybe EntryPoint
-> Attrs
-> Name
-> [RetType SOACS]
-> [FParam SOACS]
-> Body SOACS
-> FunDef SOACS
forall rep.
Maybe EntryPoint
-> Attrs
-> Name
-> [RetType rep]
-> [FParam rep]
-> Body rep
-> FunDef rep
FunDef Maybe EntryPoint
entry Attrs
attrs Name
name [RetType SOACS]
rtp [FParam SOACS]
args (Body SOACS -> FunDef SOACS) -> m (Body SOACS) -> m (FunDef SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Name (FunDef SOACS) -> Body SOACS -> m (Body SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Map Name (FunDef SOACS) -> Body SOACS -> m (Body SOACS)
inlineInBody Map Name (FunDef SOACS)
fdmap Body SOACS
body

inlineFunction ::
  MonadFreshNames m =>
  Pat Type ->
  StmAux dec ->
  [(SubExp, Diet)] ->
  (Safety, SrcLoc, [SrcLoc]) ->
  FunDef SOACS ->
  m (Stms SOACS)
inlineFunction :: Pat Type
-> StmAux dec
-> [(SubExp, Diet)]
-> (Safety, SrcLoc, [SrcLoc])
-> FunDef SOACS
-> m (Stms SOACS)
inlineFunction Pat Type
pat StmAux dec
aux [(SubExp, Diet)]
args (Safety
safety, SrcLoc
loc, [SrcLoc]
locs) FunDef SOACS
fun = do
  Body BodyDec SOACS
_ Stms SOACS
stms Result
res <-
    Body SOACS -> m (Body SOACS)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Body rep -> m (Body rep)
renameBody (Body SOACS -> m (Body SOACS)) -> Body SOACS -> m (Body SOACS)
forall a b. (a -> b) -> a -> b
$ Stms SOACS -> Result -> Body SOACS
forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody (Stms SOACS
param_stms Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> Stms SOACS
body_stms) (Body SOACS -> Result
forall rep. Body rep -> Result
bodyResult (FunDef SOACS -> Body SOACS
forall rep. FunDef rep -> Body rep
funDefBody FunDef SOACS
fun))
  Stms SOACS -> m (Stms SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms SOACS -> m (Stms SOACS)) -> Stms SOACS -> m (Stms SOACS)
forall a b. (a -> b) -> a -> b
$ Stms SOACS
stms Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> [Stm SOACS] -> Stms SOACS
forall rep. [Stm rep] -> Stms rep
stmsFromList ((Ident -> SubExpRes -> Stm SOACS)
-> [Ident] -> Result -> [Stm SOACS]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Ident -> SubExpRes -> Stm SOACS
forall rep. Buildable rep => Ident -> SubExpRes -> Stm rep
bindSubExpRes (Pat Type -> [Ident]
forall dec. Typed dec => Pat dec -> [Ident]
patIdents Pat Type
pat) Result
res)
  where
    param_stms :: Stms SOACS
param_stms =
      [Stm SOACS] -> Stms SOACS
forall rep. [Stm rep] -> Stms rep
stmsFromList ([Stm SOACS] -> Stms SOACS) -> [Stm SOACS] -> Stms SOACS
forall a b. (a -> b) -> a -> b
$
        Certs -> Stm SOACS -> Stm SOACS
forall rep. Certs -> Stm rep -> Stm rep
certify (StmAux dec -> Certs
forall dec. StmAux dec -> Certs
stmAuxCerts StmAux dec
aux)
          (Stm SOACS -> Stm SOACS) -> [Stm SOACS] -> [Stm SOACS]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Ident -> SubExp -> Stm SOACS)
-> [Ident] -> [SubExp] -> [Stm SOACS]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Ident -> SubExp -> Stm SOACS
forall rep. Buildable rep => Ident -> SubExp -> Stm rep
bindSubExp ((Param DeclType -> Ident) -> [Param DeclType] -> [Ident]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> Ident
forall dec. Typed dec => Param dec -> Ident
paramIdent ([Param DeclType] -> [Ident]) -> [Param DeclType] -> [Ident]
forall a b. (a -> b) -> a -> b
$ FunDef SOACS -> [FParam SOACS]
forall rep. FunDef rep -> [FParam rep]
funDefParams FunDef SOACS
fun) (((SubExp, Diet) -> SubExp) -> [(SubExp, Diet)] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (SubExp, Diet) -> SubExp
forall a b. (a, b) -> a
fst [(SubExp, Diet)]
args)

    body_stms :: Stms SOACS
body_stms =
      Attrs -> Safety -> [SrcLoc] -> Stms SOACS -> Stms SOACS
addLocations (StmAux dec -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux dec
aux) Safety
safety ((SrcLoc -> Bool) -> [SrcLoc] -> [SrcLoc]
forall a. (a -> Bool) -> [a] -> [a]
filter SrcLoc -> Bool
notmempty (SrcLoc
loc SrcLoc -> [SrcLoc] -> [SrcLoc]
forall a. a -> [a] -> [a]
: [SrcLoc]
locs)) (Stms SOACS -> Stms SOACS) -> Stms SOACS -> Stms SOACS
forall a b. (a -> b) -> a -> b
$
        Body SOACS -> Stms SOACS
forall rep. Body rep -> Stms rep
bodyStms (Body SOACS -> Stms SOACS) -> Body SOACS -> Stms SOACS
forall a b. (a -> b) -> a -> b
$ FunDef SOACS -> Body SOACS
forall rep. FunDef rep -> Body rep
funDefBody FunDef SOACS
fun

    -- Note that the sizes of arrays may not be correct at this
    -- point - it is crucial that we run copy propagation before
    -- the type checker sees this!
    bindSubExp :: Ident -> SubExp -> Stm rep
bindSubExp Ident
ident SubExp
se =
      [Ident] -> Exp rep -> Stm rep
forall rep. Buildable rep => [Ident] -> Exp rep -> Stm rep
mkLet [Ident
ident] (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp rep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se

    bindSubExpRes :: Ident -> SubExpRes -> Stm rep
bindSubExpRes Ident
ident (SubExpRes Certs
cs SubExp
se) =
      Certs -> Stm rep -> Stm rep
forall rep. Certs -> Stm rep -> Stm rep
certify Certs
cs (Stm rep -> Stm rep) -> Stm rep -> Stm rep
forall a b. (a -> b) -> a -> b
$ Ident -> SubExp -> Stm rep
forall rep. Buildable rep => Ident -> SubExp -> Stm rep
bindSubExp Ident
ident SubExp
se

    notmempty :: SrcLoc -> Bool
notmempty = (Loc -> Loc -> Bool
forall a. Eq a => a -> a -> Bool
/= Loc
forall a. Monoid a => a
mempty) (Loc -> Bool) -> (SrcLoc -> Loc) -> SrcLoc -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf

inlineInStms ::
  MonadFreshNames m =>
  M.Map Name (FunDef SOACS) ->
  Stms SOACS ->
  m (Stms SOACS)
inlineInStms :: Map Name (FunDef SOACS) -> Stms SOACS -> m (Stms SOACS)
inlineInStms Map Name (FunDef SOACS)
fdmap Stms SOACS
stms =
  Body SOACS -> Stms SOACS
forall rep. Body rep -> Stms rep
bodyStms (Body SOACS -> Stms SOACS) -> m (Body SOACS) -> m (Stms SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Name (FunDef SOACS) -> Body SOACS -> m (Body SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Map Name (FunDef SOACS) -> Body SOACS -> m (Body SOACS)
inlineInBody Map Name (FunDef SOACS)
fdmap (Stms SOACS -> Result -> Body SOACS
forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody Stms SOACS
stms [])

inlineInBody ::
  MonadFreshNames m =>
  M.Map Name (FunDef SOACS) ->
  Body SOACS ->
  m (Body SOACS)
inlineInBody :: Map Name (FunDef SOACS) -> Body SOACS -> m (Body SOACS)
inlineInBody Map Name (FunDef SOACS)
fdmap = Body SOACS -> m (Body SOACS)
onBody
  where
    inline :: [Stm SOACS] -> m (Stms SOACS)
inline (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Apply Name
fname [(SubExp, Diet)]
args [RetType SOACS]
_ (Safety, SrcLoc, [SrcLoc])
what) : [Stm SOACS]
rest)
      | Just FunDef SOACS
fd <- Name -> Map Name (FunDef SOACS) -> Maybe (FunDef SOACS)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
fname Map Name (FunDef SOACS)
fdmap,
        Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Attr
"noinline" Attr -> Attrs -> Bool
`inAttrs` FunDef SOACS -> Attrs
forall rep. FunDef rep -> Attrs
funDefAttrs FunDef SOACS
fd,
        Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Attr
"noinline" Attr -> Attrs -> Bool
`inAttrs` StmAux () -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux ()
StmAux (ExpDec SOACS)
aux =
          Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
(<>) (Stms SOACS -> Stms SOACS -> Stms SOACS)
-> m (Stms SOACS) -> m (Stms SOACS -> Stms SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Pat Type
-> StmAux ()
-> [(SubExp, Diet)]
-> (Safety, SrcLoc, [SrcLoc])
-> FunDef SOACS
-> m (Stms SOACS)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Pat Type
-> StmAux dec
-> [(SubExp, Diet)]
-> (Safety, SrcLoc, [SrcLoc])
-> FunDef SOACS
-> m (Stms SOACS)
inlineFunction Pat Type
Pat (LetDec SOACS)
pat StmAux ()
StmAux (ExpDec SOACS)
aux [(SubExp, Diet)]
args (Safety, SrcLoc, [SrcLoc])
what FunDef SOACS
fd m (Stms SOACS -> Stms SOACS) -> m (Stms SOACS) -> m (Stms SOACS)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Stm SOACS] -> m (Stms SOACS)
inline [Stm SOACS]
rest
    inline (stm :: Stm SOACS
stm@(Let Pat (LetDec SOACS)
_ StmAux (ExpDec SOACS)
_ BasicOp {}) : [Stm SOACS]
rest) =
      (Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm Stm SOACS
stm Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<>) (Stms SOACS -> Stms SOACS) -> m (Stms SOACS) -> m (Stms SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Stm SOACS] -> m (Stms SOACS)
inline [Stm SOACS]
rest
    inline (Stm SOACS
stm : [Stm SOACS]
rest) =
      Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
(<>) (Stms SOACS -> Stms SOACS -> Stms SOACS)
-> m (Stms SOACS) -> m (Stms SOACS -> Stms SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm (Stm SOACS -> Stms SOACS) -> m (Stm SOACS) -> m (Stms SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Stm SOACS -> m (Stm SOACS)
onStm Stm SOACS
stm) m (Stms SOACS -> Stms SOACS) -> m (Stms SOACS) -> m (Stms SOACS)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Stm SOACS] -> m (Stms SOACS)
inline [Stm SOACS]
rest
    inline [] =
      Stms SOACS -> m (Stms SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Stms SOACS
forall a. Monoid a => a
mempty

    onBody :: Body SOACS -> m (Body SOACS)
onBody (Body BodyDec SOACS
dec Stms SOACS
stms Result
res) =
      BodyDec SOACS -> Stms SOACS -> Result -> Body SOACS
forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
Body BodyDec SOACS
dec (Stms SOACS -> Result -> Body SOACS)
-> m (Stms SOACS) -> m (Result -> Body SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Stm SOACS] -> m (Stms SOACS)
inline (Stms SOACS -> [Stm SOACS]
forall rep. Stms rep -> [Stm rep]
stmsToList Stms SOACS
stms) m (Result -> Body SOACS) -> m Result -> m (Body SOACS)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Result -> m Result
forall (f :: * -> *) a. Applicative f => a -> f a
pure Result
res
    onStm :: Stm SOACS -> m (Stm SOACS)
onStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux Exp SOACS
e) = Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Exp SOACS -> Stm SOACS) -> m (Exp SOACS) -> m (Stm SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper SOACS SOACS m -> Exp SOACS -> m (Exp SOACS)
forall (m :: * -> *) frep trep.
(Applicative m, Monad m) =>
Mapper frep trep m -> Exp frep -> m (Exp trep)
mapExpM Mapper SOACS SOACS m
inliner Exp SOACS
e

    inliner :: Mapper SOACS SOACS m
inliner =
      Mapper SOACS SOACS m
forall (m :: * -> *) rep. Monad m => Mapper rep rep m
identityMapper
        { mapOnBody :: Scope SOACS -> Body SOACS -> m (Body SOACS)
mapOnBody = (Body SOACS -> m (Body SOACS))
-> Scope SOACS -> Body SOACS -> m (Body SOACS)
forall a b. a -> b -> a
const Body SOACS -> m (Body SOACS)
onBody,
          mapOnOp :: Op SOACS -> m (Op SOACS)
mapOnOp = Op SOACS -> m (Op SOACS)
SOAC SOACS -> m (SOAC SOACS)
onSOAC
        }

    onSOAC :: SOAC SOACS -> m (SOAC SOACS)
onSOAC =
      SOACMapper SOACS SOACS m -> SOAC SOACS -> m (SOAC SOACS)
forall (m :: * -> *) frep trep.
(Applicative m, Monad m) =>
SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
mapSOACM SOACMapper Any Any m
forall (m :: * -> *) rep. Monad m => SOACMapper rep rep m
identitySOACMapper {mapOnSOACLambda :: Lambda SOACS -> m (Lambda SOACS)
mapOnSOACLambda = Lambda SOACS -> m (Lambda SOACS)
onLambda}

    onLambda :: Lambda SOACS -> m (Lambda SOACS)
onLambda (Lambda [LParam SOACS]
params Body SOACS
body [Type]
ret) =
      [LParam SOACS] -> Body SOACS -> [Type] -> Lambda SOACS
forall rep. [LParam rep] -> Body rep -> [Type] -> Lambda rep
Lambda [LParam SOACS]
params (Body SOACS -> [Type] -> Lambda SOACS)
-> m (Body SOACS) -> m ([Type] -> Lambda SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Body SOACS -> m (Body SOACS)
onBody Body SOACS
body m ([Type] -> Lambda SOACS) -> m [Type] -> m (Lambda SOACS)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Type] -> m [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Type]
ret

-- Propagate source locations and attributes to the inlined
-- statements.  Attributes are propagated only when applicable (this
-- probably means that every supported attribute needs to be handled
-- specially here).
addLocations :: Attrs -> Safety -> [SrcLoc] -> Stms SOACS -> Stms SOACS
addLocations :: Attrs -> Safety -> [SrcLoc] -> Stms SOACS -> Stms SOACS
addLocations Attrs
attrs Safety
caller_safety [SrcLoc]
more_locs = (Stm SOACS -> Stm SOACS) -> Stms SOACS -> Stms SOACS
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Stm SOACS -> Stm SOACS
onStm
  where
    onStm :: Stm SOACS -> Stm SOACS
onStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Apply Name
fname [(SubExp, Diet)]
args [RetType SOACS]
t (Safety
safety, SrcLoc
loc, [SrcLoc]
locs))) =
      Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux' (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
forall a b. (a -> b) -> a -> b
$
        Name
-> [(SubExp, Diet)]
-> [RetType SOACS]
-> (Safety, SrcLoc, [SrcLoc])
-> Exp SOACS
forall rep.
Name
-> [(SubExp, Diet)]
-> [RetType rep]
-> (Safety, SrcLoc, [SrcLoc])
-> Exp rep
Apply Name
fname [(SubExp, Diet)]
args [RetType SOACS]
t (Safety -> Safety -> Safety
forall a. Ord a => a -> a -> a
min Safety
caller_safety Safety
safety, SrcLoc
loc, [SrcLoc]
locs [SrcLoc] -> [SrcLoc] -> [SrcLoc]
forall a. [a] -> [a] -> [a]
++ [SrcLoc]
more_locs)
      where
        aux' :: StmAux (ExpDec SOACS)
aux' = StmAux (ExpDec SOACS)
aux {stmAuxAttrs :: Attrs
stmAuxAttrs = Attrs
attrs Attrs -> Attrs -> Attrs
forall a. Semigroup a => a -> a -> a
<> StmAux (ExpDec SOACS) -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux (ExpDec SOACS)
aux}
    onStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (BasicOp (Assert SubExp
cond ErrorMsg SubExp
desc (SrcLoc
loc, [SrcLoc]
locs)))) =
      Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
pat (Attrs -> StmAux () -> StmAux ()
forall dec. Attrs -> StmAux dec -> StmAux dec
withAttrs (Attrs -> Attrs
attrsForAssert Attrs
attrs) StmAux ()
StmAux (ExpDec SOACS)
aux) (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
forall a b. (a -> b) -> a -> b
$
        case Safety
caller_safety of
          Safety
Safe -> BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> ErrorMsg SubExp -> (SrcLoc, [SrcLoc]) -> BasicOp
Assert SubExp
cond ErrorMsg SubExp
desc (SrcLoc
loc, [SrcLoc]
locs [SrcLoc] -> [SrcLoc] -> [SrcLoc]
forall a. [a] -> [a] -> [a]
++ [SrcLoc]
more_locs)
          Safety
Unsafe -> BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
Constant PrimValue
UnitValue
    onStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Op Op SOACS
soac)) =
      Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
pat (Attrs -> StmAux () -> StmAux ()
forall dec. Attrs -> StmAux dec -> StmAux dec
withAttrs Attrs
attrs' StmAux ()
StmAux (ExpDec SOACS)
aux) (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
forall a b. (a -> b) -> a -> b
$
        Op SOACS -> Exp SOACS
forall rep. Op rep -> Exp rep
Op (Op SOACS -> Exp SOACS) -> Op SOACS -> Exp SOACS
forall a b. (a -> b) -> a -> b
$
          Identity (SOAC SOACS) -> SOAC SOACS
forall a. Identity a -> a
runIdentity (Identity (SOAC SOACS) -> SOAC SOACS)
-> Identity (SOAC SOACS) -> SOAC SOACS
forall a b. (a -> b) -> a -> b
$
            SOACMapper SOACS SOACS Identity
-> SOAC SOACS -> Identity (SOAC SOACS)
forall (m :: * -> *) frep trep.
(Applicative m, Monad m) =>
SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
mapSOACM
              SOACMapper Any Any Identity
forall (m :: * -> *) rep. Monad m => SOACMapper rep rep m
identitySOACMapper
                { mapOnSOACLambda :: Lambda SOACS -> Identity (Lambda SOACS)
mapOnSOACLambda = Lambda SOACS -> Identity (Lambda SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Lambda SOACS -> Identity (Lambda SOACS))
-> (Lambda SOACS -> Lambda SOACS)
-> Lambda SOACS
-> Identity (Lambda SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lambda SOACS -> Lambda SOACS
onLambda
                }
              Op SOACS
SOAC SOACS
soac
      where
        attrs' :: Attrs
attrs' = Attrs
attrs Attrs -> Attrs -> Attrs
`withoutAttrs` Attrs
for_assert
        for_assert :: Attrs
for_assert = Attrs -> Attrs
attrsForAssert Attrs
attrs
        onLambda :: Lambda SOACS -> Lambda SOACS
onLambda Lambda SOACS
lam =
          Lambda SOACS
lam {lambdaBody :: Body SOACS
lambdaBody = Attrs -> Body SOACS -> Body SOACS
onBody Attrs
for_assert (Body SOACS -> Body SOACS) -> Body SOACS -> Body SOACS
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep. Lambda rep -> Body rep
lambdaBody Lambda SOACS
lam}
    onStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux Exp SOACS
e) =
      Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
forall a b. (a -> b) -> a -> b
$ Exp SOACS -> Exp SOACS
onExp Exp SOACS
e

    onExp :: Exp SOACS -> Exp SOACS
onExp =
      Mapper SOACS SOACS Identity -> Exp SOACS -> Exp SOACS
forall frep trep. Mapper frep trep Identity -> Exp frep -> Exp trep
mapExp
        Mapper SOACS SOACS Identity
forall (m :: * -> *) rep. Monad m => Mapper rep rep m
identityMapper
          { mapOnBody :: Scope SOACS -> Body SOACS -> Identity (Body SOACS)
mapOnBody = (Body SOACS -> Identity (Body SOACS))
-> Scope SOACS -> Body SOACS -> Identity (Body SOACS)
forall a b. a -> b -> a
const ((Body SOACS -> Identity (Body SOACS))
 -> Scope SOACS -> Body SOACS -> Identity (Body SOACS))
-> (Body SOACS -> Identity (Body SOACS))
-> Scope SOACS
-> Body SOACS
-> Identity (Body SOACS)
forall a b. (a -> b) -> a -> b
$ Body SOACS -> Identity (Body SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Body SOACS -> Identity (Body SOACS))
-> (Body SOACS -> Body SOACS)
-> Body SOACS
-> Identity (Body SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Attrs -> Body SOACS -> Body SOACS
onBody Attrs
attrs
          }

    withAttrs :: Attrs -> StmAux dec -> StmAux dec
withAttrs Attrs
attrs' StmAux dec
aux = StmAux dec
aux {stmAuxAttrs :: Attrs
stmAuxAttrs = Attrs
attrs' Attrs -> Attrs -> Attrs
forall a. Semigroup a => a -> a -> a
<> StmAux dec -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux dec
aux}

    onBody :: Attrs -> Body SOACS -> Body SOACS
onBody Attrs
attrs' Body SOACS
body =
      Body SOACS
body
        { bodyStms :: Stms SOACS
bodyStms =
            Attrs -> Safety -> [SrcLoc] -> Stms SOACS -> Stms SOACS
addLocations Attrs
attrs' Safety
caller_safety [SrcLoc]
more_locs (Stms SOACS -> Stms SOACS) -> Stms SOACS -> Stms SOACS
forall a b. (a -> b) -> a -> b
$
              Body SOACS -> Stms SOACS
forall rep. Body rep -> Stms rep
bodyStms Body SOACS
body
        }

-- | Remove functions not ultimately called from an entry point or a
-- constant.
removeDeadFunctionsF :: Prog SOACS -> Prog SOACS
removeDeadFunctionsF :: Prog SOACS -> Prog SOACS
removeDeadFunctionsF Prog SOACS
prog =
  let cg :: CallGraph
cg = Prog SOACS -> CallGraph
buildCallGraph Prog SOACS
prog
      live_funs :: [FunDef SOACS]
live_funs = (FunDef SOACS -> Bool) -> [FunDef SOACS] -> [FunDef SOACS]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Name -> CallGraph -> Bool
`isFunInCallGraph` CallGraph
cg) (Name -> Bool) -> (FunDef SOACS -> Name) -> FunDef SOACS -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FunDef SOACS -> Name
forall rep. FunDef rep -> Name
funDefName) ([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
   in Prog SOACS
prog {progFuns :: [FunDef SOACS]
progFuns = [FunDef SOACS]
live_funs}

-- | Inline all functions and remove the resulting dead functions.
inlineAggressively :: Pass SOACS SOACS
inlineAggressively :: Pass SOACS SOACS
inlineAggressively =
  Pass :: forall fromrep torep.
String
-> String
-> (Prog fromrep -> PassM (Prog torep))
-> Pass fromrep torep
Pass
    { passName :: String
passName = String
"Inline aggressively",
      passDescription :: String
passDescription = String
"Aggressively inline and remove resulting dead functions.",
      passFunction :: Prog SOACS -> PassM (Prog SOACS)
passFunction =
        SimpleOps SOACS -> Prog SOACS -> PassM (Prog SOACS)
forall rep.
SimplifiableRep rep =>
SimpleOps rep -> Prog rep -> PassM (Prog rep)
copyPropagateInProg SimpleOps SOACS
simpleSOACS (Prog SOACS -> PassM (Prog SOACS))
-> (Prog SOACS -> Prog SOACS) -> Prog SOACS -> PassM (Prog SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog SOACS -> Prog SOACS
removeDeadFunctionsF (Prog SOACS -> PassM (Prog SOACS))
-> (Prog SOACS -> PassM (Prog SOACS))
-> Prog SOACS
-> PassM (Prog SOACS)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Prog SOACS -> PassM (Prog SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Prog SOACS -> m (Prog SOACS)
aggInlineFunctions
    }

-- | Inline some functions and remove the resulting dead functions.
inlineConservatively :: Pass SOACS SOACS
inlineConservatively :: Pass SOACS SOACS
inlineConservatively =
  Pass :: forall fromrep torep.
String
-> String
-> (Prog fromrep -> PassM (Prog torep))
-> Pass fromrep torep
Pass
    { passName :: String
passName = String
"Inline conservatively",
      passDescription :: String
passDescription = String
"Conservatively inline and remove resulting dead functions.",
      passFunction :: Prog SOACS -> PassM (Prog SOACS)
passFunction =
        SimpleOps SOACS -> Prog SOACS -> PassM (Prog SOACS)
forall rep.
SimplifiableRep rep =>
SimpleOps rep -> Prog rep -> PassM (Prog rep)
copyPropagateInProg SimpleOps SOACS
simpleSOACS (Prog SOACS -> PassM (Prog SOACS))
-> (Prog SOACS -> Prog SOACS) -> Prog SOACS -> PassM (Prog SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog SOACS -> Prog SOACS
removeDeadFunctionsF (Prog SOACS -> PassM (Prog SOACS))
-> (Prog SOACS -> PassM (Prog SOACS))
-> Prog SOACS
-> PassM (Prog SOACS)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Prog SOACS -> PassM (Prog SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Prog SOACS -> m (Prog SOACS)
consInlineFunctions
    }

-- | @removeDeadFunctions prog@ removes the functions that are unreachable from
-- the main function from the program.
removeDeadFunctions :: Pass SOACS SOACS
removeDeadFunctions :: Pass SOACS SOACS
removeDeadFunctions =
  Pass :: forall fromrep torep.
String
-> String
-> (Prog fromrep -> PassM (Prog torep))
-> Pass fromrep torep
Pass
    { passName :: String
passName = String
"Remove dead functions",
      passDescription :: String
passDescription = String
"Remove the functions that are unreachable from entry points",
      passFunction :: Prog SOACS -> PassM (Prog SOACS)
passFunction = Prog SOACS -> PassM (Prog SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Prog SOACS -> PassM (Prog SOACS))
-> (Prog SOACS -> Prog SOACS) -> Prog SOACS -> PassM (Prog SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog SOACS -> Prog SOACS
removeDeadFunctionsF
    }