Traversal.hpp

#pragma once

#include <dune/common/version.hh>
#include <dune/common/rangeutilities.hh>
#include <dune/common/std/type_traits.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.hh from Dune 2.7

namespace AMDiS {
namespace Traversal {

  template <class Tree>
  using HasDynamicChildAccess = decltype(std::declval<Tree>().child(0u));


template <class Tree, class TreePath, class Pre, class Leaf, class Post>
void forEachNode(Tree&& tree, TreePath treePath, Pre&& preFunc, Leaf&& leafFunc, Post&& postFunc)
{
  using TreeType = std::decay_t<Tree>;
  if constexpr (TreeType::isLeaf) {
    // For leaf nodes just visit using the leaf function.
    leafFunc(tree, treePath);
  } else {
    preFunc(tree, treePath);
    if constexpr(Dune::Std::is_detected_v<HasDynamicChildAccess,TreeType>) {
      // Specialization for dynamic traversal
      for (std::size_t i = 0; i < tree.degree(); ++i) {
        auto childTreePath = Dune::TypeTree::push_back(treePath, i);
        forEachNode(tree.child(i), childTreePath, preFunc, leafFunc, postFunc);
      }
    } else {
      // Specialization for static traversal
      const auto indices = std::make_index_sequence<std::size_t(TreeType::degree())>{};
      Ranges::forIndices(indices, [&](auto i) {
        auto childTreePath = Dune::TypeTree::push_back(treePath, i);
        forEachNode(tree.child(i), childTreePath, preFunc, leafFunc, postFunc);
      });
    }
    postFunc(tree, treePath);
  }
}



template <class Tree, class Pre, class Leaf, class Post>
void forEachNode(Tree&& tree, Pre&& preFunc, Leaf&& leafFunc, Post&& postFunc)
{
  auto root = Dune::TypeTree::hybridTreePath();
  forEachNode(tree, root, preFunc, leafFunc, postFunc);
}


template <class Tree, class Inner, class Leaf>
void forEachNode(Tree&& tree, Inner&& innerFunc, Leaf&& leafFunc)
{
  auto root = Dune::TypeTree::hybridTreePath();
  forEachNode(tree, root, innerFunc, leafFunc, NoOp{});
}


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


template <class Tree, class Leaf>
void forEachLeafNode(Tree&& tree, Leaf&& leafFunc)
{
  auto root = Dune::TypeTree::hybridTreePath();
  forEachNode(tree, root, NoOp{}, leafFunc, NoOp{});
}

}} // end namespace AMDiS::Traversal