Cleanup MPI buffer

physics/Disk/include/accretion.hpp

@@ -33,7 +33,7 @@ void computeAccretionCondition(size_t first, size_t last, Dataset& d, StarData&
 template<typename StarData>
 void exchangeAndAccreteOnStar(StarData& star, double minDt_m1, int rank)
 {
-    const auto [m_accreted, p_accreted, m_removed, p_removed] = buffered_mpi_allreduce_sum(
+    const auto [m_accreted, p_accreted, m_removed, p_removed] = buffer::mpiAllreduceSum(
         star.accreted_local.mass, star.accreted_local.momentum, star.removed_local.mass, star.removed_local.momentum);
 
     const double m_star_new = m_accreted + star.m;

physics/Disk/include/buffer_reduce.hpp

@@ -17,9 +17,11 @@
 
 namespace disk
 {
-
-template<typename Fn, typename Tval, size_t N, typename Tuple>
-requires(N == flattened_size_v<Tuple>) void for_each_buffer(Fn&& f, std::array<Tval, N>& buffer, Tuple&& args_tuple)
+namespace buffer
+{
+template<typename Fn, typename Tuple, array_type buffer_type>
+requires(std::tuple_size_v<std::decay_t<buffer_type>> ==
+         flattened_size_v<Tuple>) void for_each_buffer(Fn&& f, buffer_type&& buffer, Tuple&& args_tuple)
 {
     size_t i_buffer      = 0;
     auto   access_buffer = [&i_buffer, &f, &buffer](auto&& arg)
@@ -41,24 +43,46 @@ requires(N == flattened_size_v<Tuple>) void for_each_buffer(Fn&& f, std::array<T
     for_each_tuple([&](auto&& res) { access_buffer(std::forward<decltype(res)>(res)); }, args_tuple);
 }
 
-//! @brief Copy arguments of arithmetic type and of array of this type into a buffer;
-//! to collect multiple MPI calls into one.
-template<arithmetic_or_arrays... T>
-requires same_value_types<T...> auto buffered_mpi_allreduce_sum(const T&... args)
+template<bufferable_types... T>
+auto makeBuffer(const T&... args)
 {
     using value_type             = std::common_type_t<value_type_t<std::decay_t<T>>...>;
     constexpr size_t buffer_size = flattened_size_v<std::tuple<T...>>;
 
     std::array<value_type, buffer_size> buffer;
-
     for_each_buffer([](auto& buf_element, const auto& value) { buf_element = value; }, buffer, std::tie(args...));
+    return buffer;
+}
+
+template<bufferable_types... Ts, typename T, size_t N>
+requires(N == flattened_size_v<std::tuple<Ts...>>) auto extractBuffer(const std::array<T, N>& buffer)
+{
+    using RetTuple = std::tuple<std::remove_reference_t<Ts>...>;
+    if constexpr (sizeof...(Ts) == 1)
+    {
+        std::tuple_element_t<0, RetTuple> result;
+        for_each_buffer([](const auto& buf_element, auto& res) { res = buf_element; }, buffer, std::tie(result));
+        return result;
+    }
+    else
+    {
+        RetTuple result;
+        for_each_buffer([](const auto& buf_element, auto& res) { res = buf_element; }, buffer, result);
+        return result;
+    }
+}
 
-    MPI_Allreduce(MPI_IN_PLACE, buffer.data(), buffer_size, MpiType<value_type>{}, MPI_SUM, MPI_COMM_WORLD);
+//! @brief Copy arguments of the same arithmetic type and of array of this type into a buffer,
+//! to collect multiple MPI calls into one; returns a tuple if there is more than one argument.
+template<bufferable_types... T>
+auto mpiAllreduceSum(const T&... args)
+{
+    auto buffer = makeBuffer(args...);
 
-    std::tuple<std::remove_reference_t<T>...> result;
-    for_each_buffer([](const auto& buf_element, auto& res) { res = buf_element; }, buffer, result);
+    MPI_Allreduce(MPI_IN_PLACE, buffer.data(), buffer.size(), MpiType<value_type_t<decltype(buffer)>>{}, MPI_SUM,
+                  MPI_COMM_WORLD);
 
-    if constexpr (sizeof...(T) == 1) { return std::get<0>(result); }
-    else { return result; }
+    return extractBuffer<T...>(buffer);
 }
+} // namespace buffer
 } // namespace disk

physics/Disk/include/buffer_reduce_concepts.hpp

@@ -13,14 +13,15 @@
 
 namespace disk
 {
-
+namespace buffer
+{
 template<typename T>
 concept array_type = requires(T t, std::size_t i)
 {
     t[i];
     t.size();
     std::tuple_size_v<T>;
-    typename T::value_type;
+    typename std::remove_reference_t<T>::value_type;
 };
 
 template<typename T>
@@ -58,6 +59,9 @@ struct is_same_value_types<T1, Ts...> : std::bool_constant<(std::is_same_v<value
 template<typename... T>
 concept same_value_types = is_same_value_types<std::decay_t<T>...>::value;
 
+template<typename... T>
+concept bufferable_types = arithmetic_or_arrays<T...>&& same_value_types<T...>;
+
 template<typename... T>
 struct flattened_size : std::integral_constant<std::size_t, 0>
 {
@@ -82,4 +86,5 @@ struct flattened_size<std::tuple<T...>> : std::integral_constant<std::size_t, fl
 template<typename T>
 inline constexpr std::size_t flattened_size_v = flattened_size<std::decay_t<T>>::value;
 
+} // namespace buffer
 } // namespace disk

physics/Disk/include/exchange_star_position.hpp

@@ -4,8 +4,6 @@
 
 #pragma once
 
-#include <array>
-#include <mpi.h>
 #include "buffer_reduce.hpp"
 #include "cstone/primitives/mpi_wrappers.hpp"
 
@@ -18,7 +16,7 @@ void computeAndExchangeStarPosition(StarData& star, double dt, double dt_m1)
 {
     if (star.fixed_star == 1) { return; }
 
-    const auto global_force = buffered_mpi_allreduce_sum(star.force_local);
+    const auto global_force = buffer::mpiAllreduceSum(star.force_local);
 
     star.potential = global_force[0];