chunk_maintenance.h

#pragma once
 
#include <tess/core/shape.h>
#include <tess/experimental/registered_maintenance.h>
#include <tess/storage/chunk_meta.h>
#include <tess/storage/residency.h>
#include <tess/storage/sparse_world.h>
#include <tess/storage/world.h>
 
#include <atomic>
#include <concepts>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <limits>
#include <memory>
#include <optional>
#include <type_traits>
#include <utility>
 
namespace tess::experimental::maintenance {

enum class ChunkMarkResult : std::uint8_t {
  Accepted,
  CapacityExhausted,
  Stalled,
  InvalidMask,
  Missing,
  Released,
};

enum class ChunkProductState : std::uint8_t {
  Unavailable,
  Stale,
  Current,
};

struct ChunkProductToken {
  ChunkKey key{};
  ContentVersion content_version{};
  ResidencyGeneration residency_generation{};
 
  friend constexpr auto operator==(ChunkProductToken,
                                   ChunkProductToken) noexcept
      -> bool = default;
};

template <typename Product>
struct ChunkProductView {
  const Product* value = nullptr;
  ChunkProductToken token{};
  ChunkProductState state = ChunkProductState::Unavailable;
};

enum class ChunkResidencyStatus : std::uint8_t {
  Ready,
  Missing,
  NotIdle,
  Released,
};

struct ChunkResidencyResult {
  ChunkResidencyStatus status = ChunkResidencyStatus::Missing;
  ResidencyHandle handle{};
};

enum class ChunkEvictionResult : std::uint8_t {
  Evicted,
  Missing,
  NotIdle,
  Released,
};

enum class ChunkAdapterReleaseResult : std::uint8_t {
  Released,
  NotIdle,
  AlreadyReleased,
};

template <typename World, typename Product, typename Rebuild,
          typename Backend = DirtyBitScheduler>
  requires std::default_initializable<Product> &&
           std::invocable<Rebuild&, const World&, ChunkKey, DirtyObservation,
                          Product&>
class ChunkMaintenanceAdapter final {
  static constexpr bool is_dense =
      std::same_as<typename World::residency_type, AlwaysResident>;
  static constexpr bool is_sparse =
      std::same_as<typename World::residency_type, SparseResident>;
  static_assert(is_dense || is_sparse,
                "ChunkMaintenanceAdapter supports tess dense and sparse "
                "worlds only");
 
  struct Task final : MaintenanceTask {
    ChunkMaintenanceAdapter* owner = nullptr;
    std::size_t slot = 0;
 
    void run(MaintenanceBudget& budget) override {
      owner->run_slot(slot, budget);
    }
  };
 
  struct Slot {
    Product product{};
    Task task{};
    ChunkProductToken token{};
    ChunkKey key{};
    ResidencyGeneration residency_generation{};
    bool bound = false;
    bool built = false;
    std::atomic<bool> retry_debt = false;
  };
 
 public:
  ChunkMaintenanceAdapter(World& world, DirtyMask owned_dirty_mask,
                          Rebuild rebuild, std::size_t backend_capacity = 0)
      : world_(&world),
        owned_dirty_mask_(owned_dirty_mask),
        rebuild_(std::move(rebuild)),
        slot_count_(world_slot_count(world)),
        slots_(std::make_unique<Slot[]>(slot_count_)),
        scheduler_(slot_count_,
                   backend_capacity == 0 ? slot_count_ : backend_capacity),
        handles_(std::make_unique<MaintenanceHandle[]>(slot_count_)) {
    if (owned_dirty_mask_.empty()) {
      ::tess::detail::fail_fast(
          "ChunkMaintenanceAdapter requires a nonzero owned dirty mask");
    }
    for (std::size_t slot = 0; slot < slot_count_; ++slot) {
      slots_[slot].task.owner = this;
      slots_[slot].task.slot = slot;
      const auto handle = scheduler_.register_task(slots_[slot].task);
      if (!handle.has_value()) {
        ::tess::detail::fail_fast(
            "ChunkMaintenanceAdapter fixed registry capacity mismatch");
      }
      handles_[slot] = handle.value();
      if constexpr (is_dense) {
        bind_slot(slot, ChunkKey{slot}, {});
      }
    }
    if constexpr (is_sparse) {
      bind_current_sparse_residency();
    }
    scheduler_.seal();
  }
 
  ChunkMaintenanceAdapter(const ChunkMaintenanceAdapter&) = delete;
  auto operator=(const ChunkMaintenanceAdapter&)
      -> ChunkMaintenanceAdapter& = delete;
  ChunkMaintenanceAdapter(ChunkMaintenanceAdapter&&) = delete;
  auto operator=(ChunkMaintenanceAdapter&&)
      -> ChunkMaintenanceAdapter& = delete;
 
  ~ChunkMaintenanceAdapter() = default;

  [[nodiscard]] auto mark_dirty(ChunkKey key, DirtyMask mask, Box3 bounds)
      -> ChunkMarkResult {
    if (released_) {
      return ChunkMarkResult::Released;
    }
    if (mask.empty() || !(mask & ~owned_dirty_mask_).empty()) {
      return ChunkMarkResult::InvalidMask;
    }
    const auto slot = bound_slot(key);
    if (!slot.has_value()) {
      return ChunkMarkResult::Missing;
    }
    world_->mark_dirty(key, mask, bounds);
    return map_mark_result(schedule_slot(slot.value()));
  }

  [[nodiscard]] auto retry(ChunkKey key) -> std::optional<ScheduleResult> {
    if (released_) {
      return std::nullopt;
    }
    const auto slot = bound_slot(key);
    if (!slot.has_value()) {
      return std::nullopt;
    }
    return schedule_slot(slot.value());
  }

  [[nodiscard]] auto run_some(MaintenanceBudget budget) -> DrainResult {
    transition_ready_.store(false, std::memory_order_release);
#if TESS_HAS_EXCEPTIONS
    try {
      const auto backend_result = scheduler_.run_some(budget);
      const auto result = finalize_drain(backend_result, {});
      transition_ready_.store(result == DrainResult::Idle,
                              std::memory_order_release);
      return result;
    } catch (...) {
      transition_ready_.store(false, std::memory_order_release);
      throw;
    }
#else
    const auto backend_result = scheduler_.run_some(budget);
    const auto result = finalize_drain(backend_result, {});
    transition_ready_.store(result == DrainResult::Idle,
                            std::memory_order_release);
    return result;
#endif
  }

  [[nodiscard]] auto flush() -> DrainResult {
    transition_ready_.store(false, std::memory_order_release);
#if TESS_HAS_EXCEPTIONS
    try {
      static_cast<void>(reoffer_retry_debt());
      const auto backend_result = scheduler_.flush();
      const auto follow_up = reoffer_retry_debt();
      const auto result = finalize_drain(backend_result, follow_up);
      transition_ready_.store(result == DrainResult::Idle,
                              std::memory_order_release);
      return result;
    } catch (...) {
      transition_ready_.store(false, std::memory_order_release);
      throw;
    }
#else
    static_cast<void>(reoffer_retry_debt());
    const auto backend_result = scheduler_.flush();
    const auto follow_up = reoffer_retry_debt();
    const auto result = finalize_drain(backend_result, follow_up);
    transition_ready_.store(result == DrainResult::Idle,
                            std::memory_order_release);
    return result;
#endif
  }

  [[nodiscard]] auto product(ChunkKey key) const -> ChunkProductView<Product> {
    const auto slot_index = bound_slot(key);
    if (!slot_index.has_value()) {
      return {};
    }
    const auto& slot = slots_[slot_index.value()];
    if (!slot.built || slot.token.key != key) {
      return {};
    }
    const auto state = token_current(slot.token, slot_index.value())
                           ? ChunkProductState::Current
                           : ChunkProductState::Stale;
    return ChunkProductView<Product>{&slot.product, slot.token, state};
  }

  [[nodiscard]] auto current(ChunkProductToken token) const noexcept -> bool {
    const auto slot = bound_slot(token.key);
    return slot.has_value() && slots_[slot.value()].built &&
           slots_[slot.value()].token == token &&
           token_current(token, slot.value());
  }

  [[nodiscard]] auto metrics() const noexcept -> MaintenanceMetrics {
    return scheduler_.metrics();
  }

  [[nodiscard]] auto ensure_resident(ChunkKey key) -> ChunkResidencyResult
    requires(is_sparse)
  {
    if (released_) {
      return ChunkResidencyResult{ChunkResidencyStatus::Released, {}};
    }
    if (!transition_ready_.load(std::memory_order_acquire)) {
      return ChunkResidencyResult{ChunkResidencyStatus::NotIdle, {}};
    }
    const auto handle = world_->ensure_resident(key);
    if (!handle.generation.valid()) {
      return ChunkResidencyResult{ChunkResidencyStatus::Missing, {}};
    }
    const auto ref = world_->resident_ref(key);
    if (ref.meta == nullptr || ref.slot >= slot_count_ ||
        ref.generation != handle.generation) {
      ::tess::detail::fail_fast(
          "ChunkMaintenanceAdapter sparse residency binding mismatch");
    }
    bind_slot(ref.slot, key, ref.generation);
    return ChunkResidencyResult{ChunkResidencyStatus::Ready, handle};
  }

  [[nodiscard]] auto evict(ChunkKey key) -> ChunkEvictionResult
    requires(is_sparse)
  {
    if (released_) {
      return ChunkEvictionResult::Released;
    }
    if (!transition_ready_.load(std::memory_order_acquire)) {
      return ChunkEvictionResult::NotIdle;
    }
    const auto ref = world_->resident_ref(key);
    if (ref.meta == nullptr) {
      return ChunkEvictionResult::Missing;
    }
    if (!world_->evict(key)) {
      return ChunkEvictionResult::Missing;
    }
    unbind_slot(ref.slot);
    return ChunkEvictionResult::Evicted;
  }

  [[nodiscard]] auto reconcile_residency() -> ChunkResidencyStatus
    requires(is_sparse)
  {
    if (released_) {
      return ChunkResidencyStatus::Released;
    }
    if (!transition_ready_.load(std::memory_order_acquire)) {
      return ChunkResidencyStatus::NotIdle;
    }
    for (std::size_t slot = 0; slot < slot_count_; ++slot) {
      unbind_slot(slot);
    }
    bind_current_sparse_residency();
    return ChunkResidencyStatus::Ready;
  }

  [[nodiscard]] auto try_release() -> ChunkAdapterReleaseResult {
    if (released_) {
      return ChunkAdapterReleaseResult::AlreadyReleased;
    }
    if (!transition_ready_.load(std::memory_order_acquire)) {
      return ChunkAdapterReleaseResult::NotIdle;
    }
    for (std::size_t slot = 0; slot < slot_count_; ++slot) {
      const auto result = scheduler_.try_release(handles_[slot]);
      if (result != ReleaseResult::Released) {
        return ChunkAdapterReleaseResult::NotIdle;
      }
    }
    released_ = true;
    transition_ready_.store(false, std::memory_order_release);
    return ChunkAdapterReleaseResult::Released;
  }
 
 private:
  [[nodiscard]] static auto world_slot_count(const World& world)
      -> std::size_t {
    if constexpr (is_dense) {
      static_assert(
          World::chunk_count <=
          static_cast<std::uint64_t>(std::numeric_limits<std::size_t>::max()));
      return static_cast<std::size_t>(World::chunk_count);
    } else {
      return world.capacity();
    }
  }
 
  [[nodiscard]] auto bound_slot(ChunkKey key) const noexcept
      -> std::optional<std::size_t> {
    if constexpr (is_dense) {
      if (key.value >= World::chunk_count) {
        return std::nullopt;
      }
      return static_cast<std::size_t>(key.value);
    } else {
      const auto ref = world_->resident_ref(key);
      if (ref.meta == nullptr || ref.slot >= slot_count_) {
        return std::nullopt;
      }
      const auto& slot = slots_[ref.slot];
      if (!slot.bound || slot.key != key ||
          slot.residency_generation != ref.generation) {
        return std::nullopt;
      }
      return ref.slot;
    }
  }
 
  void bind_slot(std::size_t slot_index, ChunkKey key,
                 ResidencyGeneration generation) noexcept {
    auto& slot = slots_[slot_index];
    if (!slot.bound || slot.key != key ||
        slot.residency_generation != generation) {
      slot.built = false;
      slot.retry_debt.store(false, std::memory_order_release);
    }
    slot.bound = true;
    slot.key = key;
    slot.residency_generation = generation;
  }
 
  void unbind_slot(std::size_t slot_index) noexcept {
    auto& slot = slots_[slot_index];
    slot.bound = false;
    slot.built = false;
    slot.residency_generation = {};
    slot.retry_debt.store(false, std::memory_order_release);
  }
 
  void bind_current_sparse_residency()
    requires(is_sparse)
  {
    for (const auto key : world_->resident_chunk_keys()) {
      const auto ref = world_->resident_ref(key);
      if (ref.meta == nullptr || ref.slot >= slot_count_) {
        ::tess::detail::fail_fast(
            "ChunkMaintenanceAdapter sparse residency reconciliation "
            "mismatch");
      }
      bind_slot(ref.slot, key, ref.generation);
    }
  }
 
  [[nodiscard]] auto schedule_slot(std::size_t slot) -> ScheduleResult {
    transition_ready_.store(false, std::memory_order_release);
    const auto claimed_debt =
        slots_[slot].retry_debt.exchange(false, std::memory_order_acq_rel);
    const auto result = scheduler_.schedule(handles_[slot]);
    if (claimed_debt && result != ScheduleResult::Accepted) {
      slots_[slot].retry_debt.store(true, std::memory_order_release);
    }
    return result;
  }
 
  struct RetryOfferSummary {
    bool offered = false;
    bool stalled = false;
  };
 
  [[nodiscard]] auto reoffer_retry_debt() -> RetryOfferSummary {
    auto summary = RetryOfferSummary{};
    for (std::size_t slot = 0; slot < slot_count_; ++slot) {
      if (!slots_[slot].retry_debt.load(std::memory_order_acquire)) {
        continue;
      }
      summary.offered = true;
      const auto result = schedule_slot(slot);
      summary.stalled = summary.stalled || result == ScheduleResult::Stalled;
    }
    return summary;
  }
 
  [[nodiscard]] auto has_retry_debt() const noexcept -> bool {
    for (std::size_t slot = 0; slot < slot_count_; ++slot) {
      if (slots_[slot].retry_debt.load(std::memory_order_acquire)) {
        return true;
      }
    }
    return false;
  }
 
  [[nodiscard]] auto finalize_drain(DrainResult result,
                                    RetryOfferSummary follow_up) const noexcept
      -> DrainResult {
    if (result == DrainResult::Stalled || follow_up.stalled) {
      return DrainResult::Stalled;
    }
    if (follow_up.offered || has_retry_debt()) {
      return DrainResult::BudgetExhausted;
    }
    return result;
  }
 
  [[nodiscard]] static auto map_mark_result(ScheduleResult result) noexcept
      -> ChunkMarkResult {
    switch (result) {
      case ScheduleResult::Accepted:
        return ChunkMarkResult::Accepted;
      case ScheduleResult::CapacityExhausted:
        return ChunkMarkResult::CapacityExhausted;
      case ScheduleResult::Stalled:
        return ChunkMarkResult::Stalled;
    }
    ::tess::detail::fail_fast(
        "ChunkMaintenanceAdapter received an invalid schedule result");
  }
 
  [[nodiscard]] auto sparse_binding_current(
      const Slot& slot, std::size_t slot_index) const noexcept -> bool {
    if constexpr (is_dense) {
      static_cast<void>(slot);
      static_cast<void>(slot_index);
      return true;
    } else {
      const auto ref = world_->resident_ref(slot.key);
      return ref.meta != nullptr && ref.slot == slot_index &&
             ref.generation == slot.residency_generation;
    }
  }
 
  [[nodiscard]] auto token_current(ChunkProductToken token,
                                   std::size_t slot_index) const noexcept
      -> bool {
    const auto& slot = slots_[slot_index];
    if (!sparse_binding_current(slot, slot_index) || slot.key != token.key ||
        slot.residency_generation != token.residency_generation) {
      return false;
    }
    return world_->meta(token.key).content_version == token.content_version &&
           (world_->dirty_mask(token.key) & owned_dirty_mask_).empty();
  }
 
  void run_slot(std::size_t slot_index, MaintenanceBudget& budget) {
    auto& slot = slots_[slot_index];
    if (!slot.bound || !budget.consume()) {
      return;
    }
    if (!sparse_binding_current(slot, slot_index)) {
      return;
    }
    const auto observed = world_->observe_dirty(slot.key, owned_dirty_mask_);
    if (observed.mask.empty() &&
        (!slot.built ||
         slot.token.content_version == observed.content_version)) {
      slot.retry_debt.store(false, std::memory_order_release);
      return;
    }
    if (!sparse_binding_current(slot, slot_index)) {
      return;
    }
    std::invoke(rebuild_, std::as_const(*world_), slot.key, observed,
                slot.product);
    if (!sparse_binding_current(slot, slot_index)) {
      return;
    }
    slot.token = ChunkProductToken{slot.key, observed.content_version,
                                   slot.residency_generation};
    slot.built = true;
    if (world_->clear_dirty_observed(slot.key, observed)) {
      slot.retry_debt.store(false, std::memory_order_release);
      return;
    }
    if (schedule_slot(slot_index) != ScheduleResult::Accepted) {
      slot.retry_debt.store(true, std::memory_order_release);
    }
  }
 
  World* world_;
  DirtyMask owned_dirty_mask_;
  Rebuild rebuild_;
  std::size_t slot_count_;
  std::unique_ptr<Slot[]> slots_;
  RegisteredScheduler<Backend> scheduler_;
  std::unique_ptr<MaintenanceHandle[]> handles_;
  std::atomic<bool> transition_ready_ = false;
  bool released_ = false;
};
 
}  // namespace tess::experimental::maintenance