-- |
--
-- Functions for generic traversals across Futhark syntax trees.  The
-- motivation for this module came from dissatisfaction with rewriting
-- the same trivial tree recursions for every module.  A possible
-- alternative would be to use normal \"Scrap your
-- boilerplate\"-techniques, but these are rejected for two reasons:
--
--    * They are too slow.
--
--    * More importantly, they do not tell you whether you have missed
--      some cases.
--
-- Instead, this module defines various traversals of the Futhark syntax
-- tree.  The implementation is rather tedious, but the interface is
-- easy to use.
--
-- A traversal of the Futhark syntax tree is expressed as a record of
-- functions expressing the operations to be performed on the various
-- types of nodes.
--
-- The "Futhark.Transform.Rename" module is a simple example of how to
-- use this facility.
module Futhark.IR.Traversals
  ( -- * Mapping
    Mapper (..),
    identityMapper,
    mapExpM,
    mapExp,

    -- * Walking
    Walker (..),
    identityWalker,
    walkExpM,

    -- * Ops
    TraverseOpStms (..),
    OpStmsTraverser,
    traverseLambdaStms,
  )
where

import Control.Monad
import Control.Monad.Identity
import Data.Bitraversable
import Data.Foldable (traverse_)
import Data.List.NonEmpty (NonEmpty (..))
import Futhark.IR.Prop.Scope
import Futhark.IR.Prop.Types (mapOnType)
import Futhark.IR.Syntax

-- | Express a monad mapping operation on a syntax node.  Each element
-- of this structure expresses the operation to be performed on a
-- given child.
data Mapper frep trep m = Mapper
  { forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp :: SubExp -> m SubExp,
    -- | Most bodies are enclosed in a scope, which is passed along
    -- for convenience.
    forall frep trep (m :: * -> *).
Mapper frep trep m -> Scope trep -> Body frep -> m (Body trep)
mapOnBody :: Scope trep -> Body frep -> m (Body trep),
    forall frep trep (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName :: VName -> m VName,
    forall frep trep (m :: * -> *).
Mapper frep trep m -> RetType frep -> m (RetType trep)
mapOnRetType :: RetType frep -> m (RetType trep),
    forall frep trep (m :: * -> *).
Mapper frep trep m -> BranchType frep -> m (BranchType trep)
mapOnBranchType :: BranchType frep -> m (BranchType trep),
    forall frep trep (m :: * -> *).
Mapper frep trep m -> FParam frep -> m (FParam trep)
mapOnFParam :: FParam frep -> m (FParam trep),
    forall frep trep (m :: * -> *).
Mapper frep trep m -> LParam frep -> m (LParam trep)
mapOnLParam :: LParam frep -> m (LParam trep),
    forall frep trep (m :: * -> *).
Mapper frep trep m -> Op frep -> m (Op trep)
mapOnOp :: Op frep -> m (Op trep)
  }

-- | A mapper that simply returns the tree verbatim.
identityMapper :: forall rep m. (Monad m) => Mapper rep rep m
identityMapper :: forall rep (m :: * -> *). Monad m => Mapper rep rep m
identityMapper =
  Mapper
    { mapOnSubExp :: SubExp -> m SubExp
mapOnSubExp = SubExp -> m SubExp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
      mapOnBody :: Scope rep -> Body rep -> m (Body rep)
mapOnBody = (Body rep -> m (Body rep)) -> Scope rep -> Body rep -> m (Body rep)
forall a b. a -> b -> a
const Body rep -> m (Body rep)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
      mapOnVName :: VName -> m VName
mapOnVName = VName -> m VName
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
      mapOnRetType :: RetType rep -> m (RetType rep)
mapOnRetType = RetType rep -> m (RetType rep)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
      mapOnBranchType :: BranchType rep -> m (BranchType rep)
mapOnBranchType = BranchType rep -> m (BranchType rep)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
      mapOnFParam :: FParam rep -> m (FParam rep)
mapOnFParam = FParam rep -> m (FParam rep)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
      mapOnLParam :: LParam rep -> m (LParam rep)
mapOnLParam = LParam rep -> m (LParam rep)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
      mapOnOp :: Op rep -> m (Op rep)
mapOnOp = Op rep -> m (Op rep)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    }

-- | Map a monadic action across the immediate children of an
-- expression.  Importantly, the mapping does not descend recursively
-- into subexpressions.  The mapping is done left-to-right.
mapExpM ::
  (Monad m) =>
  Mapper frep trep m ->
  Exp frep ->
  m (Exp trep)
mapExpM :: forall (m :: * -> *) frep trep.
Monad m =>
Mapper frep trep m -> Exp frep -> m (Exp trep)
mapExpM Mapper frep trep m
tv (BasicOp (SubExp SubExp
se)) =
  BasicOp -> Exp trep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp trep) -> m BasicOp -> m (Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> m SubExp -> m BasicOp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
se)
mapExpM Mapper frep trep m
tv (BasicOp (ArrayLit [SubExp]
els Type
rowt)) =
  BasicOp -> Exp trep
forall rep. BasicOp -> Exp rep
BasicOp
    (BasicOp -> Exp trep) -> m BasicOp -> m (Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( [SubExp] -> Type -> BasicOp
ArrayLit
            ([SubExp] -> Type -> BasicOp) -> m [SubExp] -> m (Type -> BasicOp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> m SubExp) -> [SubExp] -> m [SubExp]
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 (Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) [SubExp]
els
            m (Type -> BasicOp) -> m Type -> m BasicOp
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (SubExp -> m SubExp) -> Type -> m Type
forall (m :: * -> *) u.
Monad m =>
(SubExp -> m SubExp) -> TypeBase Shape u -> m (TypeBase Shape u)
mapOnType (Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) Type
rowt
        )
mapExpM Mapper frep trep m
tv (BasicOp (BinOp BinOp
bop SubExp
x SubExp
y)) =
  BasicOp -> Exp trep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp trep) -> m BasicOp -> m (Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (BinOp -> SubExp -> SubExp -> BasicOp
BinOp BinOp
bop (SubExp -> SubExp -> BasicOp) -> m SubExp -> m (SubExp -> BasicOp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
x m (SubExp -> BasicOp) -> m SubExp -> m BasicOp
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
y)
mapExpM Mapper frep trep m
tv (BasicOp (CmpOp CmpOp
op SubExp
x SubExp
y)) =
  BasicOp -> Exp trep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp trep) -> m BasicOp -> m (Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp CmpOp
op (SubExp -> SubExp -> BasicOp) -> m SubExp -> m (SubExp -> BasicOp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
x m (SubExp -> BasicOp) -> m SubExp -> m BasicOp
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
y)
mapExpM Mapper frep trep m
tv (BasicOp (ConvOp ConvOp
conv SubExp
x)) =
  BasicOp -> Exp trep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp trep) -> m BasicOp -> m (Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ConvOp -> SubExp -> BasicOp
ConvOp ConvOp
conv (SubExp -> BasicOp) -> m SubExp -> m BasicOp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
x)
mapExpM Mapper frep trep m
tv (BasicOp (UnOp UnOp
op SubExp
x)) =
  BasicOp -> Exp trep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp trep) -> m BasicOp -> m (Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (UnOp -> SubExp -> BasicOp
UnOp UnOp
op (SubExp -> BasicOp) -> m SubExp -> m BasicOp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
x)
mapExpM Mapper frep trep m
tv (Match [SubExp]
ses [Case (Body frep)]
cases Body frep
defbody (MatchDec [BranchType frep]
ts MatchSort
s)) =
  [SubExp]
-> [Case (Body trep)]
-> Body trep
-> MatchDec (BranchType trep)
-> Exp trep
forall rep.
[SubExp]
-> [Case (Body rep)]
-> Body rep
-> MatchDec (BranchType rep)
-> Exp rep
Match
    ([SubExp]
 -> [Case (Body trep)]
 -> Body trep
 -> MatchDec (BranchType trep)
 -> Exp trep)
-> m [SubExp]
-> m ([Case (Body trep)]
      -> Body trep -> MatchDec (BranchType trep) -> Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> m SubExp) -> [SubExp] -> m [SubExp]
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 (Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) [SubExp]
ses
    m ([Case (Body trep)]
   -> Body trep -> MatchDec (BranchType trep) -> Exp trep)
-> m [Case (Body trep)]
-> m (Body trep -> MatchDec (BranchType trep) -> Exp trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Case (Body frep) -> m (Case (Body trep)))
-> [Case (Body frep)] -> m [Case (Body trep)]
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 Case (Body frep) -> m (Case (Body trep))
mapOnCase [Case (Body frep)]
cases
    m (Body trep -> MatchDec (BranchType trep) -> Exp trep)
-> m (Body trep) -> m (MatchDec (BranchType trep) -> Exp trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper frep trep m -> Scope trep -> Body frep -> m (Body trep)
forall frep trep (m :: * -> *).
Mapper frep trep m -> Scope trep -> Body frep -> m (Body trep)
mapOnBody Mapper frep trep m
tv Scope trep
forall a. Monoid a => a
mempty Body frep
defbody
    m (MatchDec (BranchType trep) -> Exp trep)
-> m (MatchDec (BranchType trep)) -> m (Exp trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ([BranchType trep] -> MatchSort -> MatchDec (BranchType trep)
forall rt. [rt] -> MatchSort -> MatchDec rt
MatchDec ([BranchType trep] -> MatchSort -> MatchDec (BranchType trep))
-> m [BranchType trep]
-> m (MatchSort -> MatchDec (BranchType trep))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (BranchType frep -> m (BranchType trep))
-> [BranchType frep] -> m [BranchType trep]
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 (Mapper frep trep m -> BranchType frep -> m (BranchType trep)
forall frep trep (m :: * -> *).
Mapper frep trep m -> BranchType frep -> m (BranchType trep)
mapOnBranchType Mapper frep trep m
tv) [BranchType frep]
ts m (MatchSort -> MatchDec (BranchType trep))
-> m MatchSort -> m (MatchDec (BranchType trep))
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> MatchSort -> m MatchSort
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MatchSort
s)
  where
    mapOnCase :: Case (Body frep) -> m (Case (Body trep))
mapOnCase (Case [Maybe PrimValue]
vs Body frep
body) = [Maybe PrimValue] -> Body trep -> Case (Body trep)
forall body. [Maybe PrimValue] -> body -> Case body
Case [Maybe PrimValue]
vs (Body trep -> Case (Body trep))
-> m (Body trep) -> m (Case (Body trep))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> Scope trep -> Body frep -> m (Body trep)
forall frep trep (m :: * -> *).
Mapper frep trep m -> Scope trep -> Body frep -> m (Body trep)
mapOnBody Mapper frep trep m
tv Scope trep
forall a. Monoid a => a
mempty Body frep
body
mapExpM Mapper frep trep m
tv (Apply Name
fname [(SubExp, Diet)]
args [(RetType frep, RetAls)]
ret (Safety, SrcLoc, [SrcLoc])
loc) = do
  [(SubExp, Diet)]
args' <- [(SubExp, Diet)]
-> ((SubExp, Diet) -> m (SubExp, Diet)) -> m [(SubExp, Diet)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(SubExp, Diet)]
args (((SubExp, Diet) -> m (SubExp, Diet)) -> m [(SubExp, Diet)])
-> ((SubExp, Diet) -> m (SubExp, Diet)) -> m [(SubExp, Diet)]
forall a b. (a -> b) -> a -> b
$ \(SubExp
arg, Diet
d) ->
    (,) (SubExp -> Diet -> (SubExp, Diet))
-> m SubExp -> m (Diet -> (SubExp, Diet))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
arg m (Diet -> (SubExp, Diet)) -> m Diet -> m (SubExp, Diet)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Diet -> m Diet
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Diet
d
  Name
-> [(SubExp, Diet)]
-> [(RetType trep, RetAls)]
-> (Safety, SrcLoc, [SrcLoc])
-> Exp trep
forall rep.
Name
-> [(SubExp, Diet)]
-> [(RetType rep, RetAls)]
-> (Safety, SrcLoc, [SrcLoc])
-> Exp rep
Apply Name
fname [(SubExp, Diet)]
args' ([(RetType trep, RetAls)]
 -> (Safety, SrcLoc, [SrcLoc]) -> Exp trep)
-> m [(RetType trep, RetAls)]
-> m ((Safety, SrcLoc, [SrcLoc]) -> Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((RetType frep, RetAls) -> m (RetType trep, RetAls))
-> [(RetType frep, RetAls)] -> m [(RetType trep, RetAls)]
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 ((RetType frep -> m (RetType trep))
-> (RetAls -> m RetAls)
-> (RetType frep, RetAls)
-> m (RetType trep, RetAls)
forall (f :: * -> *) a c b d.
Applicative f =>
(a -> f c) -> (b -> f d) -> (a, b) -> f (c, d)
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse (Mapper frep trep m -> RetType frep -> m (RetType trep)
forall frep trep (m :: * -> *).
Mapper frep trep m -> RetType frep -> m (RetType trep)
mapOnRetType Mapper frep trep m
tv) RetAls -> m RetAls
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure) [(RetType frep, RetAls)]
ret m ((Safety, SrcLoc, [SrcLoc]) -> Exp trep)
-> m (Safety, SrcLoc, [SrcLoc]) -> m (Exp trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Safety, SrcLoc, [SrcLoc]) -> m (Safety, SrcLoc, [SrcLoc])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Safety, SrcLoc, [SrcLoc])
loc
mapExpM Mapper frep trep m
tv (BasicOp (Index VName
arr Slice SubExp
slice)) =
  BasicOp -> Exp trep
forall rep. BasicOp -> Exp rep
BasicOp
    (BasicOp -> Exp trep) -> m BasicOp -> m (Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( VName -> Slice SubExp -> BasicOp
Index
            (VName -> Slice SubExp -> BasicOp)
-> m VName -> m (Slice SubExp -> BasicOp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> VName -> m VName
forall frep trep (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
arr
            m (Slice SubExp -> BasicOp) -> m (Slice SubExp) -> m BasicOp
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (SubExp -> m SubExp) -> Slice SubExp -> m (Slice SubExp)
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) -> Slice a -> f (Slice b)
traverse (Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) Slice SubExp
slice
        )
mapExpM Mapper frep trep m
tv (BasicOp (Update Safety
safety VName
arr Slice SubExp
slice SubExp
se)) =
  BasicOp -> Exp trep
forall rep. BasicOp -> Exp rep
BasicOp
    (BasicOp -> Exp trep) -> m BasicOp -> m (Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( Safety -> VName -> Slice SubExp -> SubExp -> BasicOp
Update Safety
safety
            (VName -> Slice SubExp -> SubExp -> BasicOp)
-> m VName -> m (Slice SubExp -> SubExp -> BasicOp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> VName -> m VName
forall frep trep (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
arr
            m (Slice SubExp -> SubExp -> BasicOp)
-> m (Slice SubExp) -> m (SubExp -> BasicOp)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (SubExp -> m SubExp) -> Slice SubExp -> m (Slice SubExp)
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) -> Slice a -> f (Slice b)
traverse (Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) Slice SubExp
slice
            m (SubExp -> BasicOp) -> m SubExp -> m BasicOp
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
se
        )
mapExpM Mapper frep trep m
tv (BasicOp (FlatIndex VName
arr FlatSlice SubExp
slice)) =
  BasicOp -> Exp trep
forall rep. BasicOp -> Exp rep
BasicOp
    (BasicOp -> Exp trep) -> m BasicOp -> m (Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( VName -> FlatSlice SubExp -> BasicOp
FlatIndex
            (VName -> FlatSlice SubExp -> BasicOp)
-> m VName -> m (FlatSlice SubExp -> BasicOp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> VName -> m VName
forall frep trep (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
arr
            m (FlatSlice SubExp -> BasicOp)
-> m (FlatSlice SubExp) -> m BasicOp
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (SubExp -> m SubExp) -> FlatSlice SubExp -> m (FlatSlice SubExp)
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) -> FlatSlice a -> f (FlatSlice b)
traverse (Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) FlatSlice SubExp
slice
        )
mapExpM Mapper frep trep m
tv (BasicOp (FlatUpdate VName
arr FlatSlice SubExp
slice VName
se)) =
  BasicOp -> Exp trep
forall rep. BasicOp -> Exp rep
BasicOp
    (BasicOp -> Exp trep) -> m BasicOp -> m (Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( VName -> FlatSlice SubExp -> VName -> BasicOp
FlatUpdate
            (VName -> FlatSlice SubExp -> VName -> BasicOp)
-> m VName -> m (FlatSlice SubExp -> VName -> BasicOp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> VName -> m VName
forall frep trep (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
arr
            m (FlatSlice SubExp -> VName -> BasicOp)
-> m (FlatSlice SubExp) -> m (VName -> BasicOp)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (SubExp -> m SubExp) -> FlatSlice SubExp -> m (FlatSlice SubExp)
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) -> FlatSlice a -> f (FlatSlice b)
traverse (Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) FlatSlice SubExp
slice
            m (VName -> BasicOp) -> m VName -> m BasicOp
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper frep trep m -> VName -> m VName
forall frep trep (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
se
        )
mapExpM Mapper frep trep m
tv (BasicOp (Iota SubExp
n SubExp
x SubExp
s IntType
et)) =
  BasicOp -> Exp trep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp trep) -> m BasicOp -> m (Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> SubExp -> SubExp -> IntType -> BasicOp
Iota (SubExp -> SubExp -> SubExp -> IntType -> BasicOp)
-> m SubExp -> m (SubExp -> SubExp -> IntType -> BasicOp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
n m (SubExp -> SubExp -> IntType -> BasicOp)
-> m SubExp -> m (SubExp -> IntType -> BasicOp)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
x m (SubExp -> IntType -> BasicOp)
-> m SubExp -> m (IntType -> BasicOp)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
s m (IntType -> BasicOp) -> m IntType -> m BasicOp
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> IntType -> m IntType
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IntType
et)
mapExpM Mapper frep trep m
tv (BasicOp (Replicate Shape
shape SubExp
vexp)) =
  BasicOp -> Exp trep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp trep) -> m BasicOp -> m (Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Shape -> SubExp -> BasicOp
Replicate (Shape -> SubExp -> BasicOp) -> m Shape -> m (SubExp -> BasicOp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> Shape -> m Shape
forall (m :: * -> *) frep trep.
Monad m =>
Mapper frep trep m -> Shape -> m Shape
mapOnShape Mapper frep trep m
tv Shape
shape m (SubExp -> BasicOp) -> m SubExp -> m BasicOp
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
vexp)
mapExpM Mapper frep trep m
tv (BasicOp (Scratch PrimType
t [SubExp]
shape)) =
  BasicOp -> Exp trep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp trep) -> m BasicOp -> m (Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PrimType -> [SubExp] -> BasicOp
Scratch PrimType
t ([SubExp] -> BasicOp) -> m [SubExp] -> m BasicOp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> m SubExp) -> [SubExp] -> m [SubExp]
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 (Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) [SubExp]
shape)
mapExpM Mapper frep trep m
tv (BasicOp (Reshape ReshapeKind
kind Shape
shape VName
arrexp)) =
  BasicOp -> Exp trep
forall rep. BasicOp -> Exp rep
BasicOp
    (BasicOp -> Exp trep) -> m BasicOp -> m (Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( ReshapeKind -> Shape -> VName -> BasicOp
Reshape ReshapeKind
kind
            (Shape -> VName -> BasicOp) -> m Shape -> m (VName -> BasicOp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> m SubExp) -> Shape -> m Shape
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) -> ShapeBase a -> m (ShapeBase b)
mapM (Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) Shape
shape
            m (VName -> BasicOp) -> m VName -> m BasicOp
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper frep trep m -> VName -> m VName
forall frep trep (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
arrexp
        )
mapExpM Mapper frep trep m
tv (BasicOp (Rearrange [Int]
perm VName
e)) =
  BasicOp -> Exp trep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp trep) -> m BasicOp -> m (Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([Int] -> VName -> BasicOp
Rearrange [Int]
perm (VName -> BasicOp) -> m VName -> m BasicOp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> VName -> m VName
forall frep trep (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
e)
mapExpM Mapper frep trep m
tv (BasicOp (Concat Int
i (VName
x :| [VName]
ys) SubExp
size)) = do
  VName
x' <- Mapper frep trep m -> VName -> m VName
forall frep trep (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
x
  [VName]
ys' <- (VName -> m VName) -> [VName] -> m [VName]
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 (Mapper frep trep m -> VName -> m VName
forall frep trep (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv) [VName]
ys
  SubExp
size' <- Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
size
  Exp trep -> m (Exp trep)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp trep -> m (Exp trep)) -> Exp trep -> m (Exp trep)
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp trep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp trep) -> BasicOp -> Exp trep
forall a b. (a -> b) -> a -> b
$ Int -> NonEmpty VName -> SubExp -> BasicOp
Concat Int
i (VName
x' VName -> [VName] -> NonEmpty VName
forall a. a -> [a] -> NonEmpty a
:| [VName]
ys') SubExp
size'
mapExpM Mapper frep trep m
tv (BasicOp (Manifest [Int]
perm VName
e)) =
  BasicOp -> Exp trep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp trep) -> m BasicOp -> m (Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([Int] -> VName -> BasicOp
Manifest [Int]
perm (VName -> BasicOp) -> m VName -> m BasicOp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> VName -> m VName
forall frep trep (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
e)
mapExpM Mapper frep trep m
tv (BasicOp (Assert SubExp
e ErrorMsg SubExp
msg (SrcLoc, [SrcLoc])
loc)) =
  BasicOp -> Exp trep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp trep) -> m BasicOp -> m (Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> ErrorMsg SubExp -> (SrcLoc, [SrcLoc]) -> BasicOp
Assert (SubExp -> ErrorMsg SubExp -> (SrcLoc, [SrcLoc]) -> BasicOp)
-> m SubExp -> m (ErrorMsg SubExp -> (SrcLoc, [SrcLoc]) -> BasicOp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
e m (ErrorMsg SubExp -> (SrcLoc, [SrcLoc]) -> BasicOp)
-> m (ErrorMsg SubExp) -> m ((SrcLoc, [SrcLoc]) -> BasicOp)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (SubExp -> m SubExp) -> ErrorMsg SubExp -> m (ErrorMsg SubExp)
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) -> ErrorMsg a -> f (ErrorMsg b)
traverse (Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) ErrorMsg SubExp
msg m ((SrcLoc, [SrcLoc]) -> BasicOp)
-> m (SrcLoc, [SrcLoc]) -> m BasicOp
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (SrcLoc, [SrcLoc]) -> m (SrcLoc, [SrcLoc])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SrcLoc, [SrcLoc])
loc)
mapExpM Mapper frep trep m
tv (BasicOp (Opaque OpaqueOp
op SubExp
e)) =
  BasicOp -> Exp trep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp trep) -> m BasicOp -> m (Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (OpaqueOp -> SubExp -> BasicOp
Opaque OpaqueOp
op (SubExp -> BasicOp) -> m SubExp -> m BasicOp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
e)
mapExpM Mapper frep trep m
tv (BasicOp (UpdateAcc VName
v [SubExp]
is [SubExp]
ses)) =
  BasicOp -> Exp trep
forall rep. BasicOp -> Exp rep
BasicOp
    (BasicOp -> Exp trep) -> m BasicOp -> m (Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( VName -> [SubExp] -> [SubExp] -> BasicOp
UpdateAcc
            (VName -> [SubExp] -> [SubExp] -> BasicOp)
-> m VName -> m ([SubExp] -> [SubExp] -> BasicOp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> VName -> m VName
forall frep trep (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
v
            m ([SubExp] -> [SubExp] -> BasicOp)
-> m [SubExp] -> m ([SubExp] -> BasicOp)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (SubExp -> m SubExp) -> [SubExp] -> m [SubExp]
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 (Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) [SubExp]
is
            m ([SubExp] -> BasicOp) -> m [SubExp] -> m BasicOp
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (SubExp -> m SubExp) -> [SubExp] -> m [SubExp]
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 (Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) [SubExp]
ses
        )
mapExpM Mapper frep trep m
tv (WithAcc [WithAccInput frep]
inputs Lambda frep
lam) =
  [WithAccInput trep] -> Lambda trep -> Exp trep
forall rep. [WithAccInput rep] -> Lambda rep -> Exp rep
WithAcc ([WithAccInput trep] -> Lambda trep -> Exp trep)
-> m [WithAccInput trep] -> m (Lambda trep -> Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (WithAccInput frep -> m (WithAccInput trep))
-> [WithAccInput frep] -> m [WithAccInput trep]
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 WithAccInput frep -> m (WithAccInput trep)
forall {t :: * -> * -> *} {t :: * -> *} {t :: * -> *}
       {t :: * -> *}.
(Bitraversable t, Traversable t, Traversable t, Traversable t) =>
(Shape, t VName, t (t (Lambda frep) (t SubExp)))
-> m (Shape, t VName, t (t (Lambda trep) (t SubExp)))
onInput [WithAccInput frep]
inputs m (Lambda trep -> Exp trep) -> m (Lambda trep) -> m (Exp trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper frep trep m -> Lambda frep -> m (Lambda trep)
forall (m :: * -> *) frep trep.
Monad m =>
Mapper frep trep m -> Lambda frep -> m (Lambda trep)
mapOnLambda Mapper frep trep m
tv Lambda frep
lam
  where
    onInput :: (Shape, t VName, t (t (Lambda frep) (t SubExp)))
-> m (Shape, t VName, t (t (Lambda trep) (t SubExp)))
onInput (Shape
shape, t VName
vs, t (t (Lambda frep) (t SubExp))
op) =
      (,,)
        (Shape
 -> t VName
 -> t (t (Lambda trep) (t SubExp))
 -> (Shape, t VName, t (t (Lambda trep) (t SubExp))))
-> m Shape
-> m (t VName
      -> t (t (Lambda trep) (t SubExp))
      -> (Shape, t VName, t (t (Lambda trep) (t SubExp))))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> Shape -> m Shape
forall (m :: * -> *) frep trep.
Monad m =>
Mapper frep trep m -> Shape -> m Shape
mapOnShape Mapper frep trep m
tv Shape
shape
        m (t VName
   -> t (t (Lambda trep) (t SubExp))
   -> (Shape, t VName, t (t (Lambda trep) (t SubExp))))
-> m (t VName)
-> m (t (t (Lambda trep) (t SubExp))
      -> (Shape, t VName, t (t (Lambda trep) (t SubExp))))
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (VName -> m VName) -> t VName -> m (t VName)
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) -> t a -> m (t b)
mapM (Mapper frep trep m -> VName -> m VName
forall frep trep (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv) t VName
vs
        m (t (t (Lambda trep) (t SubExp))
   -> (Shape, t VName, t (t (Lambda trep) (t SubExp))))
-> m (t (t (Lambda trep) (t SubExp)))
-> m (Shape, t VName, t (t (Lambda trep) (t SubExp)))
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (t (Lambda frep) (t SubExp) -> m (t (Lambda trep) (t SubExp)))
-> t (t (Lambda frep) (t SubExp))
-> m (t (t (Lambda trep) (t SubExp)))
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) -> t a -> f (t b)
traverse ((Lambda frep -> m (Lambda trep))
-> (t SubExp -> m (t SubExp))
-> t (Lambda frep) (t SubExp)
-> m (t (Lambda trep) (t SubExp))
forall (f :: * -> *) a c b d.
Applicative f =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse (Mapper frep trep m -> Lambda frep -> m (Lambda trep)
forall (m :: * -> *) frep trep.
Monad m =>
Mapper frep trep m -> Lambda frep -> m (Lambda trep)
mapOnLambda Mapper frep trep m
tv) ((SubExp -> m SubExp) -> t SubExp -> m (t SubExp)
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) -> t a -> m (t b)
mapM (Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv))) t (t (Lambda frep) (t SubExp))
op
mapExpM Mapper frep trep m
tv (Loop [(FParam frep, SubExp)]
merge LoopForm
form Body frep
loopbody) = do
  [Param (FParamInfo trep)]
params' <- (FParam frep -> m (Param (FParamInfo trep)))
-> [FParam frep] -> m [Param (FParamInfo trep)]
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 (Mapper frep trep m -> FParam frep -> m (Param (FParamInfo trep))
forall frep trep (m :: * -> *).
Mapper frep trep m -> FParam frep -> m (FParam trep)
mapOnFParam Mapper frep trep m
tv) [FParam frep]
params
  LoopForm
form' <- Mapper frep trep m -> LoopForm -> m LoopForm
forall (m :: * -> *) frep trep.
Monad m =>
Mapper frep trep m -> LoopForm -> m LoopForm
mapOnLoopForm Mapper frep trep m
tv LoopForm
form
  let scope :: Scope trep
scope = LoopForm -> Scope trep
forall rep. LoopForm -> Scope rep
scopeOfLoopForm LoopForm
form' Scope trep -> Scope trep -> Scope trep
forall a. Semigroup a => a -> a -> a
<> [Param (FParamInfo trep)] -> Scope trep
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams [Param (FParamInfo trep)]
params'
  [(Param (FParamInfo trep), SubExp)]
-> LoopForm -> Body trep -> Exp trep
forall rep.
[(FParam rep, SubExp)] -> LoopForm -> Body rep -> Exp rep
Loop
    ([(Param (FParamInfo trep), SubExp)]
 -> LoopForm -> Body trep -> Exp trep)
-> m [(Param (FParamInfo trep), SubExp)]
-> m (LoopForm -> Body trep -> Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([Param (FParamInfo trep)]
-> [SubExp] -> [(Param (FParamInfo trep), SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param (FParamInfo trep)]
params' ([SubExp] -> [(Param (FParamInfo trep), SubExp)])
-> m [SubExp] -> m [(Param (FParamInfo trep), SubExp)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> m SubExp) -> [SubExp] -> m [SubExp]
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 (Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) [SubExp]
args)
    m (LoopForm -> Body trep -> Exp trep)
-> m LoopForm -> m (Body trep -> Exp trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> LoopForm -> m LoopForm
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure LoopForm
form'
    m (Body trep -> Exp trep) -> m (Body trep) -> m (Exp trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper frep trep m -> Scope trep -> Body frep -> m (Body trep)
forall frep trep (m :: * -> *).
Mapper frep trep m -> Scope trep -> Body frep -> m (Body trep)
mapOnBody Mapper frep trep m
tv Scope trep
scope Body frep
loopbody
  where
    ([FParam frep]
params, [SubExp]
args) = [(FParam frep, SubExp)] -> ([FParam frep], [SubExp])
forall a b. [(a, b)] -> ([a], [b])
unzip [(FParam frep, SubExp)]
merge
mapExpM Mapper frep trep m
tv (Op Op frep
op) =
  Op trep -> Exp trep
forall rep. Op rep -> Exp rep
Op (Op trep -> Exp trep) -> m (Op trep) -> m (Exp trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> Op frep -> m (Op trep)
forall frep trep (m :: * -> *).
Mapper frep trep m -> Op frep -> m (Op trep)
mapOnOp Mapper frep trep m
tv Op frep
op

mapOnShape :: (Monad m) => Mapper frep trep m -> Shape -> m Shape
mapOnShape :: forall (m :: * -> *) frep trep.
Monad m =>
Mapper frep trep m -> Shape -> m Shape
mapOnShape Mapper frep trep m
tv (Shape [SubExp]
ds) = [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape ([SubExp] -> Shape) -> m [SubExp] -> m Shape
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> m SubExp) -> [SubExp] -> m [SubExp]
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 (Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) [SubExp]
ds

mapOnLoopForm ::
  (Monad m) =>
  Mapper frep trep m ->
  LoopForm ->
  m LoopForm
mapOnLoopForm :: forall (m :: * -> *) frep trep.
Monad m =>
Mapper frep trep m -> LoopForm -> m LoopForm
mapOnLoopForm Mapper frep trep m
tv (ForLoop VName
i IntType
it SubExp
bound) =
  VName -> IntType -> SubExp -> LoopForm
ForLoop (VName -> IntType -> SubExp -> LoopForm)
-> m VName -> m (IntType -> SubExp -> LoopForm)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> VName -> m VName
forall frep trep (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
i m (IntType -> SubExp -> LoopForm)
-> m IntType -> m (SubExp -> LoopForm)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> IntType -> m IntType
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IntType
it m (SubExp -> LoopForm) -> m SubExp -> m LoopForm
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
bound
mapOnLoopForm Mapper frep trep m
tv (WhileLoop VName
cond) =
  VName -> LoopForm
WhileLoop (VName -> LoopForm) -> m VName -> m LoopForm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper frep trep m -> VName -> m VName
forall frep trep (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
cond

mapOnLambda ::
  (Monad m) =>
  Mapper frep trep m ->
  Lambda frep ->
  m (Lambda trep)
mapOnLambda :: forall (m :: * -> *) frep trep.
Monad m =>
Mapper frep trep m -> Lambda frep -> m (Lambda trep)
mapOnLambda Mapper frep trep m
tv (Lambda [LParam frep]
params [Type]
ret Body frep
body) = do
  [Param (LParamInfo trep)]
params' <- (LParam frep -> m (Param (LParamInfo trep)))
-> [LParam frep] -> m [Param (LParamInfo trep)]
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 (Mapper frep trep m -> LParam frep -> m (Param (LParamInfo trep))
forall frep trep (m :: * -> *).
Mapper frep trep m -> LParam frep -> m (LParam trep)
mapOnLParam Mapper frep trep m
tv) [LParam frep]
params
  [Param (LParamInfo trep)] -> [Type] -> Body trep -> Lambda trep
forall rep. [LParam rep] -> [Type] -> Body rep -> Lambda rep
Lambda [Param (LParamInfo trep)]
params'
    ([Type] -> Body trep -> Lambda trep)
-> m [Type] -> m (Body trep -> Lambda trep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Type -> m Type) -> [Type] -> m [Type]
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 ((SubExp -> m SubExp) -> Type -> m Type
forall (m :: * -> *) u.
Monad m =>
(SubExp -> m SubExp) -> TypeBase Shape u -> m (TypeBase Shape u)
mapOnType (Mapper frep trep m -> SubExp -> m SubExp
forall frep trep (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv)) [Type]
ret
    m (Body trep -> Lambda trep) -> m (Body trep) -> m (Lambda trep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper frep trep m -> Scope trep -> Body frep -> m (Body trep)
forall frep trep (m :: * -> *).
Mapper frep trep m -> Scope trep -> Body frep -> m (Body trep)
mapOnBody Mapper frep trep m
tv ([Param (LParamInfo trep)] -> Scope trep
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams [Param (LParamInfo trep)]
params') Body frep
body

-- | Like 'mapExpM', but in the 'Identity' monad.
mapExp :: Mapper frep trep Identity -> Exp frep -> Exp trep
mapExp :: forall frep trep. Mapper frep trep Identity -> Exp frep -> Exp trep
mapExp Mapper frep trep Identity
m = Identity (Exp trep) -> Exp trep
forall a. Identity a -> a
runIdentity (Identity (Exp trep) -> Exp trep)
-> (Exp frep -> Identity (Exp trep)) -> Exp frep -> Exp trep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mapper frep trep Identity -> Exp frep -> Identity (Exp trep)
forall (m :: * -> *) frep trep.
Monad m =>
Mapper frep trep m -> Exp frep -> m (Exp trep)
mapExpM Mapper frep trep Identity
m

-- | Express a monad expression on a syntax node.  Each element of
-- this structure expresses the action to be performed on a given
-- child.
data Walker rep m = Walker
  { forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp :: SubExp -> m (),
    forall rep (m :: * -> *).
Walker rep m -> Scope rep -> Body rep -> m ()
walkOnBody :: Scope rep -> Body rep -> m (),
    forall rep (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName :: VName -> m (),
    forall rep (m :: * -> *). Walker rep m -> RetType rep -> m ()
walkOnRetType :: RetType rep -> m (),
    forall rep (m :: * -> *). Walker rep m -> BranchType rep -> m ()
walkOnBranchType :: BranchType rep -> m (),
    forall rep (m :: * -> *). Walker rep m -> FParam rep -> m ()
walkOnFParam :: FParam rep -> m (),
    forall rep (m :: * -> *). Walker rep m -> LParam rep -> m ()
walkOnLParam :: LParam rep -> m (),
    forall rep (m :: * -> *). Walker rep m -> Op rep -> m ()
walkOnOp :: Op rep -> m ()
  }

-- | A no-op traversal.
identityWalker :: forall rep m. (Monad m) => Walker rep m
identityWalker :: forall rep (m :: * -> *). Monad m => Walker rep m
identityWalker =
  Walker
    { walkOnSubExp :: SubExp -> m ()
walkOnSubExp = m () -> SubExp -> m ()
forall a b. a -> b -> a
const (m () -> SubExp -> m ()) -> m () -> SubExp -> m ()
forall a b. (a -> b) -> a -> b
$ () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (),
      walkOnBody :: Scope rep -> Body rep -> m ()
walkOnBody = (Body rep -> m ()) -> Scope rep -> Body rep -> m ()
forall a b. a -> b -> a
const ((Body rep -> m ()) -> Scope rep -> Body rep -> m ())
-> (Body rep -> m ()) -> Scope rep -> Body rep -> m ()
forall a b. (a -> b) -> a -> b
$ m () -> Body rep -> m ()
forall a b. a -> b -> a
const (m () -> Body rep -> m ()) -> m () -> Body rep -> m ()
forall a b. (a -> b) -> a -> b
$ () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (),
      walkOnVName :: VName -> m ()
walkOnVName = m () -> VName -> m ()
forall a b. a -> b -> a
const (m () -> VName -> m ()) -> m () -> VName -> m ()
forall a b. (a -> b) -> a -> b
$ () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (),
      walkOnRetType :: RetType rep -> m ()
walkOnRetType = m () -> RetType rep -> m ()
forall a b. a -> b -> a
const (m () -> RetType rep -> m ()) -> m () -> RetType rep -> m ()
forall a b. (a -> b) -> a -> b
$ () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (),
      walkOnBranchType :: BranchType rep -> m ()
walkOnBranchType = m () -> BranchType rep -> m ()
forall a b. a -> b -> a
const (m () -> BranchType rep -> m ()) -> m () -> BranchType rep -> m ()
forall a b. (a -> b) -> a -> b
$ () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (),
      walkOnFParam :: FParam rep -> m ()
walkOnFParam = m () -> FParam rep -> m ()
forall a b. a -> b -> a
const (m () -> FParam rep -> m ()) -> m () -> FParam rep -> m ()
forall a b. (a -> b) -> a -> b
$ () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (),
      walkOnLParam :: LParam rep -> m ()
walkOnLParam = m () -> LParam rep -> m ()
forall a b. a -> b -> a
const (m () -> LParam rep -> m ()) -> m () -> LParam rep -> m ()
forall a b. (a -> b) -> a -> b
$ () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (),
      walkOnOp :: Op rep -> m ()
walkOnOp = m () -> Op rep -> m ()
forall a b. a -> b -> a
const (m () -> Op rep -> m ()) -> m () -> Op rep -> m ()
forall a b. (a -> b) -> a -> b
$ () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    }

walkOnShape :: (Monad m) => Walker rep m -> Shape -> m ()
walkOnShape :: forall (m :: * -> *) rep. Monad m => Walker rep m -> Shape -> m ()
walkOnShape Walker rep m
tv (Shape [SubExp]
ds) = (SubExp -> m ()) -> [SubExp] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) [SubExp]
ds

walkOnType :: (Monad m) => Walker rep m -> Type -> m ()
walkOnType :: forall (m :: * -> *) rep. Monad m => Walker rep m -> Type -> m ()
walkOnType Walker rep m
_ Prim {} = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
walkOnType Walker rep m
tv (Acc VName
acc Shape
ispace [Type]
ts NoUniqueness
_) = do
  Walker rep m -> VName -> m ()
forall rep (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
acc
  (SubExp -> m ()) -> Shape -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) Shape
ispace
  (Type -> m ()) -> [Type] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker rep m -> Type -> m ()
forall (m :: * -> *) rep. Monad m => Walker rep m -> Type -> m ()
walkOnType Walker rep m
tv) [Type]
ts
walkOnType Walker rep m
_ Mem {} = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
walkOnType Walker rep m
tv (Array PrimType
_ Shape
shape NoUniqueness
_) = Walker rep m -> Shape -> m ()
forall (m :: * -> *) rep. Monad m => Walker rep m -> Shape -> m ()
walkOnShape Walker rep m
tv Shape
shape

walkOnLoopForm :: (Monad m) => Walker rep m -> LoopForm -> m ()
walkOnLoopForm :: forall (m :: * -> *) rep.
Monad m =>
Walker rep m -> LoopForm -> m ()
walkOnLoopForm Walker rep m
tv (ForLoop VName
i IntType
_ SubExp
bound) = do
  Walker rep m -> VName -> m ()
forall rep (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
i
  Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
bound
walkOnLoopForm Walker rep m
tv (WhileLoop VName
cond) =
  Walker rep m -> VName -> m ()
forall rep (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
cond

walkOnLambda :: (Monad m) => Walker rep m -> Lambda rep -> m ()
walkOnLambda :: forall (m :: * -> *) rep.
Monad m =>
Walker rep m -> Lambda rep -> m ()
walkOnLambda Walker rep m
tv (Lambda [LParam rep]
params [Type]
ret Body rep
body) = do
  (LParam rep -> m ()) -> [LParam rep] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker rep m -> LParam rep -> m ()
forall rep (m :: * -> *). Walker rep m -> LParam rep -> m ()
walkOnLParam Walker rep m
tv) [LParam rep]
params
  Walker rep m -> Scope rep -> Body rep -> m ()
forall rep (m :: * -> *).
Walker rep m -> Scope rep -> Body rep -> m ()
walkOnBody Walker rep m
tv ([LParam rep] -> Scope rep
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams [LParam rep]
params) Body rep
body
  (Type -> m ()) -> [Type] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker rep m -> Type -> m ()
forall (m :: * -> *) rep. Monad m => Walker rep m -> Type -> m ()
walkOnType Walker rep m
tv) [Type]
ret

-- | As 'mapExpM', but do not construct a result AST.
walkExpM :: (Monad m) => Walker rep m -> Exp rep -> m ()
walkExpM :: forall (m :: * -> *) rep.
Monad m =>
Walker rep m -> Exp rep -> m ()
walkExpM Walker rep m
tv (BasicOp (SubExp SubExp
se)) =
  Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
se
walkExpM Walker rep m
tv (BasicOp (ArrayLit [SubExp]
els Type
rowt)) =
  (SubExp -> m ()) -> [SubExp] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) [SubExp]
els m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker rep m -> Type -> m ()
forall (m :: * -> *) rep. Monad m => Walker rep m -> Type -> m ()
walkOnType Walker rep m
tv Type
rowt
walkExpM Walker rep m
tv (BasicOp (BinOp BinOp
_ SubExp
x SubExp
y)) =
  Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
x m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
y
walkExpM Walker rep m
tv (BasicOp (CmpOp CmpOp
_ SubExp
x SubExp
y)) =
  Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
x m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
y
walkExpM Walker rep m
tv (BasicOp (ConvOp ConvOp
_ SubExp
x)) =
  Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
x
walkExpM Walker rep m
tv (BasicOp (UnOp UnOp
_ SubExp
x)) =
  Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
x
walkExpM Walker rep m
tv (Match [SubExp]
ses [Case (Body rep)]
cases Body rep
defbody (MatchDec [BranchType rep]
ts MatchSort
_)) = do
  (SubExp -> m ()) -> [SubExp] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) [SubExp]
ses
  (Case (Body rep) -> m ()) -> [Case (Body rep)] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker rep m -> Scope rep -> Body rep -> m ()
forall rep (m :: * -> *).
Walker rep m -> Scope rep -> Body rep -> m ()
walkOnBody Walker rep m
tv Scope rep
forall a. Monoid a => a
mempty (Body rep -> m ())
-> (Case (Body rep) -> Body rep) -> Case (Body rep) -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Case (Body rep) -> Body rep
forall body. Case body -> body
caseBody) [Case (Body rep)]
cases
  Walker rep m -> Scope rep -> Body rep -> m ()
forall rep (m :: * -> *).
Walker rep m -> Scope rep -> Body rep -> m ()
walkOnBody Walker rep m
tv Scope rep
forall a. Monoid a => a
mempty Body rep
defbody
  (BranchType rep -> m ()) -> [BranchType rep] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker rep m -> BranchType rep -> m ()
forall rep (m :: * -> *). Walker rep m -> BranchType rep -> m ()
walkOnBranchType Walker rep m
tv) [BranchType rep]
ts
walkExpM Walker rep m
tv (Apply Name
_ [(SubExp, Diet)]
args [(RetType rep, RetAls)]
ret (Safety, SrcLoc, [SrcLoc])
_) = do
  ((SubExp, Diet) -> m ()) -> [(SubExp, Diet)] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv (SubExp -> m ())
-> ((SubExp, Diet) -> SubExp) -> (SubExp, Diet) -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExp, Diet) -> SubExp
forall a b. (a, b) -> a
fst) [(SubExp, Diet)]
args
  ((RetType rep, RetAls) -> m ()) -> [(RetType rep, RetAls)] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker rep m -> RetType rep -> m ()
forall rep (m :: * -> *). Walker rep m -> RetType rep -> m ()
walkOnRetType Walker rep m
tv (RetType rep -> m ())
-> ((RetType rep, RetAls) -> RetType rep)
-> (RetType rep, RetAls)
-> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (RetType rep, RetAls) -> RetType rep
forall a b. (a, b) -> a
fst) [(RetType rep, RetAls)]
ret
walkExpM Walker rep m
tv (BasicOp (Index VName
arr Slice SubExp
slice)) =
  Walker rep m -> VName -> m ()
forall rep (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
arr m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (SubExp -> m ()) -> Slice SubExp -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) Slice SubExp
slice
walkExpM Walker rep m
tv (BasicOp (Update Safety
_ VName
arr Slice SubExp
slice SubExp
se)) =
  Walker rep m -> VName -> m ()
forall rep (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
arr
    m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (SubExp -> m ()) -> Slice SubExp -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) Slice SubExp
slice
    m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
se
walkExpM Walker rep m
tv (BasicOp (FlatIndex VName
arr FlatSlice SubExp
slice)) =
  Walker rep m -> VName -> m ()
forall rep (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
arr m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (SubExp -> m ()) -> FlatSlice SubExp -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) FlatSlice SubExp
slice
walkExpM Walker rep m
tv (BasicOp (FlatUpdate VName
arr FlatSlice SubExp
slice VName
se)) =
  Walker rep m -> VName -> m ()
forall rep (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
arr
    m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (SubExp -> m ()) -> FlatSlice SubExp -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) FlatSlice SubExp
slice
    m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker rep m -> VName -> m ()
forall rep (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
se
walkExpM Walker rep m
tv (BasicOp (Iota SubExp
n SubExp
x SubExp
s IntType
_)) =
  Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
n m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
x m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
s
walkExpM Walker rep m
tv (BasicOp (Replicate Shape
shape SubExp
vexp)) =
  Walker rep m -> Shape -> m ()
forall (m :: * -> *) rep. Monad m => Walker rep m -> Shape -> m ()
walkOnShape Walker rep m
tv Shape
shape m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
vexp
walkExpM Walker rep m
tv (BasicOp (Scratch PrimType
_ [SubExp]
shape)) =
  (SubExp -> m ()) -> [SubExp] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) [SubExp]
shape
walkExpM Walker rep m
tv (BasicOp (Reshape ReshapeKind
_ Shape
shape VName
arrexp)) =
  (SubExp -> m ()) -> Shape -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) Shape
shape m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker rep m -> VName -> m ()
forall rep (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
arrexp
walkExpM Walker rep m
tv (BasicOp (Rearrange [Int]
_ VName
e)) =
  Walker rep m -> VName -> m ()
forall rep (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
e
walkExpM Walker rep m
tv (BasicOp (Concat Int
_ (VName
x :| [VName]
ys) SubExp
size)) =
  Walker rep m -> VName -> m ()
forall rep (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
x m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (VName -> m ()) -> [VName] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker rep m -> VName -> m ()
forall rep (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv) [VName]
ys m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
size
walkExpM Walker rep m
tv (BasicOp (Manifest [Int]
_ VName
e)) =
  Walker rep m -> VName -> m ()
forall rep (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
e
walkExpM Walker rep m
tv (BasicOp (Assert SubExp
e ErrorMsg SubExp
msg (SrcLoc, [SrcLoc])
_)) =
  Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
e m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (SubExp -> m ()) -> ErrorMsg SubExp -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) ErrorMsg SubExp
msg
walkExpM Walker rep m
tv (BasicOp (Opaque OpaqueOp
_ SubExp
e)) =
  Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
e
walkExpM Walker rep m
tv (BasicOp (UpdateAcc VName
v [SubExp]
is [SubExp]
ses)) = do
  Walker rep m -> VName -> m ()
forall rep (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
v
  (SubExp -> m ()) -> [SubExp] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) [SubExp]
is
  (SubExp -> m ()) -> [SubExp] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) [SubExp]
ses
walkExpM Walker rep m
tv (WithAcc [WithAccInput rep]
inputs Lambda rep
lam) = do
  [WithAccInput rep] -> (WithAccInput rep -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [WithAccInput rep]
inputs ((WithAccInput rep -> m ()) -> m ())
-> (WithAccInput rep -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \(Shape
shape, [VName]
vs, Maybe (Lambda rep, [SubExp])
op) -> do
    Walker rep m -> Shape -> m ()
forall (m :: * -> *) rep. Monad m => Walker rep m -> Shape -> m ()
walkOnShape Walker rep m
tv Shape
shape
    (VName -> m ()) -> [VName] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker rep m -> VName -> m ()
forall rep (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv) [VName]
vs
    ((Lambda rep, [SubExp]) -> m ((), [()]))
-> Maybe (Lambda rep, [SubExp]) -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ ((Lambda rep -> m ())
-> ([SubExp] -> m [()]) -> (Lambda rep, [SubExp]) -> m ((), [()])
forall (f :: * -> *) a c b d.
Applicative f =>
(a -> f c) -> (b -> f d) -> (a, b) -> f (c, d)
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse (Walker rep m -> Lambda rep -> m ()
forall (m :: * -> *) rep.
Monad m =>
Walker rep m -> Lambda rep -> m ()
walkOnLambda Walker rep m
tv) ((SubExp -> m ()) -> [SubExp] -> m [()]
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 (Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv))) Maybe (Lambda rep, [SubExp])
op
  Walker rep m -> Lambda rep -> m ()
forall (m :: * -> *) rep.
Monad m =>
Walker rep m -> Lambda rep -> m ()
walkOnLambda Walker rep m
tv Lambda rep
lam
walkExpM Walker rep m
tv (Loop [(FParam rep, SubExp)]
merge LoopForm
form Body rep
loopbody) = do
  (FParam rep -> m ()) -> [FParam rep] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker rep m -> FParam rep -> m ()
forall rep (m :: * -> *). Walker rep m -> FParam rep -> m ()
walkOnFParam Walker rep m
tv) [FParam rep]
params
  Walker rep m -> LoopForm -> m ()
forall (m :: * -> *) rep.
Monad m =>
Walker rep m -> LoopForm -> m ()
walkOnLoopForm Walker rep m
tv LoopForm
form
  (SubExp -> m ()) -> [SubExp] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker rep m -> SubExp -> m ()
forall rep (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) [SubExp]
args
  let scope :: Scope rep
scope = [FParam rep] -> Scope rep
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams [FParam rep]
params Scope rep -> Scope rep -> Scope rep
forall a. Semigroup a => a -> a -> a
<> LoopForm -> Scope rep
forall rep. LoopForm -> Scope rep
scopeOfLoopForm LoopForm
form
  Walker rep m -> Scope rep -> Body rep -> m ()
forall rep (m :: * -> *).
Walker rep m -> Scope rep -> Body rep -> m ()
walkOnBody Walker rep m
tv Scope rep
scope Body rep
loopbody
  where
    ([FParam rep]
params, [SubExp]
args) = [(FParam rep, SubExp)] -> ([FParam rep], [SubExp])
forall a b. [(a, b)] -> ([a], [b])
unzip [(FParam rep, SubExp)]
merge
walkExpM Walker rep m
tv (Op Op rep
op) =
  Walker rep m -> Op rep -> m ()
forall rep (m :: * -> *). Walker rep m -> Op rep -> m ()
walkOnOp Walker rep m
tv Op rep
op

-- | A function for monadically traversing any sub-statements of the
-- given op for some representation.
type OpStmsTraverser m op rep = (Scope rep -> Stms rep -> m (Stms rep)) -> op -> m op

-- | This representation supports an 'OpStmsTraverser' for its t'Op'.
-- This is used for some simplification rules.
class TraverseOpStms rep where
  -- | Transform every sub-'Stms' of this op.
  traverseOpStms :: (Monad m) => OpStmsTraverser m (Op rep) rep

-- | A helper for defining 'traverseOpStms'.
traverseLambdaStms :: (Monad m) => OpStmsTraverser m (Lambda rep) rep
traverseLambdaStms :: forall (m :: * -> *) rep.
Monad m =>
OpStmsTraverser m (Lambda rep) rep
traverseLambdaStms Scope rep -> Stms rep -> m (Stms rep)
f (Lambda [LParam rep]
ps [Type]
ret (Body BodyDec rep
dec Stms rep
stms Result
res)) =
  [LParam rep] -> [Type] -> Body rep -> Lambda rep
forall rep. [LParam rep] -> [Type] -> Body rep -> Lambda rep
Lambda [LParam rep]
ps [Type]
ret (Body rep -> Lambda rep) -> m (Body rep) -> m (Lambda rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (BodyDec rep -> Stms rep -> Result -> Body rep
forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
Body BodyDec rep
dec (Stms rep -> Result -> Body rep)
-> m (Stms rep) -> m (Result -> Body rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Scope rep -> Stms rep -> m (Stms rep)
f ([LParam rep] -> Scope rep
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams [LParam rep]
ps) Stms rep
stms m (Result -> Body rep) -> m Result -> m (Body rep)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Result -> m Result
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Result
res)