content: -- |
  --
  -- 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
    { mapOnSubExp :: SubExp -> m SubExp,
      -- | Most bodies are enclosed in a scope, which is passed along
      -- for convenience.
      mapOnBody :: Scope trep -> Body frep -> m (Body trep),
      mapOnVName :: VName -> m VName,
      mapOnRetType :: RetType frep -> m (RetType trep),
      mapOnBranchType :: BranchType frep -> m (BranchType trep),
      mapOnFParam :: FParam frep -> m (FParam trep),
      mapOnLParam :: LParam frep -> m (LParam 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 =
    Mapper
      { mapOnSubExp = pure,
        mapOnBody = const pure,
        mapOnVName = pure,
        mapOnRetType = pure,
        mapOnBranchType = pure,
        mapOnFParam = pure,
        mapOnLParam = pure,
        mapOnOp = 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 tv (BasicOp (SubExp se)) =
    BasicOp <$> (SubExp <$> mapOnSubExp tv se)
  mapExpM _ (BasicOp (ArrayVal vs t)) =
    pure $ BasicOp $ ArrayVal vs t
  mapExpM tv (BasicOp (ArrayLit els rowt)) =
    BasicOp
      <$> ( ArrayLit
              <$> mapM (mapOnSubExp tv) els
              <*> mapOnType (mapOnSubExp tv) rowt
          )
  mapExpM tv (BasicOp (BinOp bop x y)) =
    BasicOp <$> (BinOp bop <$> mapOnSubExp tv x <*> mapOnSubExp tv y)
  mapExpM tv (BasicOp (CmpOp op x y)) =
    BasicOp <$> (CmpOp op <$> mapOnSubExp tv x <*> mapOnSubExp tv y)
  mapExpM tv (BasicOp (ConvOp conv x)) =
    BasicOp <$> (ConvOp conv <$> mapOnSubExp tv x)
  mapExpM tv (BasicOp (UnOp op x)) =
    BasicOp <$> (UnOp op <$> mapOnSubExp tv x)
  mapExpM tv (Match ses cases defbody (MatchDec ts s)) =
    Match
      <$> mapM (mapOnSubExp tv) ses
      <*> mapM mapOnCase cases
      <*> mapOnBody tv mempty defbody
      <*> (MatchDec <$> mapM (mapOnBranchType tv) ts <*> pure s)
    where
      mapOnCase (Case vs body) = Case vs <$> mapOnBody tv mempty body
  mapExpM tv (Apply fname args ret safety) = do
    args' <- forM args $ \(arg, d) ->
      (,) <$> mapOnSubExp tv arg <*> pure d
    Apply fname args' <$> mapM (bitraverse (mapOnRetType tv) pure) ret <*> pure safety
  mapExpM tv (BasicOp (Index arr slice)) =
    BasicOp
      <$> ( Index
              <$> mapOnVName tv arr
              <*> traverse (mapOnSubExp tv) slice
          )
  mapExpM tv (BasicOp (Update safety arr slice se)) =
    BasicOp
      <$> ( Update safety
              <$> mapOnVName tv arr
              <*> traverse (mapOnSubExp tv) slice
              <*> mapOnSubExp tv se
          )
  mapExpM tv (BasicOp (FlatIndex arr slice)) =
    BasicOp
      <$> ( FlatIndex
              <$> mapOnVName tv arr
              <*> traverse (mapOnSubExp tv) slice
          )
  mapExpM tv (BasicOp (FlatUpdate arr slice se)) =
    BasicOp
      <$> ( FlatUpdate
              <$> mapOnVName tv arr
              <*> traverse (mapOnSubExp tv) slice
              <*> mapOnVName tv se
          )
  mapExpM tv (BasicOp (Iota n x s et)) =
    BasicOp <$> (Iota <$> mapOnSubExp tv n <*> mapOnSubExp tv x <*> mapOnSubExp tv s <*> pure et)
  mapExpM tv (BasicOp (Replicate shape vexp)) =
    BasicOp <$> (Replicate <$> mapOnShape tv shape <*> mapOnSubExp tv vexp)
  mapExpM tv (BasicOp (Scratch t shape)) =
    BasicOp <$> (Scratch t <$> mapM (mapOnSubExp tv) shape)
  mapExpM tv (BasicOp (Reshape arrexp newshape)) =
    BasicOp
      <$> ( Reshape
              <$> mapOnVName tv arrexp
              <*> mapM (mapOnSubExp tv) newshape
          )
  mapExpM tv (BasicOp (Rearrange e perm)) =
    BasicOp <$> (Rearrange <$> mapOnVName tv e <*> pure perm)
  mapExpM tv (BasicOp (Concat i (x :| ys) size)) = do
    x' <- mapOnVName tv x
    ys' <- mapM (mapOnVName tv) ys
    size' <- mapOnSubExp tv size
    pure $ BasicOp $ Concat i (x' :| ys') size'
  mapExpM tv (BasicOp (Manifest v perm)) =
    BasicOp <$> (Manifest <$> mapOnVName tv v <*> pure perm)
  mapExpM tv (BasicOp (Assert e msg)) =
    BasicOp <$> (Assert <$> mapOnSubExp tv e <*> traverse (mapOnSubExp tv) msg)
  mapExpM tv (BasicOp (Opaque op e)) =
    BasicOp <$> (Opaque op <$> mapOnSubExp tv e)
  mapExpM tv (BasicOp (UpdateAcc safety v is ses)) =
    BasicOp
      <$> ( UpdateAcc safety
              <$> mapOnVName tv v
              <*> mapM (mapOnSubExp tv) is
              <*> mapM (mapOnSubExp tv) ses
          )
  mapExpM tv (WithAcc inputs lam) =
    WithAcc <$> mapM onInput inputs <*> mapOnLambda tv lam
    where
      onInput (shape, vs, op) =
        (,,)
          <$> mapOnShape tv shape
          <*> mapM (mapOnVName tv) vs
          <*> traverse (bitraverse (mapOnLambda tv) (mapM (mapOnSubExp tv))) op
  mapExpM tv (Loop merge form loopbody) = do
    params' <- mapM (mapOnFParam tv) params
    form' <- mapOnLoopForm tv form
    let scope = scopeOfLoopForm form' <> scopeOfFParams params'
    Loop
      <$> (zip params' <$> mapM (mapOnSubExp tv) args)
      <*> pure form'
      <*> mapOnBody tv scope loopbody
    where
      (params, args) = unzip merge
  mapExpM tv (Op op) =
    Op <$> mapOnOp tv op
  
  mapOnShape :: (Monad m) => Mapper frep trep m -> Shape -> m Shape
  mapOnShape tv (Shape ds) = Shape <$> mapM (mapOnSubExp tv) ds
  
  mapOnLoopForm ::
    (Monad m) =>
    Mapper frep trep m ->
    LoopForm ->
    m LoopForm
  mapOnLoopForm tv (ForLoop i it bound) =
    ForLoop <$> mapOnVName tv i <*> pure it <*> mapOnSubExp tv bound
  mapOnLoopForm tv (WhileLoop cond) =
    WhileLoop <$> mapOnVName tv cond
  
  mapOnLambda ::
    (Monad m) =>
    Mapper frep trep m ->
    Lambda frep ->
    m (Lambda trep)
  mapOnLambda tv (Lambda params ret body) = do
    params' <- mapM (mapOnLParam tv) params
    Lambda params'
      <$> mapM (mapOnType (mapOnSubExp tv)) ret
      <*> mapOnBody tv (scopeOfLParams params') body
  
  -- | Like 'mapExpM', but in the 'Identity' monad.
  mapExp :: Mapper frep trep Identity -> Exp frep -> Exp trep
  mapExp m = runIdentity . mapExpM 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
    { walkOnSubExp :: SubExp -> m (),
      walkOnBody :: Scope rep -> Body rep -> m (),
      walkOnVName :: VName -> m (),
      walkOnRetType :: RetType rep -> m (),
      walkOnBranchType :: BranchType rep -> m (),
      walkOnFParam :: FParam rep -> m (),
      walkOnLParam :: LParam rep -> m (),
      walkOnOp :: Op rep -> m ()
    }
  
  -- | A no-op traversal.
  identityWalker :: forall rep m. (Monad m) => Walker rep m
  identityWalker =
    Walker
      { walkOnSubExp = const $ pure (),
        walkOnBody = const $ const $ pure (),
        walkOnVName = const $ pure (),
        walkOnRetType = const $ pure (),
        walkOnBranchType = const $ pure (),
        walkOnFParam = const $ pure (),
        walkOnLParam = const $ pure (),
        walkOnOp = const $ pure ()
      }
  
  walkOnShape :: (Monad m) => Walker rep m -> Shape -> m ()
  walkOnShape tv (Shape ds) = mapM_ (walkOnSubExp tv) ds
  
  walkOnType :: (Monad m) => Walker rep m -> Type -> m ()
  walkOnType _ Prim {} = pure ()
  walkOnType tv (Acc acc ispace ts _) = do
    walkOnVName tv acc
    traverse_ (walkOnSubExp tv) ispace
    mapM_ (walkOnType tv) ts
  walkOnType _ Mem {} = pure ()
  walkOnType tv (Array _ shape _) = walkOnShape tv shape
  
  walkOnLoopForm :: (Monad m) => Walker rep m -> LoopForm -> m ()
  walkOnLoopForm tv (ForLoop i _ bound) = do
    walkOnVName tv i
    walkOnSubExp tv bound
  walkOnLoopForm tv (WhileLoop cond) =
    walkOnVName tv cond
  
  walkOnLambda :: (Monad m) => Walker rep m -> Lambda rep -> m ()
  walkOnLambda tv (Lambda params ret body) = do
    mapM_ (walkOnLParam tv) params
    walkOnBody tv (scopeOfLParams params) body
    mapM_ (walkOnType tv) ret
  
  -- | As 'mapExpM', but do not construct a result AST.
  walkExpM :: (Monad m) => Walker rep m -> Exp rep -> m ()
  walkExpM tv (BasicOp (SubExp se)) =
    walkOnSubExp tv se
  walkExpM _ (BasicOp ArrayVal {}) =
    pure ()
  walkExpM tv (BasicOp (ArrayLit els rowt)) =
    mapM_ (walkOnSubExp tv) els >> walkOnType tv rowt
  walkExpM tv (BasicOp (BinOp _ x y)) =
    walkOnSubExp tv x >> walkOnSubExp tv y
  walkExpM tv (BasicOp (CmpOp _ x y)) =
    walkOnSubExp tv x >> walkOnSubExp tv y
  walkExpM tv (BasicOp (ConvOp _ x)) =
    walkOnSubExp tv x
  walkExpM tv (BasicOp (UnOp _ x)) =
    walkOnSubExp tv x
  walkExpM tv (Match ses cases defbody (MatchDec ts _)) = do
    mapM_ (walkOnSubExp tv) ses
    mapM_ (walkOnBody tv mempty . caseBody) cases
    walkOnBody tv mempty defbody
    mapM_ (walkOnBranchType tv) ts
  walkExpM tv (Apply _ args ret _) = do
    mapM_ (walkOnSubExp tv . fst) args
    mapM_ (walkOnRetType tv . fst) ret
  walkExpM tv (BasicOp (Index arr slice)) =
    walkOnVName tv arr >> traverse_ (walkOnSubExp tv) slice
  walkExpM tv (BasicOp (Update _ arr slice se)) =
    walkOnVName tv arr
      >> traverse_ (walkOnSubExp tv) slice
      >> walkOnSubExp tv se
  walkExpM tv (BasicOp (FlatIndex arr slice)) =
    walkOnVName tv arr >> traverse_ (walkOnSubExp tv) slice
  walkExpM tv (BasicOp (FlatUpdate arr slice se)) =
    walkOnVName tv arr
      >> traverse_ (walkOnSubExp tv) slice
      >> walkOnVName tv se
  walkExpM tv (BasicOp (Iota n x s _)) =
    walkOnSubExp tv n >> walkOnSubExp tv x >> walkOnSubExp tv s
  walkExpM tv (BasicOp (Replicate shape vexp)) =
    walkOnShape tv shape >> walkOnSubExp tv vexp
  walkExpM tv (BasicOp (Scratch _ shape)) =
    mapM_ (walkOnSubExp tv) shape
  walkExpM tv (BasicOp (Reshape arrexp shape)) =
    mapM_ (walkOnSubExp tv) shape >> walkOnVName tv arrexp
  walkExpM tv (BasicOp (Rearrange v _)) =
    walkOnVName tv v
  walkExpM tv (BasicOp (Concat _ (x :| ys) size)) =
    walkOnVName tv x >> mapM_ (walkOnVName tv) ys >> walkOnSubExp tv size
  walkExpM tv (BasicOp (Manifest v _)) =
    walkOnVName tv v
  walkExpM tv (BasicOp (Assert e msg)) =
    walkOnSubExp tv e >> traverse_ (walkOnSubExp tv) msg
  walkExpM tv (BasicOp (Opaque _ e)) =
    walkOnSubExp tv e
  walkExpM tv (BasicOp (UpdateAcc _ v is ses)) = do
    walkOnVName tv v
    mapM_ (walkOnSubExp tv) is
    mapM_ (walkOnSubExp tv) ses
  walkExpM tv (WithAcc inputs lam) = do
    forM_ inputs $ \(shape, vs, op) -> do
      walkOnShape tv shape
      mapM_ (walkOnVName tv) vs
      traverse_ (bitraverse (walkOnLambda tv) (mapM (walkOnSubExp tv))) op
    walkOnLambda tv lam
  walkExpM tv (Loop merge form loopbody) = do
    mapM_ (walkOnFParam tv) params
    walkOnLoopForm tv form
    mapM_ (walkOnSubExp tv) args
    let scope = scopeOfFParams params <> scopeOfLoopForm form
    walkOnBody tv scope loopbody
    where
      (params, args) = unzip merge
  walkExpM tv (Op op) =
    walkOnOp tv 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 f (Lambda ps ret (Body dec stms res)) =
    Lambda ps ret <$> (Body dec <$> f (scopeOfLParams ps) stms <*> pure res)
  