{-# LANGUAGE Safe #-}
module Futhark.IR.Traversals
(
Mapper(..)
, identityMapper
, mapExpM
, mapExp
, mapOnType
, Walker(..)
, identityWalker
, walkExpM
)
where
import Control.Monad
import Control.Monad.Identity
import qualified Data.Traversable
import Data.Foldable (traverse_)
import Futhark.IR.Syntax
import Futhark.IR.Prop.Scope
data Mapper flore tlore m = Mapper {
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp :: SubExp -> m SubExp
, Mapper flore tlore m -> Scope tlore -> Body flore -> m (Body tlore)
mapOnBody :: Scope tlore -> Body flore -> m (Body tlore)
, Mapper flore tlore m -> VName -> m VName
mapOnVName :: VName -> m VName
, Mapper flore tlore m -> RetType flore -> m (RetType tlore)
mapOnRetType :: RetType flore -> m (RetType tlore)
, Mapper flore tlore m -> BranchType flore -> m (BranchType tlore)
mapOnBranchType :: BranchType flore -> m (BranchType tlore)
, Mapper flore tlore m -> FParam flore -> m (FParam tlore)
mapOnFParam :: FParam flore -> m (FParam tlore)
, Mapper flore tlore m -> LParam flore -> m (LParam tlore)
mapOnLParam :: LParam flore -> m (LParam tlore)
, Mapper flore tlore m -> Op flore -> m (Op tlore)
mapOnOp :: Op flore -> m (Op tlore)
}
identityMapper :: Monad m => Mapper lore lore m
identityMapper :: Mapper lore lore m
identityMapper = Mapper :: forall flore tlore (m :: * -> *).
(SubExp -> m SubExp)
-> (Scope tlore -> Body flore -> m (Body tlore))
-> (VName -> m VName)
-> (RetType flore -> m (RetType tlore))
-> (BranchType flore -> m (BranchType tlore))
-> (FParam flore -> m (FParam tlore))
-> (LParam flore -> m (LParam tlore))
-> (Op flore -> m (Op tlore))
-> Mapper flore tlore m
Mapper {
mapOnSubExp :: SubExp -> m SubExp
mapOnSubExp = SubExp -> m SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return
, mapOnBody :: Scope lore -> Body lore -> m (Body lore)
mapOnBody = (Body lore -> m (Body lore))
-> Scope lore -> Body lore -> m (Body lore)
forall a b. a -> b -> a
const Body lore -> m (Body lore)
forall (m :: * -> *) a. Monad m => a -> m a
return
, mapOnVName :: VName -> m VName
mapOnVName = VName -> m VName
forall (m :: * -> *) a. Monad m => a -> m a
return
, mapOnRetType :: RetType lore -> m (RetType lore)
mapOnRetType = RetType lore -> m (RetType lore)
forall (m :: * -> *) a. Monad m => a -> m a
return
, mapOnBranchType :: BranchType lore -> m (BranchType lore)
mapOnBranchType = BranchType lore -> m (BranchType lore)
forall (m :: * -> *) a. Monad m => a -> m a
return
, mapOnFParam :: FParam lore -> m (FParam lore)
mapOnFParam = FParam lore -> m (FParam lore)
forall (m :: * -> *) a. Monad m => a -> m a
return
, mapOnLParam :: LParam lore -> m (LParam lore)
mapOnLParam = LParam lore -> m (LParam lore)
forall (m :: * -> *) a. Monad m => a -> m a
return
, mapOnOp :: Op lore -> m (Op lore)
mapOnOp = Op lore -> m (Op lore)
forall (m :: * -> *) a. Monad m => a -> m a
return
}
mapExpM :: (Applicative m, Monad m) =>
Mapper flore tlore m -> Exp flore -> m (Exp tlore)
mapExpM :: Mapper flore tlore m -> Exp flore -> m (Exp tlore)
mapExpM Mapper flore tlore m
tv (BasicOp (SubExp SubExp
se)) =
BasicOp -> Exp tlore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tlore) -> m BasicOp -> m (Exp tlore)
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 flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv SubExp
se)
mapExpM Mapper flore tlore m
tv (BasicOp (ArrayLit [SubExp]
els Type
rowt)) =
BasicOp -> Exp tlore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tlore) -> m BasicOp -> m (Exp tlore)
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)
mapM (Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv) [SubExp]
els m (Type -> BasicOp) -> m Type -> m BasicOp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
(SubExp -> m SubExp) -> Type -> m Type
forall (m :: * -> *).
Monad m =>
(SubExp -> m SubExp) -> Type -> m Type
mapOnType (Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv) Type
rowt)
mapExpM Mapper flore tlore m
tv (BasicOp (BinOp BinOp
bop SubExp
x SubExp
y)) =
BasicOp -> Exp tlore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tlore) -> m BasicOp -> m (Exp tlore)
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 flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv SubExp
x m (SubExp -> BasicOp) -> m SubExp -> m BasicOp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv SubExp
y)
mapExpM Mapper flore tlore m
tv (BasicOp (CmpOp CmpOp
op SubExp
x SubExp
y)) =
BasicOp -> Exp tlore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tlore) -> m BasicOp -> m (Exp tlore)
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 flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv SubExp
x m (SubExp -> BasicOp) -> m SubExp -> m BasicOp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv SubExp
y)
mapExpM Mapper flore tlore m
tv (BasicOp (ConvOp ConvOp
conv SubExp
x)) =
BasicOp -> Exp tlore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tlore) -> m BasicOp -> m (Exp tlore)
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 flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv SubExp
x)
mapExpM Mapper flore tlore m
tv (BasicOp (UnOp UnOp
op SubExp
x)) =
BasicOp -> Exp tlore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tlore) -> m BasicOp -> m (Exp tlore)
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 flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv SubExp
x)
mapExpM Mapper flore tlore m
tv (If SubExp
c BodyT flore
texp BodyT flore
fexp (IfDec [BranchType flore]
ts IfSort
s)) =
SubExp
-> BodyT tlore
-> BodyT tlore
-> IfDec (BranchType tlore)
-> Exp tlore
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
If (SubExp
-> BodyT tlore
-> BodyT tlore
-> IfDec (BranchType tlore)
-> Exp tlore)
-> m SubExp
-> m (BodyT tlore
-> BodyT tlore -> IfDec (BranchType tlore) -> Exp tlore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv SubExp
c m (BodyT tlore
-> BodyT tlore -> IfDec (BranchType tlore) -> Exp tlore)
-> m (BodyT tlore)
-> m (BodyT tlore -> IfDec (BranchType tlore) -> Exp tlore)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper flore tlore m
-> Scope tlore -> BodyT flore -> m (BodyT tlore)
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> Scope tlore -> Body flore -> m (Body tlore)
mapOnBody Mapper flore tlore m
tv Scope tlore
forall a. Monoid a => a
mempty BodyT flore
texp m (BodyT tlore -> IfDec (BranchType tlore) -> Exp tlore)
-> m (BodyT tlore) -> m (IfDec (BranchType tlore) -> Exp tlore)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper flore tlore m
-> Scope tlore -> BodyT flore -> m (BodyT tlore)
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> Scope tlore -> Body flore -> m (Body tlore)
mapOnBody Mapper flore tlore m
tv Scope tlore
forall a. Monoid a => a
mempty BodyT flore
fexp m (IfDec (BranchType tlore) -> Exp tlore)
-> m (IfDec (BranchType tlore)) -> m (Exp tlore)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
([BranchType tlore] -> IfSort -> IfDec (BranchType tlore)
forall rt. [rt] -> IfSort -> IfDec rt
IfDec ([BranchType tlore] -> IfSort -> IfDec (BranchType tlore))
-> m [BranchType tlore] -> m (IfSort -> IfDec (BranchType tlore))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (BranchType flore -> m (BranchType tlore))
-> [BranchType flore] -> m [BranchType tlore]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Mapper flore tlore m -> BranchType flore -> m (BranchType tlore)
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> BranchType flore -> m (BranchType tlore)
mapOnBranchType Mapper flore tlore m
tv) [BranchType flore]
ts m (IfSort -> IfDec (BranchType tlore))
-> m IfSort -> m (IfDec (BranchType tlore))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> IfSort -> m IfSort
forall (f :: * -> *) a. Applicative f => a -> f a
pure IfSort
s)
mapExpM Mapper flore tlore m
tv (Apply Name
fname [(SubExp, Diet)]
args [RetType flore]
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 flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv SubExp
arg m (Diet -> (SubExp, Diet)) -> m Diet -> m (SubExp, Diet)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Diet -> m Diet
forall (f :: * -> *) a. Applicative f => a -> f a
pure Diet
d
Name
-> [(SubExp, Diet)]
-> [RetType tlore]
-> (Safety, SrcLoc, [SrcLoc])
-> Exp tlore
forall lore.
Name
-> [(SubExp, Diet)]
-> [RetType lore]
-> (Safety, SrcLoc, [SrcLoc])
-> ExpT lore
Apply Name
fname [(SubExp, Diet)]
args' ([RetType tlore] -> (Safety, SrcLoc, [SrcLoc]) -> Exp tlore)
-> m [RetType tlore] -> m ((Safety, SrcLoc, [SrcLoc]) -> Exp tlore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (RetType flore -> m (RetType tlore))
-> [RetType flore] -> m [RetType tlore]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Mapper flore tlore m -> RetType flore -> m (RetType tlore)
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> RetType flore -> m (RetType tlore)
mapOnRetType Mapper flore tlore m
tv) [RetType flore]
ret m ((Safety, SrcLoc, [SrcLoc]) -> Exp tlore)
-> m (Safety, SrcLoc, [SrcLoc]) -> m (Exp tlore)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Safety, SrcLoc, [SrcLoc]) -> m (Safety, SrcLoc, [SrcLoc])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Safety, SrcLoc, [SrcLoc])
loc
mapExpM Mapper flore tlore m
tv (BasicOp (Index VName
arr Slice SubExp
slice)) =
BasicOp -> Exp tlore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tlore) -> m BasicOp -> m (Exp tlore)
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 flore tlore m -> VName -> m VName
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> VName -> m VName
mapOnVName Mapper flore tlore m
tv VName
arr m (Slice SubExp -> BasicOp) -> m (Slice SubExp) -> m BasicOp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
(DimIndex SubExp -> m (DimIndex SubExp))
-> Slice SubExp -> m (Slice SubExp)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((SubExp -> m SubExp) -> DimIndex SubExp -> m (DimIndex SubExp)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv)) Slice SubExp
slice)
mapExpM Mapper flore tlore m
tv (BasicOp (Update VName
arr Slice SubExp
slice SubExp
se)) =
BasicOp -> Exp tlore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tlore) -> m BasicOp -> m (Exp tlore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> Slice SubExp -> SubExp -> BasicOp
Update (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 flore tlore m -> VName -> m VName
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> VName -> m VName
mapOnVName Mapper flore tlore m
tv VName
arr m (Slice SubExp -> SubExp -> BasicOp)
-> m (Slice SubExp) -> m (SubExp -> BasicOp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
(DimIndex SubExp -> m (DimIndex SubExp))
-> Slice SubExp -> m (Slice SubExp)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((SubExp -> m SubExp) -> DimIndex SubExp -> m (DimIndex SubExp)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv)) Slice SubExp
slice m (SubExp -> BasicOp) -> m SubExp -> m BasicOp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv SubExp
se)
mapExpM Mapper flore tlore m
tv (BasicOp (Iota SubExp
n SubExp
x SubExp
s IntType
et)) =
BasicOp -> Exp tlore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tlore) -> m BasicOp -> m (Exp tlore)
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 flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv SubExp
n m (SubExp -> SubExp -> IntType -> BasicOp)
-> m SubExp -> m (SubExp -> IntType -> BasicOp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv SubExp
x m (SubExp -> IntType -> BasicOp)
-> m SubExp -> m (IntType -> BasicOp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv SubExp
s m (IntType -> BasicOp) -> m IntType -> m BasicOp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> IntType -> m IntType
forall (f :: * -> *) a. Applicative f => a -> f a
pure IntType
et)
mapExpM Mapper flore tlore m
tv (BasicOp (Replicate Shape
shape SubExp
vexp)) =
BasicOp -> Exp tlore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tlore) -> m BasicOp -> m (Exp tlore)
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 flore tlore m -> Shape -> m Shape
forall (m :: * -> *) flore tlore.
Monad m =>
Mapper flore tlore m -> Shape -> m Shape
mapOnShape Mapper flore tlore m
tv Shape
shape m (SubExp -> BasicOp) -> m SubExp -> m BasicOp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv SubExp
vexp)
mapExpM Mapper flore tlore m
tv (BasicOp (Repeat [Shape]
shapes Shape
innershape VName
v)) =
BasicOp -> Exp tlore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tlore) -> m BasicOp -> m (Exp tlore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([Shape] -> Shape -> VName -> BasicOp
Repeat ([Shape] -> Shape -> VName -> BasicOp)
-> m [Shape] -> m (Shape -> VName -> BasicOp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Shape -> m Shape) -> [Shape] -> m [Shape]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Mapper flore tlore m -> Shape -> m Shape
forall (m :: * -> *) flore tlore.
Monad m =>
Mapper flore tlore m -> Shape -> m Shape
mapOnShape Mapper flore tlore m
tv) [Shape]
shapes m (Shape -> VName -> BasicOp) -> m Shape -> m (VName -> BasicOp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
Mapper flore tlore m -> Shape -> m Shape
forall (m :: * -> *) flore tlore.
Monad m =>
Mapper flore tlore m -> Shape -> m Shape
mapOnShape Mapper flore tlore m
tv Shape
innershape m (VName -> BasicOp) -> m VName -> m BasicOp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper flore tlore m -> VName -> m VName
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> VName -> m VName
mapOnVName Mapper flore tlore m
tv VName
v)
mapExpM Mapper flore tlore m
tv (BasicOp (Scratch PrimType
t [SubExp]
shape)) =
BasicOp -> Exp tlore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tlore) -> m BasicOp -> m (Exp tlore)
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)
mapM (Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv) [SubExp]
shape)
mapExpM Mapper flore tlore m
tv (BasicOp (Reshape ShapeChange SubExp
shape VName
arrexp)) =
BasicOp -> Exp tlore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tlore) -> m BasicOp -> m (Exp tlore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ShapeChange SubExp -> VName -> BasicOp
Reshape (ShapeChange SubExp -> VName -> BasicOp)
-> m (ShapeChange SubExp) -> m (VName -> BasicOp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
(DimChange SubExp -> m (DimChange SubExp))
-> ShapeChange SubExp -> m (ShapeChange SubExp)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((SubExp -> m SubExp) -> DimChange SubExp -> m (DimChange SubExp)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
Data.Traversable.traverse (Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv)) ShapeChange SubExp
shape m (VName -> BasicOp) -> m VName -> m BasicOp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
Mapper flore tlore m -> VName -> m VName
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> VName -> m VName
mapOnVName Mapper flore tlore m
tv VName
arrexp)
mapExpM Mapper flore tlore m
tv (BasicOp (Rearrange [Int]
perm VName
e)) =
BasicOp -> Exp tlore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tlore) -> m BasicOp -> m (Exp tlore)
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 flore tlore m -> VName -> m VName
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> VName -> m VName
mapOnVName Mapper flore tlore m
tv VName
e)
mapExpM Mapper flore tlore m
tv (BasicOp (Rotate [SubExp]
es VName
e)) =
BasicOp -> Exp tlore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tlore) -> m BasicOp -> m (Exp tlore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([SubExp] -> VName -> BasicOp
Rotate ([SubExp] -> VName -> BasicOp)
-> m [SubExp] -> m (VName -> 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)
mapM (Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv) [SubExp]
es m (VName -> BasicOp) -> m VName -> m BasicOp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper flore tlore m -> VName -> m VName
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> VName -> m VName
mapOnVName Mapper flore tlore m
tv VName
e)
mapExpM Mapper flore tlore m
tv (BasicOp (Concat Int
i VName
x [VName]
ys SubExp
size)) =
BasicOp -> Exp tlore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tlore) -> m BasicOp -> m (Exp tlore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> VName -> [VName] -> SubExp -> BasicOp
Concat Int
i (VName -> [VName] -> SubExp -> BasicOp)
-> m VName -> m ([VName] -> SubExp -> BasicOp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
Mapper flore tlore m -> VName -> m VName
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> VName -> m VName
mapOnVName Mapper flore tlore m
tv VName
x m ([VName] -> SubExp -> BasicOp)
-> m [VName] -> m (SubExp -> BasicOp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (VName -> m VName) -> [VName] -> m [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Mapper flore tlore m -> VName -> m VName
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> VName -> m VName
mapOnVName Mapper flore tlore m
tv) [VName]
ys m (SubExp -> BasicOp) -> m SubExp -> m BasicOp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv SubExp
size)
mapExpM Mapper flore tlore m
tv (BasicOp (Copy VName
e)) =
BasicOp -> Exp tlore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tlore) -> m BasicOp -> m (Exp tlore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> BasicOp
Copy (VName -> BasicOp) -> m VName -> m BasicOp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper flore tlore m -> VName -> m VName
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> VName -> m VName
mapOnVName Mapper flore tlore m
tv VName
e)
mapExpM Mapper flore tlore m
tv (BasicOp (Manifest [Int]
perm VName
e)) =
BasicOp -> Exp tlore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tlore) -> m BasicOp -> m (Exp tlore)
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 flore tlore m -> VName -> m VName
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> VName -> m VName
mapOnVName Mapper flore tlore m
tv VName
e)
mapExpM Mapper flore tlore m
tv (BasicOp (Assert SubExp
e ErrorMsg SubExp
msg (SrcLoc, [SrcLoc])
loc)) =
BasicOp -> Exp tlore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tlore) -> m BasicOp -> m (Exp tlore)
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 flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv SubExp
e m (ErrorMsg SubExp -> (SrcLoc, [SrcLoc]) -> BasicOp)
-> m (ErrorMsg SubExp) -> m ((SrcLoc, [SrcLoc]) -> BasicOp)
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)
traverse (Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv) ErrorMsg SubExp
msg m ((SrcLoc, [SrcLoc]) -> BasicOp)
-> m (SrcLoc, [SrcLoc]) -> m BasicOp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (SrcLoc, [SrcLoc]) -> m (SrcLoc, [SrcLoc])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SrcLoc, [SrcLoc])
loc)
mapExpM Mapper flore tlore m
tv (BasicOp (Opaque SubExp
e)) =
BasicOp -> Exp tlore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tlore) -> m BasicOp -> m (Exp tlore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> BasicOp
Opaque (SubExp -> BasicOp) -> m SubExp -> m BasicOp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv SubExp
e)
mapExpM Mapper flore tlore m
tv (DoLoop [(FParam flore, SubExp)]
ctxmerge [(FParam flore, SubExp)]
valmerge LoopForm flore
form BodyT flore
loopbody) = do
[Param (FParamInfo tlore)]
ctxparams' <- (FParam flore -> m (Param (FParamInfo tlore)))
-> [FParam flore] -> m [Param (FParamInfo tlore)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Mapper flore tlore m
-> FParam flore -> m (Param (FParamInfo tlore))
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> FParam flore -> m (FParam tlore)
mapOnFParam Mapper flore tlore m
tv) [FParam flore]
ctxparams
[Param (FParamInfo tlore)]
valparams' <- (FParam flore -> m (Param (FParamInfo tlore)))
-> [FParam flore] -> m [Param (FParamInfo tlore)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Mapper flore tlore m
-> FParam flore -> m (Param (FParamInfo tlore))
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> FParam flore -> m (FParam tlore)
mapOnFParam Mapper flore tlore m
tv) [FParam flore]
valparams
LoopForm tlore
form' <- Mapper flore tlore m -> LoopForm flore -> m (LoopForm tlore)
forall (m :: * -> *) flore tlore.
Monad m =>
Mapper flore tlore m -> LoopForm flore -> m (LoopForm tlore)
mapOnLoopForm Mapper flore tlore m
tv LoopForm flore
form
let scope :: Scope tlore
scope = LoopForm tlore -> Scope tlore
forall lore a. Scoped lore a => a -> Scope lore
scopeOf LoopForm tlore
form' Scope tlore -> Scope tlore -> Scope tlore
forall a. Semigroup a => a -> a -> a
<> [Param (FParamInfo tlore)] -> Scope tlore
forall lore dec.
(FParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfFParams ([Param (FParamInfo tlore)]
ctxparams'[Param (FParamInfo tlore)]
-> [Param (FParamInfo tlore)] -> [Param (FParamInfo tlore)]
forall a. [a] -> [a] -> [a]
++[Param (FParamInfo tlore)]
valparams')
[(Param (FParamInfo tlore), SubExp)]
-> [(Param (FParamInfo tlore), SubExp)]
-> LoopForm tlore
-> BodyT tlore
-> Exp tlore
forall lore.
[(FParam lore, SubExp)]
-> [(FParam lore, SubExp)]
-> LoopForm lore
-> BodyT lore
-> ExpT lore
DoLoop ([(Param (FParamInfo tlore), SubExp)]
-> [(Param (FParamInfo tlore), SubExp)]
-> LoopForm tlore
-> BodyT tlore
-> Exp tlore)
-> m [(Param (FParamInfo tlore), SubExp)]
-> m ([(Param (FParamInfo tlore), SubExp)]
-> LoopForm tlore -> BodyT tlore -> Exp tlore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
([Param (FParamInfo tlore)]
-> [SubExp] -> [(Param (FParamInfo tlore), SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param (FParamInfo tlore)]
ctxparams' ([SubExp] -> [(Param (FParamInfo tlore), SubExp)])
-> m [SubExp] -> m [(Param (FParamInfo tlore), 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)
mapM (Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv) [SubExp]
ctxinits) m ([(Param (FParamInfo tlore), SubExp)]
-> LoopForm tlore -> BodyT tlore -> Exp tlore)
-> m [(Param (FParamInfo tlore), SubExp)]
-> m (LoopForm tlore -> BodyT tlore -> Exp tlore)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
([Param (FParamInfo tlore)]
-> [SubExp] -> [(Param (FParamInfo tlore), SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param (FParamInfo tlore)]
valparams' ([SubExp] -> [(Param (FParamInfo tlore), SubExp)])
-> m [SubExp] -> m [(Param (FParamInfo tlore), 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)
mapM (Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv) [SubExp]
valinits) m (LoopForm tlore -> BodyT tlore -> Exp tlore)
-> m (LoopForm tlore) -> m (BodyT tlore -> Exp tlore)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
LoopForm tlore -> m (LoopForm tlore)
forall (f :: * -> *) a. Applicative f => a -> f a
pure LoopForm tlore
form' m (BodyT tlore -> Exp tlore) -> m (BodyT tlore) -> m (Exp tlore)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper flore tlore m
-> Scope tlore -> BodyT flore -> m (BodyT tlore)
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> Scope tlore -> Body flore -> m (Body tlore)
mapOnBody Mapper flore tlore m
tv Scope tlore
scope BodyT flore
loopbody
where ([FParam flore]
ctxparams,[SubExp]
ctxinits) = [(FParam flore, SubExp)] -> ([FParam flore], [SubExp])
forall a b. [(a, b)] -> ([a], [b])
unzip [(FParam flore, SubExp)]
ctxmerge
([FParam flore]
valparams,[SubExp]
valinits) = [(FParam flore, SubExp)] -> ([FParam flore], [SubExp])
forall a b. [(a, b)] -> ([a], [b])
unzip [(FParam flore, SubExp)]
valmerge
mapExpM Mapper flore tlore m
tv (Op Op flore
op) =
Op tlore -> Exp tlore
forall lore. Op lore -> ExpT lore
Op (Op tlore -> Exp tlore) -> m (Op tlore) -> m (Exp tlore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper flore tlore m -> Op flore -> m (Op tlore)
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> Op flore -> m (Op tlore)
mapOnOp Mapper flore tlore m
tv Op flore
op
mapOnShape :: Monad m => Mapper flore tlore m -> Shape -> m Shape
mapOnShape :: Mapper flore tlore m -> Shape -> m Shape
mapOnShape Mapper flore tlore 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)
mapM (Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv) [SubExp]
ds
mapOnLoopForm :: Monad m =>
Mapper flore tlore m -> LoopForm flore -> m (LoopForm tlore)
mapOnLoopForm :: Mapper flore tlore m -> LoopForm flore -> m (LoopForm tlore)
mapOnLoopForm Mapper flore tlore m
tv (ForLoop VName
i IntType
it SubExp
bound [(LParam flore, VName)]
loop_vars) =
VName
-> IntType -> SubExp -> [(LParam tlore, VName)] -> LoopForm tlore
forall lore.
VName
-> IntType -> SubExp -> [(LParam lore, VName)] -> LoopForm lore
ForLoop (VName
-> IntType -> SubExp -> [(LParam tlore, VName)] -> LoopForm tlore)
-> m VName
-> m (IntType
-> SubExp -> [(LParam tlore, VName)] -> LoopForm tlore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper flore tlore m -> VName -> m VName
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> VName -> m VName
mapOnVName Mapper flore tlore m
tv VName
i m (IntType -> SubExp -> [(LParam tlore, VName)] -> LoopForm tlore)
-> m IntType
-> m (SubExp -> [(LParam tlore, VName)] -> LoopForm tlore)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> IntType -> m IntType
forall (f :: * -> *) a. Applicative f => a -> f a
pure IntType
it m (SubExp -> [(LParam tlore, VName)] -> LoopForm tlore)
-> m SubExp -> m ([(LParam tlore, VName)] -> LoopForm tlore)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mapper flore tlore m -> SubExp -> m SubExp
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> SubExp -> m SubExp
mapOnSubExp Mapper flore tlore m
tv SubExp
bound m ([(LParam tlore, VName)] -> LoopForm tlore)
-> m [(LParam tlore, VName)] -> m (LoopForm tlore)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
([LParam tlore] -> [VName] -> [(LParam tlore, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([LParam tlore] -> [VName] -> [(LParam tlore, VName)])
-> m [LParam tlore] -> m ([VName] -> [(LParam tlore, VName)])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (LParam flore -> m (LParam tlore))
-> [LParam flore] -> m [LParam tlore]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Mapper flore tlore m -> LParam flore -> m (LParam tlore)
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> LParam flore -> m (LParam tlore)
mapOnLParam Mapper flore tlore m
tv) [LParam flore]
loop_lparams m ([VName] -> [(LParam tlore, VName)])
-> m [VName] -> m [(LParam tlore, VName)]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (VName -> m VName) -> [VName] -> m [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Mapper flore tlore m -> VName -> m VName
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> VName -> m VName
mapOnVName Mapper flore tlore m
tv) [VName]
loop_arrs)
where ([LParam flore]
loop_lparams,[VName]
loop_arrs) = [(LParam flore, VName)] -> ([LParam flore], [VName])
forall a b. [(a, b)] -> ([a], [b])
unzip [(LParam flore, VName)]
loop_vars
mapOnLoopForm Mapper flore tlore m
tv (WhileLoop VName
cond) =
VName -> LoopForm tlore
forall lore. VName -> LoopForm lore
WhileLoop (VName -> LoopForm tlore) -> m VName -> m (LoopForm tlore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper flore tlore m -> VName -> m VName
forall flore tlore (m :: * -> *).
Mapper flore tlore m -> VName -> m VName
mapOnVName Mapper flore tlore m
tv VName
cond
mapExp :: Mapper flore tlore Identity -> Exp flore -> Exp tlore
mapExp :: Mapper flore tlore Identity -> Exp flore -> Exp tlore
mapExp Mapper flore tlore Identity
m = Identity (Exp tlore) -> Exp tlore
forall a. Identity a -> a
runIdentity (Identity (Exp tlore) -> Exp tlore)
-> (Exp flore -> Identity (Exp tlore)) -> Exp flore -> Exp tlore
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mapper flore tlore Identity -> Exp flore -> Identity (Exp tlore)
forall (m :: * -> *) flore tlore.
(Applicative m, Monad m) =>
Mapper flore tlore m -> Exp flore -> m (Exp tlore)
mapExpM Mapper flore tlore Identity
m
mapOnType :: Monad m =>
(SubExp -> m SubExp) -> Type -> m Type
mapOnType :: (SubExp -> m SubExp) -> Type -> m Type
mapOnType SubExp -> m SubExp
_ (Prim PrimType
bt) = Type -> m Type
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> m Type) -> Type -> m Type
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
bt
mapOnType SubExp -> m SubExp
_ (Mem Space
space) = Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> m Type) -> Type -> m Type
forall a b. (a -> b) -> a -> b
$ Space -> Type
forall shape u. Space -> TypeBase shape u
Mem Space
space
mapOnType SubExp -> m SubExp
f (Array PrimType
bt Shape
shape NoUniqueness
u) =
PrimType -> Shape -> NoUniqueness -> Type
forall shape u. PrimType -> shape -> u -> TypeBase shape u
Array PrimType
bt (Shape -> NoUniqueness -> Type)
-> m Shape -> m (NoUniqueness -> Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([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)
mapM SubExp -> m SubExp
f (Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shape)) m (NoUniqueness -> Type) -> m NoUniqueness -> m Type
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> NoUniqueness -> m NoUniqueness
forall (f :: * -> *) a. Applicative f => a -> f a
pure NoUniqueness
u
data Walker lore m = Walker {
Walker lore m -> SubExp -> m ()
walkOnSubExp :: SubExp -> m ()
, Walker lore m -> Body lore -> m ()
walkOnBody :: Body lore -> m ()
, Walker lore m -> VName -> m ()
walkOnVName :: VName -> m ()
, Walker lore m -> RetType lore -> m ()
walkOnRetType :: RetType lore -> m ()
, Walker lore m -> BranchType lore -> m ()
walkOnBranchType :: BranchType lore -> m ()
, Walker lore m -> FParam lore -> m ()
walkOnFParam :: FParam lore -> m ()
, Walker lore m -> LParam lore -> m ()
walkOnLParam :: LParam lore -> m ()
, Walker lore m -> Op lore -> m ()
walkOnOp :: Op lore -> m ()
}
identityWalker :: Monad m => Walker lore m
identityWalker :: Walker lore m
identityWalker = Walker :: forall lore (m :: * -> *).
(SubExp -> m ())
-> (Body lore -> m ())
-> (VName -> m ())
-> (RetType lore -> m ())
-> (BranchType lore -> m ())
-> (FParam lore -> m ())
-> (LParam lore -> m ())
-> (Op lore -> m ())
-> Walker lore m
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 (m :: * -> *) a. Monad m => a -> m a
return ()
, walkOnBody :: Body lore -> m ()
walkOnBody = m () -> Body lore -> m ()
forall a b. a -> b -> a
const (m () -> Body lore -> m ()) -> m () -> Body lore -> m ()
forall a b. (a -> b) -> a -> b
$ () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
, 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 (m :: * -> *) a. Monad m => a -> m a
return ()
, walkOnRetType :: RetType lore -> m ()
walkOnRetType = m () -> RetType lore -> m ()
forall a b. a -> b -> a
const (m () -> RetType lore -> m ()) -> m () -> RetType lore -> m ()
forall a b. (a -> b) -> a -> b
$ () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
, walkOnBranchType :: BranchType lore -> m ()
walkOnBranchType = m () -> BranchType lore -> m ()
forall a b. a -> b -> a
const (m () -> BranchType lore -> m ())
-> m () -> BranchType lore -> m ()
forall a b. (a -> b) -> a -> b
$ () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
, walkOnFParam :: FParam lore -> m ()
walkOnFParam = m () -> FParam lore -> m ()
forall a b. a -> b -> a
const (m () -> FParam lore -> m ()) -> m () -> FParam lore -> m ()
forall a b. (a -> b) -> a -> b
$ () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
, walkOnLParam :: LParam lore -> m ()
walkOnLParam = m () -> LParam lore -> m ()
forall a b. a -> b -> a
const (m () -> LParam lore -> m ()) -> m () -> LParam lore -> m ()
forall a b. (a -> b) -> a -> b
$ () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
, walkOnOp :: Op lore -> m ()
walkOnOp = m () -> Op lore -> m ()
forall a b. a -> b -> a
const (m () -> Op lore -> m ()) -> m () -> Op lore -> m ()
forall a b. (a -> b) -> a -> b
$ () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
}
walkOnShape :: Monad m => Walker lore m -> Shape -> m ()
walkOnShape :: Walker lore m -> Shape -> m ()
walkOnShape Walker lore 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 lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv) [SubExp]
ds
walkOnType :: Monad m =>
(SubExp -> m ()) -> Type -> m ()
walkOnType :: (SubExp -> m ()) -> Type -> m ()
walkOnType SubExp -> m ()
_ Prim{} = () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
walkOnType SubExp -> m ()
_ Mem{} = () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
walkOnType SubExp -> m ()
f (Array PrimType
_ Shape
shape NoUniqueness
_) = (SubExp -> m ()) -> [SubExp] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ SubExp -> m ()
f ([SubExp] -> m ()) -> [SubExp] -> m ()
forall a b. (a -> b) -> a -> b
$ Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shape
walkOnLoopForm :: Monad m => Walker lore m -> LoopForm lore -> m ()
walkOnLoopForm :: Walker lore m -> LoopForm lore -> m ()
walkOnLoopForm Walker lore m
tv (ForLoop VName
i IntType
_ SubExp
bound [(LParam lore, VName)]
loop_vars) =
Walker lore m -> VName -> m ()
forall lore (m :: * -> *). Walker lore m -> VName -> m ()
walkOnVName Walker lore m
tv VName
i m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv SubExp
bound m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
(LParam lore -> m ()) -> [LParam lore] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker lore m -> LParam lore -> m ()
forall lore (m :: * -> *). Walker lore m -> LParam lore -> m ()
walkOnLParam Walker lore m
tv) [LParam lore]
loop_lparams m () -> m () -> m ()
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 lore m -> VName -> m ()
forall lore (m :: * -> *). Walker lore m -> VName -> m ()
walkOnVName Walker lore m
tv) [VName]
loop_arrs
where ([LParam lore]
loop_lparams,[VName]
loop_arrs) = [(LParam lore, VName)] -> ([LParam lore], [VName])
forall a b. [(a, b)] -> ([a], [b])
unzip [(LParam lore, VName)]
loop_vars
walkOnLoopForm Walker lore m
tv (WhileLoop VName
cond) =
Walker lore m -> VName -> m ()
forall lore (m :: * -> *). Walker lore m -> VName -> m ()
walkOnVName Walker lore m
tv VName
cond
walkExpM :: Monad m => Walker lore m -> Exp lore -> m ()
walkExpM :: Walker lore m -> Exp lore -> m ()
walkExpM Walker lore m
tv (BasicOp (SubExp SubExp
se)) =
Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv SubExp
se
walkExpM Walker lore 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 lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv) [SubExp]
els m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (SubExp -> m ()) -> Type -> m ()
forall (m :: * -> *). Monad m => (SubExp -> m ()) -> Type -> m ()
walkOnType (Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv) Type
rowt
walkExpM Walker lore m
tv (BasicOp (BinOp BinOp
_ SubExp
x SubExp
y)) =
Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv SubExp
x m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv SubExp
y
walkExpM Walker lore m
tv (BasicOp (CmpOp CmpOp
_ SubExp
x SubExp
y)) =
Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv SubExp
x m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv SubExp
y
walkExpM Walker lore m
tv (BasicOp (ConvOp ConvOp
_ SubExp
x)) =
Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv SubExp
x
walkExpM Walker lore m
tv (BasicOp (UnOp UnOp
_ SubExp
x)) =
Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv SubExp
x
walkExpM Walker lore m
tv (If SubExp
c BodyT lore
texp BodyT lore
fexp (IfDec [BranchType lore]
ts IfSort
_)) =
Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv SubExp
c m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker lore m -> BodyT lore -> m ()
forall lore (m :: * -> *). Walker lore m -> Body lore -> m ()
walkOnBody Walker lore m
tv BodyT lore
texp m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
Walker lore m -> BodyT lore -> m ()
forall lore (m :: * -> *). Walker lore m -> Body lore -> m ()
walkOnBody Walker lore m
tv BodyT lore
fexp m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (BranchType lore -> m ()) -> [BranchType lore] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker lore m -> BranchType lore -> m ()
forall lore (m :: * -> *). Walker lore m -> BranchType lore -> m ()
walkOnBranchType Walker lore m
tv) [BranchType lore]
ts
walkExpM Walker lore m
tv (Apply Name
_ [(SubExp, Diet)]
args [RetType lore]
ret (Safety, SrcLoc, [SrcLoc])
_) =
((SubExp, Diet) -> m ()) -> [(SubExp, Diet)] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore 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 m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (RetType lore -> m ()) -> [RetType lore] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker lore m -> RetType lore -> m ()
forall lore (m :: * -> *). Walker lore m -> RetType lore -> m ()
walkOnRetType Walker lore m
tv) [RetType lore]
ret
walkExpM Walker lore m
tv (BasicOp (Index VName
arr Slice SubExp
slice)) =
Walker lore m -> VName -> m ()
forall lore (m :: * -> *). Walker lore m -> VName -> m ()
walkOnVName Walker lore m
tv VName
arr m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (DimIndex SubExp -> m ()) -> Slice SubExp -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((SubExp -> m ()) -> DimIndex SubExp -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv)) Slice SubExp
slice
walkExpM Walker lore m
tv (BasicOp (Update VName
arr Slice SubExp
slice SubExp
se)) =
Walker lore m -> VName -> m ()
forall lore (m :: * -> *). Walker lore m -> VName -> m ()
walkOnVName Walker lore m
tv VName
arr m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
(DimIndex SubExp -> m ()) -> Slice SubExp -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((SubExp -> m ()) -> DimIndex SubExp -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv)) Slice SubExp
slice m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv SubExp
se
walkExpM Walker lore m
tv (BasicOp (Iota SubExp
n SubExp
x SubExp
s IntType
_)) =
Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv SubExp
n m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv SubExp
x m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv SubExp
s
walkExpM Walker lore m
tv (BasicOp (Replicate Shape
shape SubExp
vexp)) =
Walker lore m -> Shape -> m ()
forall (m :: * -> *) lore.
Monad m =>
Walker lore m -> Shape -> m ()
walkOnShape Walker lore m
tv Shape
shape m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv SubExp
vexp
walkExpM Walker lore m
tv (BasicOp (Repeat [Shape]
shapes Shape
innershape VName
v)) =
(Shape -> m ()) -> [Shape] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker lore m -> Shape -> m ()
forall (m :: * -> *) lore.
Monad m =>
Walker lore m -> Shape -> m ()
walkOnShape Walker lore m
tv) [Shape]
shapes m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker lore m -> Shape -> m ()
forall (m :: * -> *) lore.
Monad m =>
Walker lore m -> Shape -> m ()
walkOnShape Walker lore m
tv Shape
innershape m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker lore m -> VName -> m ()
forall lore (m :: * -> *). Walker lore m -> VName -> m ()
walkOnVName Walker lore m
tv VName
v
walkExpM Walker lore 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 lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv) [SubExp]
shape
walkExpM Walker lore m
tv (BasicOp (Reshape ShapeChange SubExp
shape VName
arrexp)) =
(DimChange SubExp -> m ()) -> ShapeChange SubExp -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((SubExp -> m ()) -> DimChange SubExp -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv)) ShapeChange SubExp
shape m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker lore m -> VName -> m ()
forall lore (m :: * -> *). Walker lore m -> VName -> m ()
walkOnVName Walker lore m
tv VName
arrexp
walkExpM Walker lore m
tv (BasicOp (Rearrange [Int]
_ VName
e)) =
Walker lore m -> VName -> m ()
forall lore (m :: * -> *). Walker lore m -> VName -> m ()
walkOnVName Walker lore m
tv VName
e
walkExpM Walker lore m
tv (BasicOp (Rotate [SubExp]
es VName
e)) =
(SubExp -> m ()) -> [SubExp] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv) [SubExp]
es m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker lore m -> VName -> m ()
forall lore (m :: * -> *). Walker lore m -> VName -> m ()
walkOnVName Walker lore m
tv VName
e
walkExpM Walker lore m
tv (BasicOp (Concat Int
_ VName
x [VName]
ys SubExp
size)) =
Walker lore m -> VName -> m ()
forall lore (m :: * -> *). Walker lore m -> VName -> m ()
walkOnVName Walker lore m
tv VName
x m () -> m () -> m ()
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 lore m -> VName -> m ()
forall lore (m :: * -> *). Walker lore m -> VName -> m ()
walkOnVName Walker lore m
tv) [VName]
ys m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv SubExp
size
walkExpM Walker lore m
tv (BasicOp (Copy VName
e)) =
Walker lore m -> VName -> m ()
forall lore (m :: * -> *). Walker lore m -> VName -> m ()
walkOnVName Walker lore m
tv VName
e
walkExpM Walker lore m
tv (BasicOp (Manifest [Int]
_ VName
e)) =
Walker lore m -> VName -> m ()
forall lore (m :: * -> *). Walker lore m -> VName -> m ()
walkOnVName Walker lore m
tv VName
e
walkExpM Walker lore m
tv (BasicOp (Assert SubExp
e ErrorMsg SubExp
msg (SrcLoc, [SrcLoc])
_)) =
Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv SubExp
e m () -> m () -> m ()
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 lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv) ErrorMsg SubExp
msg
walkExpM Walker lore m
tv (BasicOp (Opaque SubExp
e)) =
Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv SubExp
e
walkExpM Walker lore m
tv (DoLoop [(FParam lore, SubExp)]
ctxmerge [(FParam lore, SubExp)]
valmerge LoopForm lore
form BodyT lore
loopbody) = do
(FParam lore -> m ()) -> [FParam lore] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker lore m -> FParam lore -> m ()
forall lore (m :: * -> *). Walker lore m -> FParam lore -> m ()
walkOnFParam Walker lore m
tv) [FParam lore]
ctxparams
(FParam lore -> m ()) -> [FParam lore] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker lore m -> FParam lore -> m ()
forall lore (m :: * -> *). Walker lore m -> FParam lore -> m ()
walkOnFParam Walker lore m
tv) [FParam lore]
valparams
Walker lore m -> LoopForm lore -> m ()
forall (m :: * -> *) lore.
Monad m =>
Walker lore m -> LoopForm lore -> m ()
walkOnLoopForm Walker lore m
tv LoopForm lore
form
(SubExp -> m ()) -> [SubExp] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv) [SubExp]
ctxinits
(SubExp -> m ()) -> [SubExp] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Walker lore m -> SubExp -> m ()
forall lore (m :: * -> *). Walker lore m -> SubExp -> m ()
walkOnSubExp Walker lore m
tv) [SubExp]
valinits
Walker lore m -> BodyT lore -> m ()
forall lore (m :: * -> *). Walker lore m -> Body lore -> m ()
walkOnBody Walker lore m
tv BodyT lore
loopbody
where ([FParam lore]
ctxparams,[SubExp]
ctxinits) = [(FParam lore, SubExp)] -> ([FParam lore], [SubExp])
forall a b. [(a, b)] -> ([a], [b])
unzip [(FParam lore, SubExp)]
ctxmerge
([FParam lore]
valparams,[SubExp]
valinits) = [(FParam lore, SubExp)] -> ([FParam lore], [SubExp])
forall a b. [(a, b)] -> ([a], [b])
unzip [(FParam lore, SubExp)]
valmerge
walkExpM Walker lore m
tv (Op Op lore
op) =
Walker lore m -> Op lore -> m ()
forall lore (m :: * -> *). Walker lore m -> Op lore -> m ()
walkOnOp Walker lore m
tv Op lore
op