#pragma once

#include <bit>

#include "../utility/macros.hpp"
#include "../utility/traits.hpp"

/**
@file wsq.hpp
@brief work-stealing queue (wsq) include file
*/

#ifndef TF_DEFAULT_BOUNDED_TASK_QUEUE_LOG_SIZE 
  /**
  @def TF_DEFAULT_BOUNDED_TASK_QUEUE_LOG_SIZE
  
  This macro defines the default size of the bounded task queue in Log2. 
  Bounded task queue is used by each worker.
  By default, the value is set to 8, allowing the queue to hold 256 tasks.
  */
  #define TF_DEFAULT_BOUNDED_TASK_QUEUE_LOG_SIZE 8
#endif

#ifndef TF_DEFAULT_UNBOUNDED_TASK_QUEUE_LOG_SIZE 
  /**
  @def TF_DEFAULT_UNBOUNDED_TASK_QUEUE_LOG_SIZE
  
  This macro defines the default size of the unbounded task queue in Log2.
  Unbounded task queue is used by the executor as an overflow region.
  By default, the value is set to 10, allowing the queue to hold 1024 tasks initially.
  */
  #define TF_DEFAULT_UNBOUNDED_TASK_QUEUE_LOG_SIZE 10
#endif

namespace tf {

// ----------------------------------------------------------------------------
// Work-stealing queue steal protocol sentinels
//
// These free functions define the sentinel value used by steal operations
// across all work-stealing queue types (BoundedWSQ, UnboundedWSQ). They encode
// the result of a steal attempt into the return value itself, avoiding any
// out-parameter or separate status type.
//
// For pointer types T:
//   wsq_empty_value<T>() = nullptr  — queue was empty (or CAS was lost)
//
// The sentinel nullptr is returned whenever steal() cannot deliver a task,
// whether due to the queue being genuinely empty or losing a concurrent CAS
// race to another thief.
//
// For non-pointer types T, std::nullopt is returned to indicate the absence
// of a value.
//
// Queue classes expose this as a static member function (empty_value) that
// delegates here, so callers can use either form.
// ----------------------------------------------------------------------------

/**
@brief returns the empty sentinel for work-stealing steal operations

For pointer types @c T, returns @c nullptr. For non-pointer types, returns
@c std::nullopt. A steal operation returning this value means no task could
be stolen (either the queue was empty or the CAS was lost to another thief).
*/
template <typename T>
constexpr auto wsq_empty_value() {
  if constexpr (std::is_pointer_v<T>) {
    return T{nullptr};
  } else {
    return std::optional<T>{std::nullopt};
  }
}

// ----------------------------------------------------------------------------
// Unbounded Work-stealing Queue (WSQ)
// ----------------------------------------------------------------------------


/**
@class: UnboundedWSQ

@tparam T data type (must be a pointer type)

@brief class to create a lock-free unbounded work-stealing queue

This class implements the work-stealing queue described in the paper,
<a href="https://www.di.ens.fr/~zappa/readings/ppopp13.pdf">
Correct and Efficient Work-Stealing for Weak Memory Models</a>.

A work-stealing queue supports a single owner thread that performs push and pop
operations, while multiple concurrent thief threads may steal tasks
from the opposite end of the queue. The implementation is designed to
operate correctly under weak memory models and uses atomic operations
with carefully chosen memory orderings to ensure correctness and
scalability.

@dotfile images/unbounded_wsq.dot

Unlike bounded queues, this queue automatically grows its internal
storage as needed, allowing it to accommodate an arbitrary number of
tasks without a fixed capacity limit.
*/
template <typename T>
class UnboundedWSQ {

  struct Array {

    size_t C;
    size_t M;
    std::atomic<T>* S;

    explicit Array(size_t c) :
      C {c},
      M {c-1},
      S {new std::atomic<T>[C]} {
    }

    ~Array() {
      delete [] S;
    }

    size_t capacity() const noexcept {
      return C;
    }

    void push(int64_t i, T o) noexcept {
      S[i & M].store(o, std::memory_order_relaxed);
    }

    T pop(int64_t i) noexcept {
      return S[i & M].load(std::memory_order_relaxed);
    }

    Array* resize(int64_t b, int64_t t) {
      Array* ptr = new Array(2*C);
      for(int64_t i=t; i!=b; ++i) {
        ptr->push(i, pop(i));
      }
      return ptr;
    }

    Array* resize(int64_t b, int64_t t, size_t N) {
      // assert(N>0);
      Array* ptr = new Array(std::bit_ceil(C + N));
      for(int64_t i=t; i!=b; ++i) {
        ptr->push(i, pop(i));
      }
      return ptr;
    }

  };

  alignas(TF_CACHELINE_SIZE) std::atomic<int64_t> _top;
  alignas(TF_CACHELINE_SIZE) std::atomic<int64_t> _bottom;
  
  // Owner-private cached upper bound on _top.  Never read by thieves.
  // Because _top is never decremented, the real occupancy can only be
  // smaller than what is computed using this cached value, so using it
  // for the overflow check is always safe.
  int64_t _cached_top {0};

  // _array on its own cache line: avoids false-sharing with _bottom when
  // thieves load _array (consume) after reading _bottom (acquire).
  alignas(TF_CACHELINE_SIZE) std::atomic<Array*> _array;
  std::vector<Array*> _garbage;

  public:
  
  /**
  @brief the return type of queue operations
  
  `value_type` represents the type returned by `pop` and `steal` operations.
  For pointer element types `T`, it is `T` itself and uses `nullptr` to
  indicate an empty result. For non-pointer types, it is `std::optional<T>`,
  where `std::nullopt` denotes the absence of a value.

  @code{.cpp}
  static_assert(std::is_same_v<tf::UnboundedWSQ<int>::value_type, std::optional<int>>);
  static_assert(std::is_same_v<tf::UnboundedWSQ<int*>::value_type, nullptr);
  @endcode
  
  This design avoids the overhead of `std::optional` for pointer types
  while providing a uniform empty-result semantics.
  */
  using value_type = std::conditional_t<std::is_pointer_v<T>, T, std::optional<T>>;

  /**
  @brief constructs the queue with the given size in the base-2 logarithm

  @param LogSize the base-2 logarithm of the queue size

  @code{.cpp}
  tf::UnboundedWSQ<int> wsq(10);  
  assert(wsq.capacity() == 1024);
  @endcode
  */
  explicit UnboundedWSQ(int64_t LogSize = TF_DEFAULT_UNBOUNDED_TASK_QUEUE_LOG_SIZE);

  /**
  @brief destructs the queue
  */
  ~UnboundedWSQ();

  /**
  @brief queries if the queue is empty at the time of this call
  
  @code{.cpp}
  tf::UnboundedWSQ<int> wsq(10);  
  assert(wsq.empty() == true);
  wsq.push(1);
  assert(wsq.empty() == false);
  @endcode
  */
  bool empty() const noexcept;

  /**
  @brief queries the number of items at the time of this call
  
  @code{.cpp}
  tf::UnboundedWSQ<int> wsq(10);  
  assert(wsq.size() == 0);
  wsq.push(1);
  assert(wsq.size() == 1);
  @endcode
  */
  size_t size() const noexcept;

  /**
  @brief queries the capacity of the queue
  
  @code{.cpp}
  tf::UnboundedWSQ<int> wsq(10);  
  assert(wsq.capacity() == 1024);
  for(int i=0; i<1025; i++){  // insert more than 1024 ints to trigger resizing
    wsq.push(i);
  }
  assert(wsq.capacity() == 2048);
  assert(wsq.size() == 1025);
  @endcode
  */
  size_t capacity() const noexcept;
  
  /**
  @brief inserts an item to the queue

  @param item the item to push to the queue
  
  This method pushes one item into the queue.
  The operation can trigger the queue to resize its capacity if more space is required.
  
  @code{.cpp}
  tf::UnboundedWSQ<int> wsq(10);  
  assert(wsq.capacity() == 1024);
  for(int i=0; i<1025; i++) {   // insert more than 1024 items to trigger resizing
    wsq.push(i);
  }
  assert(wsq.capacity() == 2048);
  assert(wsq.size() == 1025);
  @endcode
  
  Only the owner thread can insert an item to the queue.
  */
  void push(T item);
  
  /**
  @brief tries to insert a batch of items into the queue

  @tparam I input iterator type
  @param first iterator to the first item in the batch
  @param N number of items to insert beginning at @p first

  This method pushes up to @p N items from the range `[first, first + N)` into the queue. 
  The operation can trigger the queue to resize its capacity 
  if more space is required.
  The iterator `first` is updated in place and will point to the next uninserted element after the call.

  Bulk insertion is often faster than inserting elements one by one because it requires fewer atomic operations.

  @code{.cpp}
  tf::UnboundedWSQ<int> wsq(10);  
  assert(wsq.capacity() == 1024);
  std::vector<int> vec(1025, 1);
  auto first = vec.data();
  wsq.bulk_push(first, vec.size()); 
  assert(wsq.capacity() == 2048);
  assert(wsq.size() == vec.size());
  assert(std::distance(vec.begin(), first) == 1025);
  @endcode
  
  Only the owner thread can insert an item to the queue.
  */
  template <typename I>
  void bulk_push(I& first, size_t N);

  /**
  @brief pops out an item from the queue

  This method pops an item from the queue.
  If the queue is empty, empty_value() is returned.
  The elements popped out from the queue follow a last-in-first-out (LIFO) order.
  
  @code{.cpp}
  tf::UnboundedWSQ<int> wsq(10);  
  wsq.push(1);
  wsq.push(2);
  wsq.push(3);
  assert(wsq.pop().value() = 3);
  assert(wsq.pop().value() = 2);
  assert(wsq.pop().value() = 1);
  assert(wsq.pop() == std::nullopt);
  @endcode
  
  Only the owner thread can pop out an item from the queue.
  */
  value_type pop();

  /**
  @brief steals an item from the queue

  Any threads can try to steal an item from the queue.
  The return can be an empty_value() if this operation failed.
  The elements stolen from the queue follow a first-in-first-out (FIFO) order.

  @code{.cpp}
  tf::UnboundedWSQ<int> wsq(10);
  wsq.push(1);
  wsq.push(2);
  wsq.push(3);
  assert(wsq.steal().value() = 1);
  assert(wsq.steal().value() = 2);
  assert(wsq.steal().value() = 3);
  assert(wsq.steal() == std::nullopt);
  @endcode

  Multiple threads can simultaneously steal items from the queue.
  */
  value_type steal();
  
  /**
  @brief returns the empty sentinel value for the queue element type
  
  This function provides a type-appropriate empty value used to indicate
  that a pop or steal operation failed. For pointer types, the empty value
  is `nullptr` of type `T`; for non-pointer types, it is `std::nullopt` of type `std::optional<T>`.
  
  The function is implemented as a `constexpr` helper to avoid additional
  storage, runtime overhead, or code duplication across queue operations.
  
  @return an empty `value_type` representing the absence of an element.
  */
  static constexpr auto empty_value()     { return wsq_empty_value<T>();     }

  private:

  Array* _resize_array(Array* a, int64_t b, int64_t t);
  Array* _resize_array(Array* a, int64_t b, int64_t t, size_t N);
};

// Constructor
template <typename T>
UnboundedWSQ<T>::UnboundedWSQ(int64_t LogSize) {
  _top.store(0, std::memory_order_relaxed);
  _bottom.store(0, std::memory_order_relaxed);
  _array.store(new Array{(size_t{1} << LogSize)}, std::memory_order_relaxed);
  _garbage.reserve(32);
}

// Destructor
template <typename T>
UnboundedWSQ<T>::~UnboundedWSQ() {
  for(auto a : _garbage) {
    delete a;
  }
  delete _array.load();
}

// Function: empty
template <typename T>
bool UnboundedWSQ<T>::empty() const noexcept {
  int64_t t = _top.load(std::memory_order_relaxed);
  int64_t b = _bottom.load(std::memory_order_relaxed);
  return (b <= t);
}

// Function: size
template <typename T>
size_t UnboundedWSQ<T>::size() const noexcept {
  int64_t t = _top.load(std::memory_order_relaxed);
  int64_t b = _bottom.load(std::memory_order_relaxed);
  return static_cast<size_t>(b >= t ? b - t : 0);
}

// Function: push
template <typename T>
void UnboundedWSQ<T>::push(T o) {

  int64_t b = _bottom.load(std::memory_order_relaxed);
  Array* a = _array.load(std::memory_order_relaxed);

  // queue is full with one additional item (b-t+1)
  if(a->capacity() < static_cast<size_t>(b - _cached_top + 1)) [[unlikely]] {
    _cached_top = _top.load(std::memory_order_acquire);
    if(a->capacity() < static_cast<size_t>(b - _cached_top + 1)) [[unlikely]] {
      a = _resize_array(a, b, _cached_top);
    }
  }

  a->push(b, o);
  std::atomic_thread_fence(std::memory_order_release);

  // original paper uses relaxed here but tsa complains
  _bottom.store(b + 1, std::memory_order_release);
}

// Function: bulk_push
template <typename T>
template <typename I>
void UnboundedWSQ<T>::bulk_push(I& first, size_t N) {

  if(N == 0) return;

  int64_t b = _bottom.load(std::memory_order_relaxed);
  Array* a = _array.load(std::memory_order_relaxed);

  // queue is full with N additional items
  if((b - _cached_top + N) > a->capacity()) [[unlikely]] {
    _cached_top = _top.load(std::memory_order_acquire);
    if((b - _cached_top + N) > a->capacity()) [[unlikely]] {
      a = _resize_array(a, b, _cached_top, N);
    }
  }

  for(size_t i=0; i<N; ++i) {
    a->push(b++, *first++);
  }
  std::atomic_thread_fence(std::memory_order_release);

  // original paper uses relaxed here but tsa complains
  _bottom.store(b, std::memory_order_release);
}

// Function: pop
template <typename T>
typename UnboundedWSQ<T>::value_type 
UnboundedWSQ<T>::pop() {

  int64_t b = _bottom.load(std::memory_order_relaxed) - 1;
  Array* a = _array.load(std::memory_order_relaxed);
  _bottom.store(b, std::memory_order_relaxed);
  std::atomic_thread_fence(std::memory_order_seq_cst);
  int64_t t = _top.load(std::memory_order_relaxed);

  //T item {nullptr};
  auto item = empty_value();

  if(t <= b) {
    item = a->pop(b);
    if(t == b) {
      // the last item just got stolen
      if(!_top.compare_exchange_strong(t, t+1, std::memory_order_seq_cst,
                                               std::memory_order_relaxed)) {
        //item = nullptr;
        item = empty_value();
      }
      _bottom.store(b + 1, std::memory_order_relaxed);
    }
  }
  else {
    _bottom.store(b + 1, std::memory_order_relaxed);
  }

  return item;
}

// Function: steal
template <typename T>
typename UnboundedWSQ<T>::value_type 
UnboundedWSQ<T>::steal() {
  
  int64_t t = _top.load(std::memory_order_acquire);
  std::atomic_thread_fence(std::memory_order_seq_cst);
  int64_t b = _bottom.load(std::memory_order_acquire);

  //T item {nullptr};
  auto item = empty_value();

  if(t < b) {
    Array* a = _array.load(std::memory_order_consume);
    item = a->pop(t);
    if(!_top.compare_exchange_strong(t, t+1,
                                     std::memory_order_seq_cst,
                                     std::memory_order_relaxed)) {
      //return nullptr;
      return empty_value();
    }
  }

  return item;
}


// Function: capacity
template <typename T>
size_t UnboundedWSQ<T>::capacity() const noexcept {
  return _array.load(std::memory_order_relaxed)->capacity();
}

template <typename T>
typename UnboundedWSQ<T>::Array*
UnboundedWSQ<T>::_resize_array(Array* a, int64_t b, int64_t t) {
  Array* tmp = a->resize(b, t);
  _garbage.push_back(a);
  // Note: the original paper using relaxed causes t-san to complain
  _array.store(tmp, std::memory_order_release);
  return tmp;
}

template <typename T>
typename UnboundedWSQ<T>::Array*
UnboundedWSQ<T>::_resize_array(Array* a, int64_t b, int64_t t, size_t N) {
  Array* tmp = a->resize(b, t, N);
  _garbage.push_back(a);
  // Note: the original paper using relaxed causes t-san to complain
  _array.store(tmp, std::memory_order_release);
  return tmp;
}

// ----------------------------------------------------------------------------
// Bounded Work-stealing Queue (WSQ)
// ----------------------------------------------------------------------------

/**
@class: BoundedWSQ

@tparam T data type
@tparam LogSize the base-2 logarithm of the queue size

@brief class to create a lock-free bounded work-stealing queue

This class implements the work-stealing queue described in the paper,
<a href="https://www.di.ens.fr/~zappa/readings/ppopp13.pdf">
Correct and Efficient Work-Stealing for Weak Memory Models</a>.

A work-stealing queue supports a single owner thread that performs push and pop
operations, while multiple concurrent thief threads may steal tasks
from the opposite end of the queue. The implementation is designed to
operate correctly under weak memory models and uses atomic operations
with carefully chosen memory orderings to ensure correctness and
scalability.

@dotfile images/bounded_wsq.dot

The queue has a fixed capacity determined at construction time and
does not grow dynamically. When the queue is full, push operations
may fail or require external handling.
*/
template <typename T, size_t LogSize = TF_DEFAULT_BOUNDED_TASK_QUEUE_LOG_SIZE>
class BoundedWSQ {
  
  constexpr static size_t BufferSize = size_t{1} << LogSize;
  constexpr static size_t BufferMask = (BufferSize - 1);

  static_assert((BufferSize >= 2) && ((BufferSize & (BufferSize - 1)) == 0));

  alignas(TF_CACHELINE_SIZE) std::atomic<int64_t> _top {0};
  alignas(TF_CACHELINE_SIZE) std::atomic<int64_t> _bottom {0};
  alignas(TF_CACHELINE_SIZE) std::atomic<T> _buffer[BufferSize];

  public:
  
  /**
  @brief the return type of queue operations
  
  `value_type` represents the type returned by `pop` and `steal` operations.
  For pointer element types `T`, it is `T` itself and uses `nullptr` to
  indicate an empty result. For non-pointer types, it is `std::optional<T>`,
  where `std::nullopt` denotes the absence of a value.
  
  @code{.cpp}
  static_assert(std::is_same_v<tf::UnboundedWSQ<int>::value_type, std::optional<int>>);
  static_assert(std::is_same_v<tf::UnboundedWSQ<int*>::value_type, nullptr);
  @endcode
  
  This design avoids the overhead of `std::optional` for pointer types
  while providing a uniform empty-result semantics.
  */
  using value_type = std::conditional_t<std::is_pointer_v<T>, T, std::optional<T>>;
    
  /**
  @brief constructs the queue with a given capacity
  
  @code{.cpp}
  tf::BoundedWSQ<int, 10> wsq;  
  static_assert(wsq.capacity() == 1024);
  @endcode
  */
  BoundedWSQ() = default;

  /**
  @brief destructs the queue
  */
  ~BoundedWSQ() = default;
  
  /**
  @brief queries if the queue is empty at the time of this call
  
  @code{.cpp}
  tf::BoundedWSQ<int, 10> wsq;  
  assert(wsq.empty() == true);
  wsq.push(1);
  assert(wsq.empty() == false);
  @endcode
  */
  bool empty() const noexcept;
  
  /**
  @brief queries the number of items at the time of this call
  
  @code{.cpp}
  tf::BoundedWSQ<int, 10> wsq;  
  assert(wsq.size() == 0);
  wsq.push(1);
  assert(wsq.size() == 1);
  @endcode
  */
  size_t size() const noexcept;

  /**
  @brief queries the capacity of the queue
  
  The capacity of a bounded work-stealing queue is decided at compile time.

  @code{.cpp}
  tf::BoundedWSQ<int, 10> wsq;
  static_assert(wsq.capacity() == 1024);
  @endcode
  */
  constexpr size_t capacity() const;
  
  /**
  @brief tries to insert an item to the queue

  @tparam O data type 
  @param item the item to perfect-forward to the queue
  @return `true` if the insertion succeed or `false` (queue is full)
  
  This method attempts to push one item into the queue.
  If the operation succeed, it returns `true` or `false` otherwise.
  
  @code{.cpp}
  tf::BoundedWSQ<int, 10> wsq;  
  static_assert(wsq.capacity() == 1024);
  for(int i=0; i<1024; i++) {
    assert(wsq.try_push(i) == true);
  }
  assert(wsq.size() == 1024);
  assert(wsq.try_push(0) == false);
  @endcode

  Only the owner thread can insert an item to the queue. 
  */
  template <typename O>
  bool try_push(O&& item);
  
  /**
  @brief tries to insert a batch of items into the queue

  @tparam I input iterator type
  @param first iterator to the first item in the batch
  @param N number of items to insert beginning at @p first
  @return the number of items successfully inserted

  This method attempts to push up to @p N items from the range
  `[first, first + N)` into the queue. Insertion stops early if the
  queue becomes full. 
  The iterator `first` is updated in place and will point to the next uninserted element after the call.

  Bulk insertion is often faster than inserting elements one by one because it requires fewer atomic operations.
  
  @code{.cpp}
  tf::BoundedWSQ<int, 10> wsq;  
  static_assert(wsq.capacity() == 1024);
  std::vector<int> vec(1030, 1);
  auto first = vec.begin();
  assert(wsq.try_bulk_push(first, vec.size()) == wsq.capacity());
  assert(std::distance(vec.begin(), first) == wsq.capacity());
  @endcode

  Only the owner thread can insert items into the queue.
  */
  template <typename I>
  size_t try_bulk_push(I& first, size_t N);
  
  /**
  @brief pops out an item from the queue

  The method pops an item out of the queue based on a last-in-first-out (LIFO) order.
  The return can be an empty_value() if this operation failed (empty queue).

  @code{.cpp}
  tf::BoundedWSQ<int, 10> wsq;  
  wsq.push(1);
  wsq.push(2);
  wsq.push(3);
  assert(wsq.pop().value() = 3);
  assert(wsq.pop().value() = 2);
  assert(wsq.pop().value() = 1);
  assert(wsq.pop() == std::nullopt);
  @endcode

  Only the owner thread can pop out an item from the queue. 
  */
  value_type pop();
  
  /**
  @brief steals an item from the queue

  Any threads can try to steal an item from the queue.
  The return can be an empty_value() if this operation failed.
  The elements stolen from the queue follow a first-in-first-out (FIFO) order.
  
  @code{.cpp}
  tf::BoundedWSQ<int, 10> wsq;  
  wsq.push(1);
  wsq.push(2);
  wsq.push(3);
  assert(wsq.steal().value() = 1);
  assert(wsq.steal().value() = 2);
  assert(wsq.steal().value() = 3);
  assert(wsq.steal() == std::nullopt);
  @endcode
  
  Multiple threads can simultaneously steal items from the queue.
  */
  value_type steal();


  /**
  @brief returns the empty sentinel value for the queue element type

  This function provides a type-appropriate empty value used to indicate
  that a pop or steal operation failed. For pointer types, the empty value
  is `nullptr` of type `T`; for non-pointer types, it is `std::nullopt` of type `std::optional<T>`.

  The function is implemented as a `constexpr` helper to avoid additional
  storage, runtime overhead, or code duplication across queue operations.

  @return an empty `value_type` representing the absence of an element.
  */
  static constexpr auto empty_value()     { return wsq_empty_value<T>();     }
};

// Function: empty
template <typename T, size_t LogSize>
bool BoundedWSQ<T, LogSize>::empty() const noexcept {
  int64_t t = _top.load(std::memory_order_relaxed);
  int64_t b = _bottom.load(std::memory_order_relaxed);
  return b <= t;
}

// Function: size
template <typename T, size_t LogSize>
size_t BoundedWSQ<T, LogSize>::size() const noexcept {
  int64_t t = _top.load(std::memory_order_relaxed);
  int64_t b = _bottom.load(std::memory_order_relaxed);
  return static_cast<size_t>(b >= t ? b - t : 0);
}

// Function: try_push
template <typename T, size_t LogSize>
template <typename O>
bool BoundedWSQ<T, LogSize>::try_push(O&& o) {

  int64_t b = _bottom.load(std::memory_order_relaxed);
  int64_t t = _top.load(std::memory_order_acquire);

  // queue is full with one additional item (b-t+1)
  if(static_cast<size_t>(b - t + 1) > BufferSize) [[unlikely]] {
    return false;
  }
  
  _buffer[b & BufferMask].store(std::forward<O>(o), std::memory_order_relaxed);

  std::atomic_thread_fence(std::memory_order_release);
  
  // original paper uses relaxed here but tsa complains
  _bottom.store(b + 1, std::memory_order_release);

  return true;
}

// Function: try_bulk_push
template <typename T, size_t LogSize>
template <typename I>
size_t BoundedWSQ<T, LogSize>::try_bulk_push(I& first, size_t N) {

  if(N == 0) return 0;

  int64_t b = _bottom.load(std::memory_order_relaxed);
  int64_t t = _top.load(std::memory_order_acquire);

  size_t r = BufferSize - (b - t);  // remaining capacity
  size_t n = std::min(N, r);        // number of pushable elements

  if(n > 0) {
    // push n elements into the queue
    for(size_t i=0; i<n; ++i) {
      _buffer[b++ & BufferMask].store(*first++, std::memory_order_relaxed);
    }
    std::atomic_thread_fence(std::memory_order_release);
    // original paper uses relaxed here but tsa complains
    _bottom.store(b, std::memory_order_release);
  }
  
  return n;
}

// Function: pop
template <typename T, size_t LogSize>
typename BoundedWSQ<T, LogSize>::value_type 
BoundedWSQ<T, LogSize>::pop() {

  int64_t b = _bottom.load(std::memory_order_relaxed) - 1;
  _bottom.store(b, std::memory_order_relaxed);
  std::atomic_thread_fence(std::memory_order_seq_cst);
  int64_t t = _top.load(std::memory_order_relaxed);

  //T item {nullptr};
  auto item = empty_value();

  if(t <= b) {
    item = _buffer[b & BufferMask].load(std::memory_order_relaxed);
    if(t == b) {
      // the last item just got stolen
      if(!_top.compare_exchange_strong(t, t+1, 
                                       std::memory_order_seq_cst, 
                                       std::memory_order_relaxed)) {
        //item = nullptr;
        item = empty_value();
      }
      _bottom.store(b + 1, std::memory_order_relaxed);
    }
  }
  else {
    _bottom.store(b + 1, std::memory_order_relaxed);
  }

  return item;
}

// Function: steal
template <typename T, size_t LogSize>
typename BoundedWSQ<T, LogSize>::value_type 
BoundedWSQ<T, LogSize>::steal() {
  int64_t t = _top.load(std::memory_order_acquire);
  std::atomic_thread_fence(std::memory_order_seq_cst);
  int64_t b = _bottom.load(std::memory_order_acquire);
  
  //T item{nullptr};
  auto item = empty_value();

  if(t < b) {
    item = _buffer[t & BufferMask].load(std::memory_order_relaxed);
    if(!_top.compare_exchange_strong(t, t+1,
                                     std::memory_order_seq_cst,
                                     std::memory_order_relaxed)) {
      //return nullptr;
      return empty_value();
    }
  }

  return item;
}


// Function: capacity
template <typename T, size_t LogSize>
constexpr size_t BoundedWSQ<T, LogSize>::capacity() const {
  return BufferSize;
}

// ----------------------------------------------------------------------------
// Work-stealing Queue Concept
// ----------------------------------------------------------------------------

/**
@brief concept to check if a type supports bounded work-stealing queue operations

A type @c Q satisfies @c BoundedWSQLike if it exposes the three core operations
required by a bounded work-stealing queue: owner-side pop, thief-side steal, and
a fallible owner-side push.

@par Requirements

- @c Q::value_type must be defined
- @c q.pop() must be valid
- @c q.steal() must be valid
- @c q.try_push(v) must return @c bool

@par Examples

@code{.cpp}
// BoundedWSQ satisfies BoundedWSQLike
static_assert(tf::BoundedWSQLike<tf::BoundedWSQ<int*>>);

// Custom queue type that satisfies the concept
struct MyQueue {
  using value_type = int*;
  value_type pop()           { return nullptr; }
  value_type steal()         { return nullptr; }
  bool try_push(value_type)  { return true; }
};
static_assert(tf::BoundedWSQLike<MyQueue>);
@endcode
*/
template <typename Q>
concept BoundedWSQLike = requires(Q& q, typename Q::value_type v) {
  { q.steal() };
  { q.pop() };
  { q.try_push(v) } -> std::same_as<bool>;
};

// ----------------------------------------------------------------------------
// Bounded Work-stealing Queue with Priority
// ----------------------------------------------------------------------------

/**
@private

@brief default priority function for BoundedPriorityWSQ

Calls `T::priority()` on the element and returns the result as
a @c size_t. Users may supply a custom priority function in place of this
default to decouple the queue from any particular element interface.
*/
struct DefaultPriorityFn {
  template <typename T>
  size_t operator()(const T& item) const {
    if constexpr (std::is_pointer_v<T>) {
      return static_cast<size_t>(item->priority());
    } else {
      return static_cast<size_t>(item.priority());
    }
  }
};

/**
@class BoundedPriorityWSQ

@brief class to create a lock-free bounded work-stealing queue with priority support

@tparam Q           underlying bounded work-stealing queue type; must satisfy @c BoundedWSQLike
@tparam MaxPriority number of priority levels (sub-queues); must be at least 1
@tparam PriorityFn  callable @c (const element&) -> size_t that maps an item to its
                    priority index; the returned value must be in @c [0, MaxPriority).
                    Defaults to @c DefaultPriorityFn.

@c BoundedPriorityWSQ wraps an array of @p MaxPriority sub-queues of type @p Q and
routes each item to the sub-queue that corresponds to its priority.
Priority levels are numbered from @c 0 to @c MaxPriority-1 in decreasing order
of urgency: priority @c 0 is the highest and priority @c MaxPriority-1 is the
lowest. Pop and steal operations always service sub-queues in ascending index
order, so higher-priority items are consumed before lower-priority ones.

@code{.cpp}
// 3-level priority queue with a custom priority function
auto priority_fn = [](tf::Node* n) -> size_t {
  // map application-level priority to [0, 3)
  return n->priority();  // must return 0, 1, or 2
};

tf::BoundedPriorityWSQ<tf::BoundedWSQ<tf::Node*>, 3, decltype(priority_fn)> pq(priority_fn);

tf::Node* n = ...;
pq.try_push(n);

// pop always returns the highest-priority available item
auto item = pq.pop();
@endcode
*/
template <BoundedWSQLike Q, size_t MaxPriority, typename PriorityFn = DefaultPriorityFn>
class BoundedPriorityWSQ {

  static_assert(MaxPriority >= 1);

  public:

  /**
  @brief the return type of pop and steal operations

  Inherits the @c value_type of the underlying queue @p Q.
  For pointer element types, it is the pointer type itself (using @c nullptr as the
  empty sentinel); for non-pointer types it is @c std::optional<T>.
  */
  using value_type = typename Q::value_type;

  /**
  @brief constructs the queue with an optional priority function

  @param fn priority function to use; defaults to a default-constructed @p PriorityFn
  */
  explicit BoundedPriorityWSQ(PriorityFn fn = {}) : _priority_fn(std::move(fn)) {}

  /**
  @brief tries to insert an item into the sub-queue determined by its priority

  @tparam O item type (forwarded)
  @param item the item to insert; @c _priority_fn(item) determines the target sub-queue
              and must return a value in @c [0, MaxPriority)
  @return @c true if the item was inserted, @c false if the target sub-queue is full

  @code{.cpp}
  auto fn = [](tf::Node* n) -> size_t { return n->priority(); };
  tf::BoundedPriorityWSQ<tf::BoundedWSQ<tf::Node*>, 3, decltype(fn)> pq(fn);
  tf::Node* n = ...;
  bool ok = pq.try_push(n);
  @endcode

  Only the owner thread may call this method.
  */
  template <typename O>
  bool try_push(O&& item) {
    return _wsqs[_priority_fn(item)].try_push(std::forward<O>(item));
  }

  /**
  @brief tries to insert a batch of items, routing contiguous same-priority runs
         to the corresponding sub-queue via bulk push

  @tparam I random-access iterator type
  @param first iterator to the first item; advanced in place by the number of items inserted
  @param N number of items to insert starting at @p first
  @return number of items successfully inserted

  Scans @p first for contiguous runs of items that map to the same sub-queue
  (as determined by @c _priority_fn) and delegates each run to the underlying
  queue's @c try_bulk_push (a single atomic bottom update per run). Stops
  immediately if a run is only partially inserted (the target sub-queue became
  full), leaving @p first pointing at the first uninserted item.

  @code{.cpp}
  auto fn = [](tf::Node* n) -> size_t { return n->priority(); };
  tf::BoundedPriorityWSQ<tf::BoundedWSQ<tf::Node*>, 3, decltype(fn)> pq(fn);

  // Suppose items have priorities: 0 0 0 1 1 2 1 1 1 1 2 2
  // Bulk push routes:
  //   queue[0] <- 3 items (priority 0 run)
  //   queue[1] <- 2 items (priority 1 run)
  //   queue[2] <- 1 item  (priority 2 run)
  //   queue[1] <- up to 4 items (priority 1 run); stops early if queue[1] is full
  std::vector<tf::Node*> batch = ...;
  auto it = batch.begin();
  size_t inserted = pq.try_bulk_push(it, batch.size());
  // it now points to batch.begin() + inserted
  @endcode

  Only the owner thread may call this method.
  */
  template <typename I>
  size_t try_bulk_push(I& first, size_t N) {

    if(N == 0) return 0;

    size_t remaining = N;

    while(remaining > 0) {
      size_t q = _priority_fn(first[0]);
      size_t run_len = 1;
      while(run_len < remaining && _priority_fn(first[run_len]) == q) {
        ++run_len;
      }
      size_t inserted = _wsqs[q].try_bulk_push(first, run_len);
      remaining -= inserted;
      if(inserted < run_len) {
        break;
      }
    }

    return N - remaining;
  }

  /**
  @brief pops an item from the highest-priority non-empty sub-queue

  @return the popped item, or @c empty_value() if all sub-queues are empty

  Sub-queues are checked in ascending index order (index 0 is highest priority).
  The element is removed from the owner (back) end of the selected sub-queue,
  following last-in-first-out (LIFO) order within each priority level.

  @code{.cpp}
  tf::BoundedPriorityWSQ<tf::BoundedWSQ<tf::Node*>, 3> pq;
  // ... push items of various priorities ...
  auto item = pq.pop();  // returns highest-priority available item
  if(item != tf::BoundedWSQ<tf::Node*>::empty_value()) {
    // process item
  }
  @endcode

  Only the owner thread may call this method.
  */
  value_type pop() {
    value_type result = Q::empty_value();
    for(size_t p=0; p<MaxPriority && !(result = _wsqs[p].pop()); ++p);
    return result;
  }

  /**
  @brief steals an item from the highest-priority non-empty sub-queue

  @return the stolen item, or @c empty_value() if all sub-queues are empty

  Sub-queues are checked in ascending index order (index 0 is highest priority).
  The element is removed from the thief (front) end of the selected sub-queue,
  following first-in-first-out (FIFO) order within each priority level.

  @code{.cpp}
  tf::BoundedPriorityWSQ<tf::BoundedWSQ<tf::Node*>, 3> pq;
  // ... push items of various priorities ...
  auto item = pq.steal();  // any thread may steal
  if(item != tf::BoundedWSQ<tf::Node*>::empty_value()) {
    // process item
  }
  @endcode

  Any thread may call this method concurrently.
  */
  value_type steal() {
    value_type result = Q::empty_value();
    for(size_t p=0; p<MaxPriority && !(result = _wsqs[p].steal()); ++p);
    return result;
  }

  /**
  @brief queries whether all sub-queues are empty at the time of this call

  @return @c true if every sub-queue is empty, @c false otherwise

  @code{.cpp}
  tf::BoundedPriorityWSQ<tf::BoundedWSQ<tf::Node*>, 3> pq;
  assert(pq.empty() == true);
  tf::Node* n = ...;
  pq.try_push(n);
  assert(pq.empty() == false);
  @endcode
  */
  bool empty() const noexcept {
    bool empty = true;
    for(size_t i = 0; i < MaxPriority && (empty = _wsqs[i].empty()); ++i);
    return empty;
  }

  /**
  @brief queries the total number of items across all sub-queues at the time of this call

  @return sum of sizes of all @p MaxPriority sub-queues

  @code{.cpp}
  tf::BoundedPriorityWSQ<tf::BoundedWSQ<tf::Node*>, 3> pq;
  assert(pq.size() == 0);
  tf::Node* n = ...;
  pq.try_push(n);
  assert(pq.size() == 1);
  @endcode
  */
  size_t size() const noexcept {
    size_t n = 0;
    for(size_t i = 0; i < MaxPriority; ++i) {
      n += _wsqs[i].size();
    }
    return n;
  }

  /**
  @brief queries the total capacity across all sub-queues

  @return the capacity of one sub-queue multiplied by @p MaxPriority

  @code{.cpp}
  // Each BoundedWSQ<Node*, 8> holds 256 items; 3 priority levels -> 768 total
  tf::BoundedPriorityWSQ<tf::BoundedWSQ<tf::Node*>, 3> pq;
  assert(pq.capacity() == 256 * 3);
  @endcode
  */
  constexpr size_t capacity() const {
    return _wsqs[0].capacity() * MaxPriority;
  }

  /**
  @brief returns a reference to the sub-queue at the given priority level

  @param priority priority index in @c [0, MaxPriority); @c 0 is the highest priority

  @code{.cpp}
  tf::BoundedPriorityWSQ<tf::BoundedWSQ<tf::Node*>, 3> pq;
  // directly access the high-priority sub-queue
  auto item = pq[0].pop();
  assert(pq[0].size() == 0);
  @endcode
  */
  Q& operator[](size_t priority) { return _wsqs[priority]; }

  /**
  @brief returns a const reference to the sub-queue at the given priority level

  @param priority priority index in @c [0, MaxPriority); @c 0 is the highest priority

  @code{.cpp}
  const tf::BoundedPriorityWSQ<tf::BoundedWSQ<tf::Node*>, 3> pq;
  assert(pq[0].empty() == true);
  @endcode
  */
  const Q& operator[](size_t priority) const { return _wsqs[priority]; }

  private:

  [[no_unique_address]] PriorityFn _priority_fn;
  std::array<Q, MaxPriority> _wsqs;
};


}  // end of namespace tf -----------------------------------------------------
