work_efficient_async_inclusive_scan.cpp

template <typename InputIt, typename OutputIt, typename BinaryOp, typename T, typename Size>
unique_future<OutputIt>
async_inclusive_scan(InputIt first, InputIt last, OutputIt output,BinaryOp op, T init, Size chunk_size)
{
  Size const elements = std::distance(first, last);
  Size const chunks   = (1 + ((elements - 1) / chunk_size)); // Round up.

  std::vector<unique_future<T>> sweep;
  sweep.reserve(chunks);

  // Upsweep.
  for (Size chunk = 0; chunk < chunks; ++chunk)
    sweep.emplace_back(async([=] {
      auto const this_begin = chunk * chunk_size;
      auto const this_end   = std::min(elements, (chunk + 1) * chunk_size);
      return T(*--std::inclusive_scan(first + this_begin, first + this_end,
                                      output + this_begin, op));
    }));

  auto sums = co_await when_all(sweep);

  // We add in init here.
  std::inclusive_scan(sums.begin(), sums.end(), sums.begin(), op, init);

  sweep.clear();

  // Downsweep.
  for (Size chunk = 1; chunk < chunks; ++chunk)
    // Capturing sums by reference is fine here; we're going to co_await on the
    // futures, which will keep the current scope alive for them.
    sweep.emplace_back(async([=, &sums] {
      auto const this_begin = chunk * chunk_size;
      auto const this_end   = std::min(elements, (chunk + 1) * chunk_size);
      std::for_each(output + this_begin, output + this_end,
                    [=, &sums] (auto& t) { t = op(std::move(t), sums[chunk - 1]); });
      return *(output + this_end - 1);
    }));

  co_await when_all(sweep);

  co_return output + elements;
}