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.hh from Dune 2.7

namespace AMDiS {
namespace Traversal {
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();
  }

} // end namespace Impl_


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);
    const auto degree = Impl_::traversalDegree(tree);
    if constexpr (std::is_integral_v<TYPEOF(degree)>) {
      // Specialization for dynamic traversal
      for (std::size_t i = 0; i < std::size_t(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<VALUE(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