batch_unbatch.hpp

/*
 * This file is part of libtcspc
 * Copyright 2019-2026 Board of Regents of the University of Wisconsin System
 * SPDX-License-Identifier: MIT
 */
 
#pragma once
 
#include "arg_wrappers.hpp"
#include "bucket.hpp"
#include "common.hpp"
#include "introspect.hpp"
#include "processor.hpp"
 
#include <cstddef>
#include <iterator>
#include <memory>
#include <stdexcept>
#include <type_traits>
#include <utility>
 
namespace tcspc {
 
namespace internal {
 
template <typename Event, typename Downstream>
    requires processor<Downstream, bucket<Event>>
class batch {
    static_assert(
        std::is_default_constructible_v<Event>,
        "batch requires Event to be default-constructible (buckets are "
        "pre-filled with default-constructed events)");
 
    std::shared_ptr<bucket_source<Event>> bsource;
    std::size_t bsize;
 
    bucket<Event> cur_bucket;
    std::size_t n_filled = 0;
 
    Downstream downstream;
 
  public:
    explicit batch(std::shared_ptr<bucket_source<Event>> buffer_provider,
                   arg::batch_size<std::size_t> batch_size,
                   Downstream downstream)
        : bsource(std::move(buffer_provider)), bsize(batch_size.value),
          downstream(std::move(downstream)) {
        if (bsize == 0)
            throw std::invalid_argument(
                "batch processor batch_size must not be zero");
    }
 
    [[nodiscard]] auto introspect_node() const -> processor_info {
        return processor_info(this, "batch");
    }
 
    [[nodiscard]] auto introspect_graph() const -> processor_graph {
        return downstream.introspect_graph().push_entry_point(this);
    }
 
    template <typename E>
        requires std::convertible_to<std::remove_cvref_t<E>, Event>
    void handle(E &&event) {
        if (cur_bucket.empty())
            cur_bucket = bsource->bucket_of_size(bsize);
 
        cur_bucket[n_filled] = std::forward<E>(event);
        ++n_filled;
 
        if (n_filled == bsize) {
            downstream.handle(std::move(cur_bucket));
            cur_bucket = {};
            n_filled = 0;
        }
    }
 
    void flush() {
        if (n_filled > 0) {
            cur_bucket.shrink(0, n_filled);
            downstream.handle(std::move(cur_bucket));
        }
        downstream.flush();
    }
};
 
template <typename ContainerEvent>
    requires std::is_same_v<typename std::iterator_traits<
                                decltype(std::declval<ContainerEvent>()
                                             .begin())>::value_type,
                            typename std::iterator_traits<
                                decltype(std::declval<ContainerEvent>()
                                             .end())>::value_type>
using unbatch_element_t = std::iterator_traits<
    decltype(std::declval<ContainerEvent>().end())>::value_type;
 
template <typename ContainerEvent, typename Downstream>
    requires processor<Downstream, unbatch_element_t<ContainerEvent>>
class unbatch {
    Downstream downstream;
 
  public:
    explicit unbatch(Downstream downstream)
        : downstream(std::move(downstream)) {}
 
    [[nodiscard]] auto introspect_node() const -> processor_info {
        return processor_info(this, "unbatch");
    }
 
    [[nodiscard]] auto introspect_graph() const -> processor_graph {
        return downstream.introspect_graph().push_entry_point(this);
    }
 
    // Should we mark this LIBTCSPC_NOINLINE? It would be good to increase the
    // chances that the downstream call will be inlined. But preliminary tests
    // (Apple clang 14 arm64) suggest that when the downstream is simple enough
    // to inline, it will be inlined, together with this loop, into upstream;
    // conversely, if the downstream is too complex to inline, it won't be
    // inlined even if this function is marked noinline. There may be
    // borderline cases where this doesn't hold, but it is probably best to
    // leave it to the compiler.
    template <typename E>
        requires std::convertible_to<std::remove_cvref_t<E>, ContainerEvent>
    void handle(E &&event) {
        if constexpr (std::is_lvalue_reference_v<E>) {
            for (auto const &e : event)
                downstream.handle(e);
        } else {
            for (auto &e : event)
                downstream.handle(std::move(e));
        }
    }
 
    template <typename E>
        requires(
            not std::convertible_to<std::remove_cvref_t<E>, ContainerEvent> and
            handler_for<Downstream, std::remove_cvref_t<E>>)
    void handle(E &&event) {
        downstream.handle(std::forward<E>(event));
    }
 
    void flush() { downstream.flush(); }
};
 
} // namespace internal

template <typename Event, typename Downstream>
auto batch(std::shared_ptr<bucket_source<Event>> buffer_provider,
           arg::batch_size<std::size_t> batch_size, Downstream downstream) {
    return internal::batch<Event, Downstream>(
        std::move(buffer_provider), batch_size, std::move(downstream));
}

template <typename ContainerEvent, typename Downstream>
auto unbatch(Downstream downstream) {
    return internal::unbatch<ContainerEvent, Downstream>(
        std::move(downstream));
}

template <typename Event, typename Downstream>
auto process_in_batches(arg::batch_size<std::size_t> batch_size,
                        Downstream downstream) {
    return batch<Event>(
        recycling_bucket_source<Event>::create(arg::max_bucket_count<>{1}),
        batch_size, unbatch<bucket<Event>>(std::move(downstream)));
}
 
} // namespace tcspc