TreeContainer.hpp

#pragma once

#include <array>
#include <functional>
#include <type_traits>
#include <utility>

#include <dune/common/indices.hh>
#include <dune/common/tuplevector.hh>

#include <dune/typetree/treepath.hh>

#include <amdis/common/Apply.hpp>
#include <amdis/common/TypeTraits.hpp>
#include <amdis/typetree/TreePath.hpp>

// NOTE: backport of dune/typetree/treecontainer.hh

namespace AMDiS {
namespace TypeTree {

template <class Value>
class LeafNodeStorage
{
public:
  template <class V>
  LeafNodeStorage(V&& value)
    : value_(FWD(value))
  {}

  LeafNodeStorage()
    : value_()
  {}

  Value& value() { return value_; }
  Value const& value() const { return value_; }

  bool operator==(LeafNodeStorage const& other) const
  {
    return value_ == other.value_;
  }

private:
  Value value_;
};

template <class Value>
LeafNodeStorage(Value const&)
  -> LeafNodeStorage<Value>;


template <class Value, class Container>
class InnerNodeStorage
{
public:
  template <class V, class C>
  InnerNodeStorage(V&& value, C&& container)
    : value_(FWD(value))
    , container_(FWD(container))
  {}

  InnerNodeStorage()
    : value_()
    , container_()
  {}

  template <class I>
  auto& operator[](I const& i) { return container_[i]; }

  template <class I>
  auto const& operator[](I const& i) const { return container_[i]; }

  Value& value() { return value_; }
  Value const& value() const { return value_; }

  Container& container() { return container_; }
  Container const& container() const { return container_; }

  bool operator==(InnerNodeStorage const& other) const
  {
    return value_ == other.value_ && container_ == other.container_;
  }

private:
  Value value_;
  Container container_;
};

template <class Value, class Container>
InnerNodeStorage(Value const&, Container const&)
  -> InnerNodeStorage<Value, Container>;



template <class NodeToValue, bool leafOnly = false>
class ContainerFactory
{
public:

  ContainerFactory(NodeToValue nodeToValue)
    : nodeToValue_(std::move(nodeToValue))
  {}

  template <class Node>
  auto operator()(Node const& node) const
  {
    if constexpr (Node::isLeaf)
      return LeafNodeStorage{nodeToValue_(node)};
    else
    if constexpr (Node::isPower) {
      using TransformedChild = decltype((*this)(node.child(0)));
      return makeInnerNodeStorage(node,
        std::array<TransformedChild, Node::degree()>());
    }
    else
    if constexpr (Node::isComposite) {
      return makeInnerNodeStorage(node,
        Ranges::applyIndices<Node::degree()>(
        [&](auto... ii) { return Dune::makeTupleVector((*this)(node.child(ii))...); }));
    }
    else {
      static_assert(Node::isLeaf || Node::isPower || Node::isComposite,
        "Node must be one of leaf, power, composite.");
      return;
    }
  }

  template <class Node, class Container>
  auto makeInnerNodeStorage(Node const& node, Container&& container) const
  {
    if constexpr(!Node::isLeaf && leafOnly)
      return FWD(container);
    else
      return InnerNodeStorage{nodeToValue_(node), FWD(container)};
  }

private:
  NodeToValue nodeToValue_;
};


template <class Container>
class TreeContainerStorage
{
  using Self = TreeContainerStorage;

  template <class C>
  static constexpr auto
  accessByTreePath(C&& container, Dune::TypeTree::HybridTreePath<>) -> decltype(container.value())
  {
    return container.value();
  }

  template <class C, class T0, class... T>
  static constexpr decltype(auto)
  accessByTreePath(C&& container, Dune::TypeTree::HybridTreePath<T0,T...> const& path)
  {
    auto head = Dune::TypeTree::treePathEntry(path,Dune::Indices::_0);
    return accessByTreePath(container[head], pop_front(path));
  }

public:
  TreeContainerStorage()
    : container_()
  {}

  TreeContainerStorage(Container const& container)
    : container_(container)
  {}

  TreeContainerStorage(Container&& container)
    : container_(std::move(container))
  {}

  template <class... T>
  decltype(auto) operator[](Dune::TypeTree::HybridTreePath<T...> const& path) const
  {
    return accessByTreePath(container_, path);
  }

  template <class... T>
  decltype(auto) operator[](Dune::TypeTree::HybridTreePath<T...> const& path)
  {
    return accessByTreePath(container_, path);
  }

  Container const& data() const
  {
    return container_;
  }

  Container& data()
  {
    return container_;
  }

  bool operator==(TreeContainerStorage const& other) const
  {
    return container_ == other.container_;
  }

  bool operator!=(TreeContainerStorage const& other) const
  {
    return container_ != other.container_;
  }

private:
  Container container_;
};




template <bool leafOnly = false, class Tree, class NodeToValue>
auto treeContainer(Tree const& tree, NodeToValue nodeToValue)
{
  auto factory = ContainerFactory<NodeToValue,leafOnly>{nodeToValue};
  return TreeContainerStorage{factory(tree)};
}



template <class Value, bool leafOnly = false, class Tree>
auto treeContainer(Tree const& tree)
{
  return treeContainer<leafOnly>(tree, [](auto&&) { return Value{}; });
}

template <template<class> class NodeData, bool leafOnly = false, class Tree>
auto treeContainer(Tree const& tree)
{
  return treeContainer<leafOnly>(tree, [](auto&& node) { return NodeData<TYPEOF(node)>{}; });
}

template <class Value, class Tree, bool leafOnly = false>
using UniformTreeContainer
  = TYPEOF(treeContainer<Value,leafOnly>(std::declval<const Tree&>()));

template <class Value, class RowTree, class ColTree = RowTree, bool leafOnly = false>
using UniformTreeMatrix
  = UniformTreeContainer<UniformTreeContainer<Value,ColTree,leafOnly>,RowTree,leafOnly>;

template <template <class Node> class NodeData, class Tree, bool leafOnly = false>
using TreeContainer
  = TYPEOF(treeContainer<NodeData,leafOnly>(std::declval<const Tree&>()));


namespace Impl_ {

template <template <class,class> class NodeData, class Tree, bool leafOnly>
struct RowNodeData
{
  template <class RowNode>
  struct ColNodeData
  {
    template <class ColNode>
    using type = NodeData<RowNode, ColNode>;
  };

  template <class RowNode>
  using type = TreeContainer<ColNodeData<RowNode>::template type, Tree, leafOnly>;
};

} // end namespace Impl_

template <template <class,class> class NodeData, class RowTree, class ColTree = RowTree, bool leafOnly = false>
using TreeMatrix
  = TreeContainer<Impl_::RowNodeData<NodeData,ColTree,leafOnly>::template type,RowTree,leafOnly>;

}} //namespace AMDiS::TypeTree