Traversal.hpp

#pragma once

#include <dune/common/rangeutilities.hh>

#include <dune/typetree/nodetags.hh>
#include <dune/typetree/treepath.hh>

#include <amdis/common/ForEach.hpp>
#include <amdis/common/TypeTraits.hpp>

// NOTE: backport of dune/typetree/traversal.hpp from Dune 2.7

namespace AMDiS
{
  namespace Impl
  {

    template <bool dynamic = true, class Tree>
    auto traversalDegree(Tree const& tree)
    {
      if constexpr (dynamic && Tree::isPower)
        return std::size_t(tree.degree());
      else if constexpr (Tree::isPower || Tree::isComposite)
        return std::integral_constant<std::size_t, Tree::degree()>{};
      else
        return tree.degree();
    }

    template <class Tree, class TreePath, class PreFunc, class LeafFunc, class PostFunc>
    void for_each_node(Tree&& tree, TreePath treePath, PreFunc&& preFunc, LeafFunc&& leafFunc, PostFunc&& postFunc)
    {
      using TreeType = std::decay_t<Tree>;
      if constexpr (TreeType::isLeaf) {
        // If we have a leaf tree just visit it using the leaf function.
        leafFunc(tree, treePath);
      } else {
        // Otherwise visit the tree with the pre function,
        // visit all children using a static or dynamic loop, and
        // finally visit the tree with the post function.
        preFunc(tree, treePath);
        Dune::Hybrid::forEach(Dune::range(traversalDegree(tree)), [&](auto i) {
          auto childTreePath = push_back(treePath, i);
          Impl::for_each_node(tree.child(i), childTreePath, preFunc, leafFunc, postFunc);
        });
        postFunc(tree, treePath);
      }
    }

  } // end namespace Impl



  template <class Tree, class PreFunc, class LeafFunc, class PostFunc>
  void for_each_node(Tree&& tree, PreFunc&& preFunc, LeafFunc&& leafFunc, PostFunc&& postFunc)
  {
    auto root = Dune::TypeTree::hybridTreePath();
    Impl::for_each_node(tree, root, preFunc, leafFunc, postFunc);
  }


  template <class Tree, class InnerFunc, class LeafFunc>
  void for_each_node(Tree&& tree, InnerFunc&& innerFunc, LeafFunc&& leafFunc)
  {
    auto root = Dune::TypeTree::hybridTreePath();
    Impl::for_each_node(tree, root, innerFunc, leafFunc, NoOp{});
  }


  template <class Tree, class NodeFunc>
  void for_each_node(Tree&& tree, NodeFunc&& nodeFunc)
  {
    auto root = Dune::TypeTree::hybridTreePath();
    Impl::for_each_node(tree, root, nodeFunc, nodeFunc, NoOp{});
  }


  template <class Tree, class LeafFunc>
  void for_each_leaf_node(Tree&& tree, LeafFunc&& leafFunc)
  {
    auto root = Dune::TypeTree::hybridTreePath();
    Impl::for_each_node(tree, root, NoOp{}, leafFunc, NoOp{});
  }

} // end namespace AMDiS