UniqueBorderPartition.hpp

#pragma once

#include <cassert>
#include <set>

#include <dune/common/unused.hh>
#include <dune/common/version.hh>
#include <dune/grid/common/datahandleif.hh>

namespace Dune
{
  // forward declaration
  template <int dim> class UGGrid;
}

namespace AMDiS
{

  template <class Grid>
  class UniqueBorderPartition
      : public Dune::CommDataHandleIF<UniqueBorderPartition<Grid>, int>
  {
    using IdSet = typename Grid::GlobalIdSet;
    using IdType = typename IdSet::IdType;

  public:
    using EntitySet = std::set<IdType>;

  public:

    UniqueBorderPartition(int rank, Grid const& grid)
      : myrank_(rank)
      , idSet_(grid.globalIdSet())
    {}

    // Communicate all entities
    bool contains(int /*dim*/, int /*codim*/) const { return true; }

    // communicate exactly one integer, the rank
    bool fixedSize(int /*dim*/, int /*codim*/) const { return true; }

    // Always contains one int, the rank
    template <class Entity>
    std::size_t size(Entity const& e) const { return 1; }

    template <class MessageBuffer, class Entity>
    void gather(MessageBuffer& buff, Entity const& e) const
    {
      buff.write(myrank_);
    }

    template <class MessageBuffer, class Entity>
    void scatter(MessageBuffer& buff, Entity const& e, std::size_t n)
    {
      scatterImpl(buff, e, n, int_t<Entity::codimension>{});
    }

    bool contains(IdType const& id) const
    {
      return notOwner_.count(id) == 0;
    }

    int numIterations() const
    {
      return 1;
    }

  private:

    template <class MessageBuffer, class Entity, int cd>
    void scatterImpl(MessageBuffer& buff, Entity const& e, [[maybe_unused]] std::size_t n, int_t<cd>)
    {
      assert(n == 1);

      int rank = 0;
      buff.read(rank);

      // insert only border entities that are owned by other processors, i.e. rank > myrank
      // Those entities are not owned by this rank.
      if (rank > myrank_)
        notOwner_.insert(idSet_.id(e));
    }

    template <class MessageBuffer, class Entity>
    void scatterImpl(MessageBuffer& buff, Entity const& e, [[maybe_unused]] std::size_t n, int_t<0>)
    {
      assert(n == 1);

      int rank = 0;
      buff.read(rank);

      // insert only border entities that are owned by other processors, i.e. rank > myrank
      // Those entities are not owned by this rank.
      if (rank > myrank_) {
        for (int codim = 1; codim <= Grid::dimension; ++codim) {
          for (int i = 0; i < int(e.subEntities(codim)); ++i) {
            notOwner_.insert(idSet_.subId(e, i, codim));
          }
        }
      }
    }

  private:
    int myrank_;
    EntitySet notOwner_;
    IdSet const& idSet_;
  };


  template <int dim>
  class UniqueBorderPartition<Dune::UGGrid<dim>>
      : public Dune::CommDataHandleIF<UniqueBorderPartition<Dune::UGGrid<dim>>, int>
  {
    using Grid = Dune::UGGrid<dim>;
    using IdSet = typename Grid::LocalIdSet;
    using IdType = typename IdSet::IdType;

  public:
    using EntitySet = std::set<IdType>;

  public:

    UniqueBorderPartition(int rank, Grid const& grid)
      : myrank_(rank)
      , idSet_(grid.localIdSet())
    {}

    // Communicate all entities
    bool contains(int /*dim*/, int codim) const { return true; }

    // see size()
    bool fixedSize(int /*dim*/, int codim) const { return codim != 0; }

    // for codim=0 elements communicate data for all subEntities, otherwise communicate just the owner rank
    template <class Entity>
    std::size_t size(Entity const& e) const
    {
      if (Entity::codimension != 0)
        return 1;

      std::size_t s = 0;
      for (int codim = 1; codim <= Grid::dimension; ++codim)
        s += e.subEntities(codim);
      return s;
    }

    template <class MessageBuffer, class Entity>
    void gather(MessageBuffer& buff, Entity const& e) const
    {
      gatherImpl(buff, e, int_t<Entity::codimension>{});
    }

    template <class MessageBuffer, class Entity>
    void scatter(MessageBuffer& buff, Entity const& e, std::size_t n)
    {
      scatterImpl(buff, e, n, int_t<Entity::codimension>{});
    }

    bool contains(IdType const& id) const
    {
      return entityToRank_[id] == myrank_;
    }

    int numIterations() const
    {
      return Grid::dimension;
    }

  private:
    template <class MessageBuffer, class Entity, int cd>
    void gatherImpl(MessageBuffer& buff, Entity const& e, int_t<cd>) const
    {
      int& rank = entityToRank_[idSet_.id(e)];
      rank = std::max(rank, myrank_);
      buff.write(rank);
    }

    template <class MessageBuffer, class Entity>
    void gatherImpl(MessageBuffer& buff, Entity const& e, int_t<0>) const
    {
      // maybe use global unique numbering (?)
      for (int codim = 1; codim <= Grid::dimension; ++codim) {
        for (int i = 0; i < int(e.subEntities(codim)); ++i) {
          int& rank = entityToRank_[idSet_.subId(e, i, codim)];
          rank = std::max(rank, myrank_);
          buff.write(rank);
        }
      }
    }

    template <class MessageBuffer, class Entity, int cd>
    void scatterImpl(MessageBuffer& buff, Entity const& e, [[maybe_unused]] std::size_t n, int_t<cd>)
    {
      assert(n == 1);

      int rank = 0;
      buff.read(rank);

      int& storedRank = entityToRank_[idSet_.id(e)];
      storedRank = std::max(rank, storedRank);
    }

    template <class MessageBuffer, class Entity>
    void scatterImpl(MessageBuffer& buff, Entity const& e, std::size_t n, int_t<0>)
    {
      std::size_t j = 0;
      for (int codim = 1; codim <= Grid::dimension; ++codim) {
        for (int i = 0; i < int(e.subEntities(codim)); ++i, ++j) {
          assert(j < n);
          int rank = 0;
          buff.read(rank);

          int& storedRank = entityToRank_[idSet_.subId(e, i, codim)];
          storedRank = std::max(rank, storedRank);
        }
      }
    }

  private:
    int myrank_;
    mutable std::map<IdType, int> entityToRank_;
    IdSet const& idSet_;
  };

} // end namespace AMDiS