Change the type stored in lazy mpi objects

TRIQS · Sep 29, 2024 · f3772cd · f3772cd
1 parent 671f0a8
commit f3772cd
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 13 deletions.
diff --git a/c++/nda/mpi/gather.hpp b/c++/nda/mpi/gather.hpp
@@ -55,11 +55,11 @@ struct mpi::lazy<mpi::tag::gather, A> {
   /// Value type of the array/view.
   using value_type = typename std::decay_t<A>::value_type;
 
-  /// Const view type of the array/view stored in the lazy object.
-  using const_view_type = decltype(std::declval<const A>()());
+  /// Type of the array/view stored in the lazy object.
+  using stored_type = A;
 
-  /// View of the array/view to be gathered.
-  const_view_type rhs;
+  /// Array/View to be gathered.
+  stored_type rhs;
 
   /// MPI communicator.
   mpi::communicator comm;

diff --git a/c++/nda/mpi/reduce.hpp b/c++/nda/mpi/reduce.hpp
@@ -57,11 +57,11 @@ struct mpi::lazy<mpi::tag::reduce, A> {
   /// Value type of the array/view.
   using value_type = typename std::decay_t<A>::value_type;
 
-  /// Const view type of the array/view stored in the lazy object.
-  using const_view_type = decltype(std::declval<const A>()());
+  /// Type of the array/view stored in the lazy object.
+  using stored_type = A;
 
-  /// View of the array/view to be reduced.
-  const_view_type rhs;
+  /// Array/View to be reduced.
+  stored_type rhs;
 
   /// MPI communicator.
   mpi::communicator comm;
@@ -87,7 +87,7 @@ struct mpi::lazy<mpi::tag::reduce, A> {
    */
   [[nodiscard]] auto shape() const {
     if ((comm.rank() == root) || all) return rhs.shape();
-    return std::array<long, const_view_type::rank>{};
+    return std::array<long, std::remove_cvref_t<stored_type>::rank>{};
   }
 
   /**

diff --git a/c++/nda/mpi/scatter.hpp b/c++/nda/mpi/scatter.hpp
@@ -54,11 +54,11 @@ struct mpi::lazy<mpi::tag::scatter, A> {
   /// Value type of the array/view.
   using value_type = typename std::decay_t<A>::value_type;
 
-  /// Const view type of the array/view stored in the lazy object.
-  using const_view_type = decltype(std::declval<const A>()());
+  /// Type of the array/view stored in the lazy object.
+  using stored_type = A;
 
-  /// View of the array/view to be scattered.
-  const_view_type rhs;
+  /// Array/View to be scattered.
+  stored_type rhs;
 
   /// MPI communicator.
   mpi::communicator comm;

diff --git a/test/c++/nda_mpi.cpp b/test/c++/nda_mpi.cpp
@@ -339,4 +339,21 @@ TEST_F(NDAMpi, VariousCollectiveCommunications) {
   EXPECT_ARRAY_NEAR(R2, comm.size() * A);
 }
 
+TEST_F(NDAMpi, PassingTemporaryObjects) {
+  auto A         = nda::array<int, 1>{1, 2, 3};
+  auto lazy_arr  = mpi::gather(nda::array<int, 1>{1, 2, 3}, comm);
+  auto res_arr   = nda::array<int, 1>(lazy_arr);
+  auto lazy_view = mpi::gather(A(), comm);
+  auto res_view  = nda::array<int, 1>(lazy_view);
+  if (comm.rank() == 0) {
+    for (long i = 0; i < comm.size(); ++i) {
+      EXPECT_ARRAY_EQ(res_arr(nda::range(i * 3, (i + 1) * 3)), A);
+      EXPECT_ARRAY_EQ(res_view(nda::range(i * 3, (i + 1) * 3)), A);
+    }
+  } else {
+    EXPECT_TRUE(res_arr.empty());
+    EXPECT_TRUE(res_view.empty());
+  }
+}
+
 MPI_TEST_MAIN