copy_to_buckets.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 "core.hpp"
#include "errors.hpp"
#include "introspect.hpp"
#include "processor.hpp"
 
#include <algorithm>
#include <cstddef>
#include <exception>
#include <memory>
#include <span>
#include <stdexcept>
#include <type_traits>
#include <utility>
 
namespace tcspc {
 
namespace internal {
 
template <typename T, typename Downstream> class copy_to_buckets {
    static_assert(processor<Downstream, bucket<T>>);
 
    std::shared_ptr<bucket_source<T>> bsource;
 
    Downstream downstream;
 
  public:
    explicit copy_to_buckets(std::shared_ptr<bucket_source<T>> buffer_provider,
                             Downstream downstream)
        : bsource(std::move(buffer_provider)),
          downstream(std::move(downstream)) {
        if (not bsource)
            throw std::invalid_argument(
                "copy_to_buckets buffer_provider must not be null");
    }
 
    [[nodiscard]] auto introspect_node() const -> processor_info {
        return processor_info(this, "copy_to_buckets");
    }
 
    [[nodiscard]] auto introspect_graph() const -> processor_graph {
        return downstream.introspect_graph().push_entry_point(this);
    }
 
    template <typename Event>
        requires std::is_constructible_v<std::span<T const>,
                                         std::remove_cvref_t<Event>>
    void handle(Event &&event) {
        auto const event_span = std::span<T const>(event);
        auto b = bsource->bucket_of_size(event_span.size());
        std::copy(event_span.begin(), event_span.end(), b.begin());
        downstream.handle(std::move(b));
    }
 
    template <typename Event>
        requires(not std::is_constructible_v<std::span<T const>,
                                             std::remove_cvref_t<Event>> and
                 handler_for<Downstream, std::remove_cvref_t<Event>>)
    void handle(Event &&event) {
        downstream.handle(std::forward<Event>(event));
    }
 
    void flush() { downstream.flush(); }
};
 
template <typename T, typename LiveDownstream, typename BatchDownstream>
class copy_to_full_buckets {
    static_assert(processor<LiveDownstream, bucket<T const>>);
    static_assert(processor<BatchDownstream, bucket<T>>);
 
    std::shared_ptr<bucket_source<T>> bsource;
    std::size_t bsize;
 
    bucket<T> bkt;
    std::size_t filled = 0;
 
    LiveDownstream live_downstream;
    BatchDownstream batch_downstream;
 
    // Mutates 'b' only when throwing.
    void emit_live(bucket<T> &b, std::size_t start, std::size_t count) {
        if (count > 0) {
            try {
                auto v = bsource->shared_view_of(b);
                v.shrink(start, count);
                live_downstream.handle(std::move(v));
            } catch (end_of_processing const &) {
                b.shrink(0, start + count);
                batch_downstream.handle(std::move(b));
                batch_downstream.flush();
                throw;
            }
        }
    }
 
    void emit_batch(bucket<T> &&b) {
        try {
            batch_downstream.handle(std::move(b));
        } catch (end_of_processing const &) {
            live_downstream.flush();
            throw;
        }
    }
 
    void flush_batch() {
        if (not bkt.empty() && filled > 0) {
            bkt.shrink(0, filled);
            batch_downstream.handle(std::move(bkt));
        }
        batch_downstream.flush();
    }
 
  public:
    explicit copy_to_full_buckets(
        std::shared_ptr<bucket_source<T>> buffer_provider,
        arg::batch_size<std::size_t> batch_size,
        LiveDownstream live_downstream, BatchDownstream batch_downstream)
        : bsource(std::move(buffer_provider)), bsize(batch_size.value),
          live_downstream(std::move(live_downstream)),
          batch_downstream(std::move(batch_downstream)) {
        if (not bsource)
            throw std::invalid_argument(
                "copy_to_full_buckets buffer_provider must not be null");
        if constexpr (not std::is_same_v<LiveDownstream, null_sink>) {
            if (not bsource->supports_shared_views())
                throw std::invalid_argument(
                    "copy_to_full_buckets buffer_provider must support shared views");
        }
        if (bsize == 0)
            throw std::invalid_argument(
                "copy_to_full_buckets batch size must be positive");
    }
 
    [[nodiscard]] auto introspect_node() const -> processor_info {
        return processor_info(this, "copy_to_full_buckets");
    }
 
    [[nodiscard]] auto introspect_graph() const -> processor_graph {
        return merge_processor_graphs(
            live_downstream.introspect_graph().push_entry_point(this),
            batch_downstream.introspect_graph().push_entry_point(this));
    }
 
    template <typename Event>
        requires std::is_constructible_v<std::span<T const>,
                                         std::remove_cvref_t<Event>>
    void handle(Event &&event) {
        auto src = std::span<T const>(event);
        while (not src.empty()) {
            if (filled == 0 && bkt.empty())
                bkt = bsource->bucket_of_size(bsize);
            auto const dest = std::span(bkt).subspan(filled);
            auto const copy_size = std::min(src.size(), dest.size());
            std::copy_n(src.begin(), copy_size, dest.begin());
            if constexpr (not std::is_same_v<LiveDownstream, null_sink>)
                emit_live(bkt, filled, copy_size);
            filled += copy_size;
            if (filled == bsize) {
                emit_batch(std::move(bkt));
                bkt = {};
                filled = 0;
            }
            src = src.subspan(copy_size);
        }
    }
 
    template <typename Event>
        requires(not std::is_constructible_v<std::span<T const>,
                                             std::remove_cvref_t<Event>> and
                 handler_for<LiveDownstream, std::remove_cvref_t<Event>>)
    void handle(Event &&event) {
        try {
            live_downstream.handle(std::forward<Event>(event));
        } catch (end_of_processing const &) {
            flush_batch();
            throw;
        }
    }
 
    void flush() {
        std::exception_ptr end;
        try {
            live_downstream.flush();
        } catch (end_of_processing const &) {
            end = std::current_exception();
        }
        flush_batch();
        if (end)
            std::rethrow_exception(end);
    }
};
 
} // namespace internal

template <typename T, typename Downstream>
auto copy_to_buckets(std::shared_ptr<bucket_source<T>> buffer_provider,
                     Downstream downstream) {
    return internal::copy_to_buckets<T, Downstream>(std::move(buffer_provider),
                                                    std::move(downstream));
}

template <typename T, typename LiveDownstream, typename BatchDownstream>
auto copy_to_full_buckets(std::shared_ptr<bucket_source<T>> buffer_provider,
                          arg::batch_size<std::size_t> batch_size,
                          LiveDownstream live_downstream,
                          BatchDownstream batch_downstream) {
    return internal::copy_to_full_buckets<T, LiveDownstream, BatchDownstream>(
        std::move(buffer_provider), batch_size, std::move(live_downstream),
        std::move(batch_downstream));
}
 
}; // namespace tcspc