Fixed an issue with stable_distribute and stable_distribute_inplace

when the total size is less than the number of processes.
Also added a test that was failing before the fix.

include/mxx/distribution.hpp

@@ -111,15 +111,18 @@ void stable_distribute(_InIterator begin, _InIterator end, _OutIterator out, con
  size_t prefix = mxx::exscan(local_size, comm);
 
  // calculate where to send elements
+ // if there are any elements to send
  std::vector<size_t> send_counts(comm.size(), 0);
- blk_dist part(total_size, comm.size(), comm.rank());
- int first_p = part.rank_of(prefix);
- size_t left_to_send = local_size;
- for (; left_to_send > 0 && first_p < comm.size(); ++first_p) {
- size_t nsend = std::min<size_t>(part.iprefix_size(first_p) - prefix, left_to_send);
- send_counts[first_p] = nsend;
- left_to_send -= nsend;
- prefix += nsend;
+ if (local_size > 0) {
+ blk_dist part(total_size, comm.size(), comm.rank());
+ int first_p = part.rank_of(prefix);
+ size_t left_to_send = local_size;
+ for (; left_to_send > 0 && first_p < comm.size(); ++first_p) {
+ size_t nsend = std::min<size_t>(part.iprefix_size(first_p) - prefix, left_to_send);
+ send_counts[first_p] = nsend;
+ left_to_send -= nsend;
+ prefix += nsend;
+ }
  }
  std::vector<size_t> recv_counts = mxx::all2all(send_counts, comm);
  // TODO: accept iterators in mxx::all2all?

test/test_distribution.cpp

@@ -143,6 +143,16 @@ TEST(MxxDistribution, DistributeVector) {
 
  test_distribute<std::vector<int>>(size, gen, c);
  test_stable_distribute<std::vector<int>>(size, gen, c);
+
+ // create a distribution of total size smaller than
+ // the total number of processes and zero elements on the last process
+ size = (c.rank() % 2 == 0) ? 1 : 0;
+ if (c.is_last()) {
+ size = 0;
+ }
+ // XXX: test_distribute fails with an assertion error in mxx::sort
+ // test_distribute<std::vector<int>>(size, gen, c);
+ test_stable_distribute<std::vector<int>>(size, gen, c);
 }