bucket.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 "errors.hpp"
#include "introspect.hpp"
#include "move_only_any.hpp"
#include "processor.hpp"
 
#include <algorithm>
#include <cassert>
#include <condition_variable>
#include <cstddef>
#include <functional>
#include <limits>
#include <memory>
#include <mutex>
#include <ostream>
#include <span>
#include <stdexcept>
#include <type_traits>
#include <typeinfo>
#include <utility>
#include <vector>
 
namespace tcspc {

template <typename T> class bucket {
    struct owning_storage {
        // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
        std::unique_ptr<T[]> p;
 
        // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
        explicit owning_storage(std::unique_ptr<T[]> ptr)
            : p(std::move(ptr)) {}
    };
 
    std::span<T> s;
    internal::move_only_any store;
 
  public:
    bucket() noexcept = default;

    template <typename S>
    explicit bucket(std::span<T> span, S &&storage)
        : s(span), store(std::forward<S>(storage)) {}
 
    ~bucket() = default;

    bucket(bucket const &other)
        : store(std::invoke([&s = other.s] {
              using TMut = std::remove_cv_t<T>;
              // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
              std::unique_ptr<TMut[]> r(s.empty() ? nullptr
                                                  : new TMut[s.size()]);
              std::copy(s.begin(), s.end(), r.get());
              return internal::move_only_any(
                  std::in_place_type<owning_storage>, std::move(r));
          })) {
        T *ptr = internal::move_only_any_cast<owning_storage>(&store)->p.get();
        s = {ptr, other.s.size()};
    }

    bucket(bucket &&other) noexcept = default;

    auto operator=(bucket const &other) -> bucket & {
        bucket t(other);
        std::swap(*this, t);
        return *this;
    }

    auto operator=(bucket &&other) noexcept -> bucket & = default;
 
    // Equivalents of most of span's members follow. But many of the member
    // functions are duplicated to work with a const bucket (because a const
    // bucket has read-only data, unlike a const span).

    using element_type = std::span<T>::element_type;
    using value_type = std::span<T>::value_type;
    using size_type = std::span<T>::size_type;
    using difference_type = std::span<T>::difference_type;
    using pointer = std::span<T>::pointer;
    using const_pointer = std::span<T>::const_pointer;
    using reference = std::span<T>::reference;
    using const_reference = std::span<T>::const_reference;
    using iterator = std::span<T>::iterator;
    using const_iterator = std::span<T const>::iterator;
    using reverse_iterator = std::span<T>::reverse_iterator;
    using const_reverse_iterator = std::span<T const>::reverse_iterator;

    [[nodiscard]] constexpr auto begin() noexcept -> iterator {
        return s.begin();
    }

    [[nodiscard]] constexpr auto cbegin() const noexcept -> const_iterator {
        return std::span<T const>(s).begin();
    }

    [[nodiscard]] constexpr auto begin() const noexcept -> const_iterator {
        return cbegin();
    }

    [[nodiscard]] constexpr auto end() noexcept -> iterator { return s.end(); }

    [[nodiscard]] constexpr auto cend() const noexcept -> const_iterator {
        return std::span<T const>(s).end();
    }

    [[nodiscard]] constexpr auto end() const noexcept -> const_iterator {
        return cend();
    }

    [[nodiscard]] constexpr auto rbegin() noexcept -> reverse_iterator {
        return s.rbegin();
    }

    [[nodiscard]] constexpr auto crbegin() const noexcept
        -> const_reverse_iterator {
        return std::span<T const>(s).rbegin();
    }

    [[nodiscard]] constexpr auto rbegin() const noexcept
        -> const_reverse_iterator {
        return crbegin();
    }

    [[nodiscard]] constexpr auto rend() noexcept -> reverse_iterator {
        return s.rend();
    }

    [[nodiscard]] constexpr auto crend() const noexcept
        -> const_reverse_iterator {
        return std::span<T const>(s).rend();
    }

    [[nodiscard]] constexpr auto rend() const noexcept
        -> const_reverse_iterator {
        return crend();
    }

    [[nodiscard]] constexpr auto front() -> reference { return s.front(); }

    [[nodiscard]] constexpr auto front() const -> const_reference {
        return std::span<T const>(s).front();
    }

    [[nodiscard]] constexpr auto back() -> reference { return s.back(); }

    [[nodiscard]] constexpr auto back() const -> const_reference {
        return std::span<T const>(s).back();
    }

    [[nodiscard]] constexpr auto operator[](size_type idx) -> reference {
        return s[idx];
    }

    [[nodiscard]] constexpr auto operator[](size_type idx) const
        -> const_reference {
        return s[idx];
    }

    [[nodiscard]] constexpr auto at(size_type pos) -> reference {
        if (pos >= size())
            throw std::out_of_range("bucket element index out of range");
        return s[pos];
    }

    [[nodiscard]] constexpr auto at(size_type pos) const -> const_reference {
        if (pos >= size())
            throw std::out_of_range("bucket element index out of range");
        return s[pos];
    }

    [[nodiscard]] constexpr auto data() noexcept -> pointer {
        return s.data();
    }

    [[nodiscard]] constexpr auto data() const noexcept -> const_pointer {
        return s.data();
    }

    [[nodiscard]] constexpr auto size() const noexcept -> size_type {
        return s.size();
    }

    [[nodiscard]] constexpr auto size_bytes() const noexcept -> size_type {
        return s.size_bytes();
    }

    [[nodiscard]] constexpr auto empty() const noexcept -> bool {
        return s.empty();
    }

    [[nodiscard]] constexpr auto first(size_type count) -> std::span<T> {
        return s.first(count);
    }

    [[nodiscard]] constexpr auto first(size_type count) const
        -> std::span<T const> {
        return std::span<T const>(s).first(count);
    }

    [[nodiscard]] constexpr auto last(size_type count) -> std::span<T> {
        return s.last(count);
    }

    [[nodiscard]] constexpr auto last(size_type count) const
        -> std::span<T const> {
        return std::span<T const>(s).last(count);
    }

    [[nodiscard]] constexpr auto subspan(size_type offset,
                                         size_type count = std::dynamic_extent)
        -> std::span<T> {
        return s.subspan(offset, count);
    }

    [[nodiscard]] constexpr auto
    subspan(size_type offset, size_type count = std::dynamic_extent) const
        -> std::span<T const> {
        return std::span<T const>(s).subspan(offset, count);
    }
 
    // End of span member equivalents.
 
    // We do not prevent access to storage of owning instance, but the value
    // inside the move_only_any will not be accessible because owning_storage
    // is a private type.

    template <typename S>
    [[nodiscard]] auto check_storage_type() const noexcept -> bool {
        return store.type() == typeid(S);
    }

    template <typename S> [[nodiscard]] auto storage() const -> S const & {
        return internal::move_only_any_cast<S const &>(store);
    }

    template <typename S> [[nodiscard]] auto extract_storage() -> S {
        s = {};
        S ret = internal::move_only_any_cast<S>(std::move(store));
        store = {};
        return ret;
    }

    void shrink(std::size_t start, std::size_t count = std::dynamic_extent) {
        s = s.subspan(start, count);
    }
 
    // In contrast to span, bucket equality is deep equality.

    friend constexpr auto operator==(bucket const &lhs, bucket const &rhs)
        -> bool {
        return lhs.s.size() == rhs.s.size() &&
               std::equal(lhs.s.begin(), lhs.s.end(), rhs.s.begin());
    }

    friend auto operator<<(std::ostream &stream, bucket const &bkt)
        -> std::ostream & {
        static constexpr std::size_t num_to_print = 10;
        auto const size = bkt.s.size();
        stream << "bucket(size=" << size;
        if constexpr (std::is_same_v<std::remove_cv_t<T>, std::byte>) {
            for (std::size_t i = 0; i < std::min(size, num_to_print - 1); ++i)
                stream << ", " << std::to_integer<int>(bkt.s[i]);
            if (size > num_to_print)
                stream << ", ...";
            if (size >= num_to_print)
                stream << ", " << std::to_integer<int>(bkt.s[size - 1]);
        } else {
            for (std::size_t i = 0; i < std::min(size, num_to_print - 1); ++i)
                stream << ", " << bkt.s[i];
            if (size > num_to_print)
                stream << ", ...";
            if (size >= num_to_print)
                stream << ", " << bkt.s[size - 1];
        }
        return stream << ')';
    }
};

template <typename T>
[[nodiscard]] auto ad_hoc_bucket(std::span<T> s) -> bucket<T> {
    struct ad_hoc_storage {};
    return bucket<T>(s, ad_hoc_storage{});
}

template <typename T> struct bucket_source {
    virtual ~bucket_source() = default;

    virtual auto bucket_of_size(std::size_t size) -> bucket<T> = 0;

    [[nodiscard]] virtual auto supports_shared_views() const noexcept -> bool {
        return false;
    }

    [[nodiscard]] virtual auto
    shared_view_of([[maybe_unused]] bucket<T> const &bkt) -> bucket<T const> {
        throw std::logic_error(
            "this bucket source does not support shared views");
    }
};

template <typename T>
class new_delete_bucket_source final : public bucket_source<T> {
    new_delete_bucket_source() = default;
 
  public:
    static auto create() -> std::shared_ptr<bucket_source<T>> {
        static std::shared_ptr<bucket_source<T>> instance(
            new new_delete_bucket_source());
        return instance;
    }

    auto bucket_of_size(std::size_t size) -> bucket<T> override {
        // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
        std::unique_ptr<T[]> p(new T[size]);
        return bucket<T>{std::span(p.get(), size), std::move(p)};
    }
};

template <typename T>
class sharable_new_delete_bucket_source final : public bucket_source<T> {
    sharable_new_delete_bucket_source() = default;
 
  public:
    static auto create() -> std::shared_ptr<bucket_source<T>> {
        static std::shared_ptr<bucket_source<T>> instance(
            new sharable_new_delete_bucket_source());
        return instance;
    }

    auto bucket_of_size(std::size_t size) -> bucket<T> override {
        // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
        std::shared_ptr<T[]> p(new T[size]);
        return bucket<T>{std::span(p.get(), size), std::move(p)};
    }

    [[nodiscard]] auto supports_shared_views() const noexcept
        -> bool override {
        return true;
    }

    [[nodiscard]] auto shared_view_of(bucket<T> const &bkt)
        -> bucket<T const> override {
        // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
        auto storage = bkt.template storage<std::shared_ptr<T[]>>();
        return bucket<T const>{std::span<T const>(bkt), std::move(storage)};
    }
};

template <typename T, bool Blocking = false, bool ClearRecycled = false>
class recycling_bucket_source final
    : public bucket_source<T>,
      public std::enable_shared_from_this<
          recycling_bucket_source<T, Blocking, ClearRecycled>> {
    std::mutex mutex;
    std::condition_variable not_empty_condition;
    std::size_t max_buckets;
    std::size_t max_recycled;
    std::size_t bucket_count = 0;
    std::vector<std::unique_ptr<std::vector<T>>> recyclable;
 
    struct bucket_storage {
        std::shared_ptr<recycling_bucket_source> source;
        std::unique_ptr<std::vector<T>> storage;
 
        ~bucket_storage() {
            if (not source)
                return;
 
            assert(storage);
 
            // If "large" bucket, release the underlying storage (but do go on
            // to "recycle", so that a thread waiting on bucket_of_size() will
            // get a bucket).
            if (source->max_recycled > 0 &&
                storage->size() > source->max_recycled)
                *storage = std::vector<T>();
 
            if constexpr (ClearRecycled)
                storage->clear();
 
            {
                auto const lock = std::lock_guard(source->mutex);
                source->recyclable.push_back(std::move(storage));
            }
 
            if constexpr (Blocking) {
                // Is it safe to wake up the source when this bucket_storage
                // might be the last thing holding on to it? Yes, because if
                // that is the case there is no thread waiting on the source.
                source->not_empty_condition.notify_one();
            }
        }
 
        explicit bucket_storage(
            std::shared_ptr<recycling_bucket_source> source,
            std::unique_ptr<std::vector<T>> &&storage)
            : source(std::move(source)), storage(std::move(storage)) {}
 
        bucket_storage(bucket_storage const &) = delete;
        auto operator=(bucket_storage const &) = delete;
 
        bucket_storage(bucket_storage &&) noexcept = default;
        auto operator=(bucket_storage &&) noexcept
            -> bucket_storage & = default;
    };
 
    explicit recycling_bucket_source(
        arg::max_bucket_count<> max_bucket_count,
        arg::max_recycled_size<> max_recycled_size)
        : max_buckets(max_bucket_count.value),
          max_recycled(max_recycled_size.value) {}
 
  public:
    static auto create(
        arg::max_bucket_count<> max_bucket_count =
            arg::max_bucket_count{std::numeric_limits<std::size_t>::max()},
        arg::max_recycled_size<> max_recycled_size = arg::max_recycled_size<>{
            0}) -> std::shared_ptr<bucket_source<T>> {
        return std::shared_ptr<recycling_bucket_source>(
            new recycling_bucket_source(max_bucket_count, max_recycled_size));
    }

    auto bucket_of_size(std::size_t size) -> bucket<T> override {
        std::unique_ptr<std::vector<T>> p;
        {
            auto lock = std::unique_lock(mutex);
            if (recyclable.empty() && bucket_count < max_buckets) {
                ++bucket_count;
            } else {
                if constexpr (Blocking) {
                    not_empty_condition.wait(
                        lock, [&] { return not recyclable.empty(); });
                } else if (recyclable.empty()) {
                    throw buffer_overflow_error(
                        "recycling bucket source exhausted");
                }
                p = std::move(recyclable.back());
                recyclable.pop_back();
            }
        }
        if (not p)
            p = std::make_unique<std::vector<T>>();
        p->resize(size);
        auto const spn = std::span(p->data(), p->size());
        return bucket<T>{
            spn, bucket_storage(this->shared_from_this(), std::move(p))};
    }
};

template <typename T, bool Blocking = false, bool ClearRecycled = false>
class sharable_recycling_bucket_source final
    : public bucket_source<T>,
      public std::enable_shared_from_this<
          sharable_recycling_bucket_source<T, Blocking, ClearRecycled>> {
    std::mutex mutex;
    std::condition_variable not_empty_condition;
    std::size_t max_buckets;
    std::size_t max_recycled;
    std::size_t bucket_count = 0;
    std::vector<std::unique_ptr<std::vector<T>>> recyclable;
 
    // Wrap storage in private type to forbid observation/extraction.
    struct bucket_storage {
        std::shared_ptr<std::vector<T>> storage; // With custom deleter.
    };
 
    explicit sharable_recycling_bucket_source(
        arg::max_bucket_count<> max_bucket_count,
        arg::max_recycled_size<> max_recycled_size)
        : max_buckets(max_bucket_count.value),
          max_recycled(max_recycled_size.value) {}
 
  public:
    static auto create(
        arg::max_bucket_count<> max_bucket_count =
            arg::max_bucket_count{std::numeric_limits<std::size_t>::max()},
        arg::max_recycled_size<> max_recycled_size = arg::max_recycled_size<>{
            0}) -> std::shared_ptr<bucket_source<T>> {
        return std::shared_ptr<sharable_recycling_bucket_source>(
            new sharable_recycling_bucket_source(max_bucket_count,
                                                 max_recycled_size));
    }

    auto bucket_of_size(std::size_t size) -> bucket<T> override {
        std::unique_ptr<std::vector<T>> p;
        {
            auto lock = std::unique_lock(mutex);
            if (recyclable.empty() && bucket_count < max_buckets) {
                ++bucket_count;
            } else {
                if constexpr (Blocking) {
                    not_empty_condition.wait(
                        lock, [&] { return not recyclable.empty(); });
                } else if (recyclable.empty()) {
                    throw buffer_overflow_error(
                        "sharable recycling bucket source exhausted");
                }
                p = std::move(recyclable.back());
                recyclable.pop_back();
            }
        }
        if (not p)
            p = std::make_unique<std::vector<T>>();
        p->resize(size);
        auto const spn = std::span(*p);
        std::shared_ptr<std::vector<T>> shptr{
            p.release(),
            [self = this->shared_from_this()](std::vector<T> *pv) {
                if (not pv)
                    return;
                if (self->max_recycled > 0 && pv->size() > self->max_recycled)
                    *pv = std::vector<T>();
                if constexpr (ClearRecycled)
                    pv->clear();
                {
                    auto const lock = std::lock_guard(self->mutex);
                    self->recyclable.emplace_back(pv);
                }
                if constexpr (Blocking)
                    self->not_empty_condition.notify_one();
            }
 
        };
        return bucket<T>{spn, bucket_storage{shptr}};
    }

    [[nodiscard]] auto supports_shared_views() const noexcept
        -> bool override {
        return true;
    }

    [[nodiscard]] auto shared_view_of(bucket<T> const &bkt)
        -> bucket<T const> override {
        auto storage = bkt.template storage<bucket_storage>();
        return bucket<T const>{std::span<T const>(bkt), std::move(storage)};
    }
};
 
namespace internal {
 
template <typename Event, typename Downstream> class extract_bucket {
    static_assert(
        processor<Downstream, decltype(std::declval<Event>().data_bucket)>);
 
    Downstream downstream;
 
  public:
    explicit extract_bucket(Downstream downstream)
        : downstream(std::move(downstream)) {}
 
    [[nodiscard]] auto introspect_node() const -> processor_info {
        return processor_info(this, "extract_bucket");
    }
 
    [[nodiscard]] auto introspect_graph() const -> processor_graph {
        return downstream.introspect_graph().push_entry_point(this);
    }
 
    void handle(Event const &event) { downstream.handle(event.data_bucket); }
 
    void handle(Event &&event) {
        downstream.handle(std::move(event).data_bucket);
    }
 
    void flush() { downstream.flush(); }
};
 
} // namespace internal

template <typename Event, typename Downstream>
auto extract_bucket(Downstream downstream) {
    return internal::extract_bucket<Event, Downstream>(std::move(downstream));
}
 
} // namespace tcspc