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
{

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

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

    template <class C, class... T>
    static constexpr decltype(auto)
    accessByTreePath(C&& container, Dune::TypeTree::HybridTreePath<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_;
  };


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

      ValueAndContainer()
        : 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==(ValueAndContainer const& other) const
      {
        return value_ == other.value_ && container_ == other.container_;
      }

    private:
      Value value_;
      Container container_;
    };

    template <class Value>
    class ValueAndContainer<Value,void>
    {
    public:
      template <class V, Dune::disableCopyMove<ValueAndContainer,V> = 0>
      ValueAndContainer(V&& value)
        : value_(FWD(value))
      {}

      ValueAndContainer()
        : value_()
      {}

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

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

    private:
      Value value_;
    };

    struct Ignore{};

    template <class Container>
    class ValueAndContainer<Ignore,Container>
    {
    public:
      template <class C>
      ValueAndContainer(Ignore, C&& container)
        : container_(FWD(container))
      {}

      ValueAndContainer()
        : 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]; }

      Ignore value() { return {}; }
      Ignore value() const { return {}; }

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

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

    private:
      Container container_;
    };


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

    template <class Value>
    ValueAndContainer(Value const&)
      -> ValueAndContainer<Value,void>;



    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 ValueAndContainer{value(node)};
        else
        if constexpr (Node::isPower) {
          using TransformedChild = decltype((*this)(node.child(0)));
          return ValueAndContainer{
            value(node),
            std::array<TransformedChild, Node::degree()>()};
        }
        else
        if constexpr (Node::isComposite) {
          return ValueAndContainer{
            value(node),
            Tools::apply_indices<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>
      decltype(auto) value(Node const& node) const
      {
        if constexpr(!Node::isLeaf && leafOnly)
          return Ignore{};
        else
          return nodeToValue_(node);
      }

    private:
      NodeToValue nodeToValue_;
    };

  } // end namespace Impl



  template <bool leafOnly = false, class Tree, class NodeToValue>
  auto makeTreeContainer(Tree const& tree, NodeToValue&& nodeToValue)
  {
    auto f = std::ref(nodeToValue);
    auto factory = Impl::ContainerFactory<decltype(f),leafOnly>(f);
    return TreeContainerStorage{factory(tree)};
  }



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

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

  template <class Value, class Tree, bool leafOnly = false>
  using UniformTreeContainer
    = TYPEOF(makeTreeContainer<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(makeTreeContainer<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