{-# LANGUAGE TypeFamilies #-}

-- | Preprocess the program before flattening.  This rewrites SOAC forms
-- that flatten does not want to see directly, while leaving the result in
-- SOACS form so the normal flattening pipeline can continue afterwards.
module Futhark.Pass.Flatten.PreProcess
  ( shouldDissectForm,
    preprocessProg,
    preprocessBody,
    preprocessStms,
    preprocessStm,
    preprocessLambda,
    runSimplifiedBuilder,
  )
where

import Data.Maybe (isNothing)
import Futhark.Builder
import Futhark.IR.SOACS
import Futhark.IR.SOACS.Simplify
import Futhark.Pass
import Futhark.Tools
import Futhark.Transform.FirstOrderTransform qualified as FOT
import Futhark.Transform.ISRWIM (irwim, iswim)

shouldDissectForm :: ScremaForm SOACS -> Bool
shouldDissectForm :: ScremaForm SOACS -> Bool
shouldDissectForm ScremaForm SOACS
form =
  Maybe (Lambda SOACS) -> Bool
forall a. Maybe a -> Bool
isNothing (ScremaForm SOACS -> Maybe (Lambda SOACS)
forall rep. ScremaForm rep -> Maybe (Lambda rep)
isMapSOAC ScremaForm SOACS
form)
    Bool -> Bool -> Bool
&& Maybe [Reduce SOACS] -> Bool
forall a. Maybe a -> Bool
isNothing (ScremaForm SOACS -> Maybe [Reduce SOACS]
forall rep. ScremaForm rep -> Maybe [Reduce rep]
isReduceSOAC ScremaForm SOACS
form)
    Bool -> Bool -> Bool
&& Maybe [Scan SOACS] -> Bool
forall a. Maybe a -> Bool
isNothing (ScremaForm SOACS -> Maybe [Scan SOACS]
forall rep. ScremaForm rep -> Maybe [Scan rep]
isScanSOAC ScremaForm SOACS
form)
    Bool -> Bool -> Bool
&& Maybe ([Reduce SOACS], Lambda SOACS) -> Bool
forall a. Maybe a -> Bool
isNothing (ScremaForm SOACS -> Maybe ([Reduce SOACS], Lambda SOACS)
forall rep. ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
isRedomapSOAC ScremaForm SOACS
form)
    Bool -> Bool -> Bool
&& Maybe ([Scan SOACS], Lambda SOACS) -> Bool
forall a. Maybe a -> Bool
isNothing (ScremaForm SOACS -> Maybe ([Scan SOACS], Lambda SOACS)
forall rep. ScremaForm rep -> Maybe ([Scan rep], Lambda rep)
isScanomapSOAC ScremaForm SOACS
form)
    Bool -> Bool -> Bool
&& Maybe (Lambda SOACS, [Scan SOACS], Lambda SOACS) -> Bool
forall a. Maybe a -> Bool
isNothing (ScremaForm SOACS
-> Maybe (Lambda SOACS, [Scan SOACS], Lambda SOACS)
forall rep.
ScremaForm rep -> Maybe (Lambda rep, [Scan rep], Lambda rep)
isMaposcanomapSOAC ScremaForm SOACS
form)

runSimplifiedBuilder ::
  (MonadFreshNames m) =>
  Scope SOACS ->
  BuilderT SOACS m a ->
  m (Stms SOACS)
runSimplifiedBuilder :: forall (m :: * -> *) a.
MonadFreshNames m =>
Scope SOACS -> BuilderT SOACS m a -> m (Stms SOACS)
runSimplifiedBuilder Scope SOACS
scope BuilderT SOACS m a
m =
  (Stms SOACS, Stms SOACS) -> Stms SOACS
forall a b. (a, b) -> a
fst ((Stms SOACS, Stms SOACS) -> Stms SOACS)
-> m (Stms SOACS, Stms SOACS) -> m (Stms SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> BuilderT SOACS m (Stms SOACS)
-> Scope SOACS -> m (Stms SOACS, Stms SOACS)
forall (m :: * -> *) rep a.
MonadFreshNames m =>
BuilderT rep m a -> Scope rep -> m (a, Stms rep)
runBuilderT (Stms SOACS -> BuilderT SOACS m (Stms SOACS)
forall (m :: * -> *).
(HasScope SOACS m, MonadFreshNames m) =>
Stms SOACS -> m (Stms SOACS)
simplifyStms (Stms SOACS -> BuilderT SOACS m (Stms SOACS))
-> BuilderT SOACS m (Stms SOACS) -> BuilderT SOACS m (Stms SOACS)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BuilderT SOACS m a
-> BuilderT SOACS m (Stms (Rep (BuilderT SOACS m)))
forall (m :: * -> *) a. MonadBuilder m => m a -> m (Stms (Rep m))
collectStms_ BuilderT SOACS m a
m) Scope SOACS
scope

-- | Rewrite a SOAC form that flattening does not handle directly into one it
-- does, recursively preprocessing the result (which may itself contain further
-- SOACs). Returns 'Nothing' for statements that need no rewriting - those are
-- handled structurally by 'preprocessStm'.
rewriteSoacStm ::
  (MonadFreshNames m) =>
  Scope SOACS ->
  Stm SOACS ->
  Maybe (m (Stms SOACS))
rewriteSoacStm :: forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Stm SOACS -> Maybe (m (Stms SOACS))
rewriteSoacStm Scope SOACS
scope (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Op Op SOACS
soac))
  | Attr
"sequential_outer" Attr -> Attrs -> Bool
`inAttrs` StmAux () -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux ()
StmAux (ExpDec SOACS)
aux =
      m (Stms SOACS) -> Maybe (m (Stms SOACS))
forall a. a -> Maybe a
Just (m (Stms SOACS) -> Maybe (m (Stms SOACS)))
-> m (Stms SOACS) -> Maybe (m (Stms SOACS))
forall a b. (a -> b) -> a -> b
$
        Scope SOACS -> Stms SOACS -> m (Stms SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Stms SOACS -> m (Stms SOACS)
preprocessStms Scope SOACS
scope (Stms SOACS -> m (Stms SOACS)) -> m (Stms SOACS) -> m (Stms SOACS)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Scope SOACS -> BuilderT SOACS m () -> m (Stms SOACS)
forall (m :: * -> *) a.
MonadFreshNames m =>
Scope SOACS -> BuilderT SOACS m a -> m (Stms SOACS)
runSimplifiedBuilder Scope SOACS
scope (Pat (LetDec (Rep (BuilderT SOACS m)))
-> SOAC (Rep (BuilderT SOACS m)) -> BuilderT SOACS m ()
forall (m :: * -> *).
Transformer m =>
Pat (LetDec (Rep m)) -> SOAC (Rep m) -> m ()
FOT.transformSOAC Pat (LetDec (Rep (BuilderT SOACS m)))
Pat (LetDec SOACS)
pat Op SOACS
SOAC (Rep (BuilderT SOACS m))
soac)
rewriteSoacStm Scope SOACS
scope (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Op (Stream SubExp
w [VName]
arrs [SubExp]
nes Lambda SOACS
lam))) = m (Stms SOACS) -> Maybe (m (Stms SOACS))
forall a. a -> Maybe a
Just (m (Stms SOACS) -> Maybe (m (Stms SOACS)))
-> m (Stms SOACS) -> Maybe (m (Stms SOACS))
forall a b. (a -> b) -> a -> b
$ do
  stms <- Scope SOACS -> BuilderT SOACS m () -> m (Stms SOACS)
forall (m :: * -> *) a.
MonadFreshNames m =>
Scope SOACS -> BuilderT SOACS m a -> m (Stms SOACS)
runSimplifiedBuilder Scope SOACS
scope (StmAux () -> BuilderT SOACS m () -> BuilderT SOACS m ()
forall (m :: * -> *) anyrep a.
MonadBuilder m =>
StmAux anyrep -> m a -> m a
auxing StmAux ()
StmAux (ExpDec SOACS)
aux (BuilderT SOACS m () -> BuilderT SOACS m ())
-> BuilderT SOACS m () -> BuilderT SOACS m ()
forall a b. (a -> b) -> a -> b
$ Pat (LetDec (Rep (BuilderT SOACS m)))
-> SubExp
-> [SubExp]
-> Lambda (Rep (BuilderT SOACS m))
-> [VName]
-> BuilderT SOACS m ()
forall (m :: * -> *).
(MonadBuilder m, Buildable (Rep m)) =>
Pat (LetDec (Rep m))
-> SubExp -> [SubExp] -> Lambda (Rep m) -> [VName] -> m ()
sequentialStreamWholeArray Pat (LetDec (Rep (BuilderT SOACS m)))
Pat (LetDec SOACS)
pat SubExp
w [SubExp]
nes Lambda (Rep (BuilderT SOACS m))
Lambda SOACS
lam [VName]
arrs)
  preprocessStms scope stms
rewriteSoacStm Scope SOACS
scope (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Op (Screma SubExp
w' [VName]
arrs' ScremaForm SOACS
form')))
  | Just [Scan SOACS]
scans <- ScremaForm SOACS -> Maybe [Scan SOACS]
forall rep. ScremaForm rep -> Maybe [Scan rep]
isScanSOAC ScremaForm SOACS
form',
    Scan Lambda SOACS
scan_lam [SubExp]
nes <- [Scan SOACS] -> Scan SOACS
forall rep. Buildable rep => [Scan rep] -> Scan rep
singleScan [Scan SOACS]
scans,
    Just BuilderT SOACS m ()
do_iswim <- Pat Type
-> SubExp
-> Lambda SOACS
-> [(SubExp, VName)]
-> Maybe (BuilderT SOACS m ())
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ SOACS) =>
Pat Type
-> SubExp -> Lambda SOACS -> [(SubExp, VName)] -> Maybe (m ())
iswim Pat Type
Pat (LetDec SOACS)
pat SubExp
w' Lambda SOACS
scan_lam ([SubExp] -> [VName] -> [(SubExp, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
nes [VName]
arrs') = m (Stms SOACS) -> Maybe (m (Stms SOACS))
forall a. a -> Maybe a
Just (m (Stms SOACS) -> Maybe (m (Stms SOACS)))
-> m (Stms SOACS) -> Maybe (m (Stms SOACS))
forall a b. (a -> b) -> a -> b
$ do
      stms <- Scope SOACS -> BuilderT SOACS m () -> m (Stms SOACS)
forall (m :: * -> *) a.
MonadFreshNames m =>
Scope SOACS -> BuilderT SOACS m a -> m (Stms SOACS)
runSimplifiedBuilder Scope SOACS
scope (BuilderT SOACS m () -> m (Stms SOACS))
-> BuilderT SOACS m () -> m (Stms SOACS)
forall a b. (a -> b) -> a -> b
$ StmAux () -> BuilderT SOACS m () -> BuilderT SOACS m ()
forall (m :: * -> *) anyrep a.
MonadBuilder m =>
StmAux anyrep -> m a -> m a
auxing StmAux ()
StmAux (ExpDec SOACS)
aux BuilderT SOACS m ()
do_iswim
      preprocessStms scope stms
  | Just [Reduce Commutativity
comm Lambda SOACS
red_fun [SubExp]
nes] <- ScremaForm SOACS -> Maybe [Reduce SOACS]
forall rep. ScremaForm rep -> Maybe [Reduce rep]
isReduceSOAC ScremaForm SOACS
form',
    let comm' :: Commutativity
comm'
          | Lambda SOACS -> Bool
forall rep. Lambda rep -> Bool
commutativeLambda Lambda SOACS
red_fun = Commutativity
Commutative
          | Bool
otherwise = Commutativity
comm,
    Just BuilderT SOACS m ()
do_irwim <- Pat Type
-> SubExp
-> Commutativity
-> Lambda SOACS
-> [(SubExp, VName)]
-> Maybe (BuilderT SOACS m ())
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ SOACS) =>
Pat Type
-> SubExp
-> Commutativity
-> Lambda SOACS
-> [(SubExp, VName)]
-> Maybe (m ())
irwim Pat Type
Pat (LetDec SOACS)
pat SubExp
w' Commutativity
comm' Lambda SOACS
red_fun ([SubExp] -> [VName] -> [(SubExp, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
nes [VName]
arrs') = m (Stms SOACS) -> Maybe (m (Stms SOACS))
forall a. a -> Maybe a
Just (m (Stms SOACS) -> Maybe (m (Stms SOACS)))
-> m (Stms SOACS) -> Maybe (m (Stms SOACS))
forall a b. (a -> b) -> a -> b
$ do
      stms <- Scope SOACS -> BuilderT SOACS m () -> m (Stms SOACS)
forall (m :: * -> *) a.
MonadFreshNames m =>
Scope SOACS -> BuilderT SOACS m a -> m (Stms SOACS)
runSimplifiedBuilder Scope SOACS
scope (BuilderT SOACS m () -> m (Stms SOACS))
-> BuilderT SOACS m () -> m (Stms SOACS)
forall a b. (a -> b) -> a -> b
$ StmAux () -> BuilderT SOACS m () -> BuilderT SOACS m ()
forall (m :: * -> *) anyrep a.
MonadBuilder m =>
StmAux anyrep -> m a -> m a
auxing StmAux ()
StmAux (ExpDec SOACS)
aux BuilderT SOACS m ()
do_irwim
      preprocessStms scope stms
  | ScremaForm SOACS -> Bool
shouldDissectForm ScremaForm SOACS
form' = m (Stms SOACS) -> Maybe (m (Stms SOACS))
forall a. a -> Maybe a
Just (m (Stms SOACS) -> Maybe (m (Stms SOACS)))
-> m (Stms SOACS) -> Maybe (m (Stms SOACS))
forall a b. (a -> b) -> a -> b
$ do
      stms <- Scope SOACS -> BuilderT SOACS m () -> m (Stms SOACS)
forall (m :: * -> *) a.
MonadFreshNames m =>
Scope SOACS -> BuilderT SOACS m a -> m (Stms SOACS)
runSimplifiedBuilder Scope SOACS
scope (StmAux () -> BuilderT SOACS m () -> BuilderT SOACS m ()
forall (m :: * -> *) anyrep a.
MonadBuilder m =>
StmAux anyrep -> m a -> m a
auxing StmAux ()
StmAux (ExpDec SOACS)
aux (BuilderT SOACS m () -> BuilderT SOACS m ())
-> BuilderT SOACS m () -> BuilderT SOACS m ()
forall a b. (a -> b) -> a -> b
$ Pat (LetDec (Rep (BuilderT SOACS m)))
-> SubExp
-> ScremaForm (Rep (BuilderT SOACS m))
-> [VName]
-> BuilderT SOACS m ()
forall (m :: * -> *).
(MonadBuilder m, Op (Rep m) ~ SOAC (Rep m), Buildable (Rep m)) =>
Pat (LetDec (Rep m))
-> SubExp -> ScremaForm (Rep m) -> [VName] -> m ()
dissectScrema Pat (LetDec (Rep (BuilderT SOACS m)))
Pat (LetDec SOACS)
pat SubExp
w' ScremaForm (Rep (BuilderT SOACS m))
ScremaForm SOACS
form' [VName]
arrs')
      preprocessStms scope stms
rewriteSoacStm Scope SOACS
_ Stm SOACS
_ = Maybe (m (Stms SOACS))
forall a. Maybe a
Nothing

preprocessStm ::
  (MonadFreshNames m) =>
  Scope SOACS ->
  Stm SOACS ->
  m (Stms SOACS)
preprocessStm :: forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Stm SOACS -> m (Stms SOACS)
preprocessStm Scope SOACS
_ Stm SOACS
stm
  | Attr
"sequential" Attr -> Attrs -> Bool
`inAttrs` StmAux () -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs (Stm SOACS -> StmAux (ExpDec SOACS)
forall rep. Stm rep -> StmAux (ExpDec rep)
stmAux Stm SOACS
stm) = Stms SOACS -> m (Stms SOACS)
forall a. a -> m a
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
$ Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm Stm SOACS
stm
preprocessStm Scope SOACS
scope Stm SOACS
stm
  | Just m (Stms SOACS)
rewritten <- Scope SOACS -> Stm SOACS -> Maybe (m (Stms SOACS))
forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Stm SOACS -> Maybe (m (Stms SOACS))
rewriteSoacStm Scope SOACS
scope Stm SOACS
stm = m (Stms SOACS)
rewritten
preprocessStm Scope SOACS
scope (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Loop [(FParam SOACS, SubExp)]
merge LoopForm
form Body SOACS
body)) = do
  let scope' :: Scope SOACS
scope' = [Param DeclType] -> Scope SOACS
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams (((Param DeclType, SubExp) -> Param DeclType)
-> [(Param DeclType, SubExp)] -> [Param DeclType]
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType, SubExp) -> Param DeclType
forall a b. (a, b) -> a
fst [(Param DeclType, SubExp)]
[(FParam SOACS, SubExp)]
merge) Scope SOACS -> Scope SOACS -> Scope SOACS
forall a. Semigroup a => a -> a -> a
<> LoopForm -> Scope SOACS
forall rep. LoopForm -> Scope rep
scopeOfLoopForm LoopForm
form Scope SOACS -> Scope SOACS -> Scope SOACS
forall a. Semigroup a => a -> a -> a
<> Scope SOACS
scope
  body' <- Scope SOACS -> Body SOACS -> m (Body SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Body SOACS -> m (Body SOACS)
preprocessBody Scope SOACS
scope' Body SOACS
body
  pure $ oneStm $ Let pat aux $ Loop merge form body'
preprocessStm Scope SOACS
scope (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Match [SubExp]
ses [Case (Body SOACS)]
cases Body SOACS
defbody MatchDec (BranchType SOACS)
dec)) = do
  cases' <- (Case (Body SOACS) -> m (Case (Body SOACS)))
-> [Case (Body SOACS)] -> m [Case (Body SOACS)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((Body SOACS -> m (Body SOACS))
-> Case (Body SOACS) -> m (Case (Body SOACS))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Case a -> f (Case b)
traverse (Scope SOACS -> Body SOACS -> m (Body SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Body SOACS -> m (Body SOACS)
preprocessBody Scope SOACS
scope)) [Case (Body SOACS)]
cases
  defbody' <- preprocessBody scope defbody
  pure $ oneStm $ Let pat aux $ Match ses cases' defbody' dec
preprocessStm Scope SOACS
scope (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (WithAcc [WithAccInput SOACS]
inputs Lambda SOACS
lam)) = do
  lam' <- Scope SOACS -> Lambda SOACS -> m (Lambda SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Lambda SOACS -> m (Lambda SOACS)
preprocessLambda Scope SOACS
scope Lambda SOACS
lam
  pure $ oneStm $ Let pat aux $ WithAcc inputs lam'
preprocessStm Scope SOACS
_ Stm SOACS
stm = Stms SOACS -> m (Stms SOACS)
forall a. a -> m a
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
$ Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm Stm SOACS
stm

preprocessStms ::
  (MonadFreshNames m) =>
  Scope SOACS ->
  Stms SOACS ->
  m (Stms SOACS)
preprocessStms :: forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Stms SOACS -> m (Stms SOACS)
preprocessStms Scope SOACS
scope Stms SOACS
stms = [Stms SOACS] -> Stms SOACS
forall a. Monoid a => [a] -> a
mconcat ([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)) -> [Stm SOACS] -> m [Stms SOACS]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Scope SOACS -> Stm SOACS -> m (Stms SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Stm SOACS -> m (Stms SOACS)
preprocessStm Scope SOACS
scope') (Stms SOACS -> [Stm SOACS]
forall rep. Stms rep -> [Stm rep]
stmsToList Stms SOACS
stms)
  where
    scope' :: Scope SOACS
scope' = Stms SOACS -> Scope SOACS
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stms SOACS
stms Scope SOACS -> Scope SOACS -> Scope SOACS
forall a. Semigroup a => a -> a -> a
<> Scope SOACS
scope

preprocessBody ::
  (MonadFreshNames m) =>
  Scope SOACS ->
  Body SOACS ->
  m (Body SOACS)
preprocessBody :: forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Body SOACS -> m (Body SOACS)
preprocessBody Scope SOACS
scope Body SOACS
body = do
  stms <- Scope SOACS -> Stms SOACS -> m (Stms SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Stms SOACS -> m (Stms SOACS)
preprocessStms Scope SOACS
scope (Stms SOACS -> m (Stms SOACS)) -> Stms SOACS -> m (Stms SOACS)
forall a b. (a -> b) -> a -> b
$ Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms Body SOACS
body
  pure $ body {bodyStms = stms}

preprocessLambda ::
  (MonadFreshNames m) =>
  Scope SOACS ->
  Lambda SOACS ->
  m (Lambda SOACS)
preprocessLambda :: forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Lambda SOACS -> m (Lambda SOACS)
preprocessLambda Scope SOACS
scope Lambda SOACS
lam = do
  body <- Scope SOACS -> Body SOACS -> m (Body SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Body SOACS -> m (Body SOACS)
preprocessBody ([Param Type] -> Scope SOACS
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam) Scope SOACS -> Scope SOACS -> Scope SOACS
forall a. Semigroup a => a -> a -> a
<> Scope SOACS
scope) (Body SOACS -> m (Body SOACS)) -> Body SOACS -> m (Body SOACS)
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam
  let lam' = Lambda SOACS
lam {lambdaBody = body}
  fst <$> runBuilderT (simplifyLambda lam') scope

preprocessFun :: Stms SOACS -> FunDef SOACS -> PassM (FunDef SOACS)
preprocessFun :: Stms SOACS -> FunDef SOACS -> PassM (FunDef SOACS)
preprocessFun Stms SOACS
consts FunDef SOACS
fd = do
  body <- Scope SOACS -> Body SOACS -> PassM (Body SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Body SOACS -> m (Body SOACS)
preprocessBody (Stms SOACS -> Scope SOACS
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stms SOACS
consts Scope SOACS -> Scope SOACS -> Scope SOACS
forall a. Semigroup a => a -> a -> a
<> FunDef SOACS -> Scope SOACS
forall rep a. Scoped rep a => a -> Scope rep
scopeOf FunDef SOACS
fd) (Body SOACS -> PassM (Body SOACS))
-> Body SOACS -> PassM (Body SOACS)
forall a b. (a -> b) -> a -> b
$ FunDef SOACS -> Body SOACS
forall rep. FunDef rep -> Body rep
funDefBody FunDef SOACS
fd
  pure $ fd {funDefBody = body}

preprocessProg :: Prog SOACS -> PassM (Prog SOACS)
preprocessProg :: Prog SOACS -> PassM (Prog SOACS)
preprocessProg =
  (Stms SOACS -> PassM (Stms SOACS))
-> (Stms SOACS -> FunDef SOACS -> PassM (FunDef SOACS))
-> Prog SOACS
-> PassM (Prog SOACS)
forall fromrep torep.
(Stms fromrep -> PassM (Stms torep))
-> (Stms torep -> FunDef fromrep -> PassM (FunDef torep))
-> Prog fromrep
-> PassM (Prog torep)
intraproceduralTransformationWithConsts
    (Scope SOACS -> Stms SOACS -> PassM (Stms SOACS)
forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Stms SOACS -> m (Stms SOACS)
preprocessStms Scope SOACS
forall a. Monoid a => a
mempty)
    Stms SOACS -> FunDef SOACS -> PassM (FunDef SOACS)
preprocessFun