//  Copyright (c) Meta Platforms, Inc. and affiliates.
//  This source code is licensed under both the GPLv2 (found in the
//  COPYING file in the root directory) and Apache 2.0 License
//  (found in the LICENSE.Apache file in the root directory).

#pragma once

#include <atomic>
#include <memory>

#include "rocksdb/rocksdb_namespace.h"

#if defined(__clang__)
#define ROCKSDB_SUPPRESS_SHARED_PTR_ATOMIC_DEPRECATION_PUSH \
  _Pragma("clang diagnostic push")                          \
      _Pragma("clang diagnostic ignored \"-Wdeprecated-declarations\"")
#define ROCKSDB_SUPPRESS_SHARED_PTR_ATOMIC_DEPRECATION_POP \
  _Pragma("clang diagnostic pop")
#elif defined(__GNUC__)
#define ROCKSDB_SUPPRESS_SHARED_PTR_ATOMIC_DEPRECATION_PUSH \
  _Pragma("GCC diagnostic push")                            \
      _Pragma("GCC diagnostic ignored \"-Wdeprecated-declarations\"")
#define ROCKSDB_SUPPRESS_SHARED_PTR_ATOMIC_DEPRECATION_POP \
  _Pragma("GCC diagnostic pop")
#else
#define ROCKSDB_SUPPRESS_SHARED_PTR_ATOMIC_DEPRECATION_PUSH
#define ROCKSDB_SUPPRESS_SHARED_PTR_ATOMIC_DEPRECATION_POP
#endif

namespace ROCKSDB_NAMESPACE {

// Background:
// std::atomic is somewhat easy to misuse:
// * Implicit conversion to T makes it easy to use an unnecessarily strong
// memory ordering (std::memory_order_seq_cst) and to hide atomic operations
// that should be evident on reading the code.
// * Similarly, defaulting to std::memory_order_seq_cst for atomic operations
// makes it easy to use unnecessarily strong orderings. (It's always safe if
// some ordering is safe, but it's better to be intentional and thoughtful when
// carefully optimizing code with atomics.) Legitimate needs for seq_cst vs.
// acq_rel are rare, such as drawing inferences across two atomics in
// implementing hazard pointers.
// * It's easy to use nonsensical (UB) combinations like store with
// std::memory_order_acquire. Getting these right in development is an
// unnecessary cognitive overhead even if they are caught by UBSAN.
//
// For such reasons, we provide wrappers below to make clear and explicit
// usage of atomics easier.

// Wrapper around std::atomic for better code clarity (see Background above).
//
// This relaxed-only wrapper is intended for atomics that are not used to
// synchronize other data across threads (only the atomic data), so can always
// used relaxed memory ordering. For example, a cross-thread counter that never
// returns the same result can be a RelaxedAtomic.
template <typename T>
class RelaxedAtomic {
 public:
  explicit RelaxedAtomic(T initial = {}) : v_(initial) {}
  void StoreRelaxed(T desired) { v_.store(desired, std::memory_order_relaxed); }
  T LoadRelaxed() const { return v_.load(std::memory_order_relaxed); }
  bool CasWeakRelaxed(T& expected, T desired) {
    return v_.compare_exchange_weak(expected, desired,
                                    std::memory_order_relaxed);
  }
  bool CasStrongRelaxed(T& expected, T desired) {
    return v_.compare_exchange_strong(expected, desired,
                                      std::memory_order_relaxed);
  }
  T ExchangeRelaxed(T desired) {
    return v_.exchange(desired, std::memory_order_relaxed);
  }
  T FetchAddRelaxed(T operand) {
    return v_.fetch_add(operand, std::memory_order_relaxed);
  }
  T FetchSubRelaxed(T operand) {
    return v_.fetch_sub(operand, std::memory_order_relaxed);
  }
  T FetchAndRelaxed(T operand) {
    return v_.fetch_and(operand, std::memory_order_relaxed);
  }
  T FetchOrRelaxed(T operand) {
    return v_.fetch_or(operand, std::memory_order_relaxed);
  }
  T FetchXorRelaxed(T operand) {
    return v_.fetch_xor(operand, std::memory_order_relaxed);
  }

 protected:
  std::atomic<T> v_;
};

// A reasonably general-purpose wrapper around std::atomic for better code
// clarity (see Background above).
//
// Operations use std::memory_order_acq_rel by default (or just acquire or just
// release for read-only and write-only operations), but relaxed operations are
// also available and can be mixed in when appropriate.
//
// Future: add std::memory_order_seqcst variants like StoreSeqCst if/when
// there's a need for them (rare). No distinct type is needed because the
// distinction between acq_rel and seq_cst is more about where it is used in
// combination with other atomics than the atomic itself.
template <typename T>
class Atomic : public RelaxedAtomic<T> {
 public:
  explicit Atomic(T initial = {}) : RelaxedAtomic<T>(initial) {}
  void Store(T desired) {
    RelaxedAtomic<T>::v_.store(desired, std::memory_order_release);
  }
  T Load() const {
    return RelaxedAtomic<T>::v_.load(std::memory_order_acquire);
  }
  bool CasWeak(T& expected, T desired) {
    return RelaxedAtomic<T>::v_.compare_exchange_weak(
        expected, desired, std::memory_order_acq_rel);
  }
  bool CasStrong(T& expected, T desired) {
    return RelaxedAtomic<T>::v_.compare_exchange_strong(
        expected, desired, std::memory_order_acq_rel);
  }
  T Exchange(T desired) {
    return RelaxedAtomic<T>::v_.exchange(desired, std::memory_order_acq_rel);
  }
  T FetchAdd(T operand) {
    return RelaxedAtomic<T>::v_.fetch_add(operand, std::memory_order_acq_rel);
  }
  T FetchSub(T operand) {
    return RelaxedAtomic<T>::v_.fetch_sub(operand, std::memory_order_acq_rel);
  }
  T FetchAnd(T operand) {
    return RelaxedAtomic<T>::v_.fetch_and(operand, std::memory_order_acq_rel);
  }
  T FetchOr(T operand) {
    return RelaxedAtomic<T>::v_.fetch_or(operand, std::memory_order_acq_rel);
  }
  T FetchXor(T operand) {
    return RelaxedAtomic<T>::v_.fetch_xor(operand, std::memory_order_acq_rel);
  }
};

// Atomic shared_ptr helper wrappers.
//
// The C++20 replacement is std::atomic<std::shared_ptr<T>>, but the fbcode
// make build still uses a libstdc++ 11.x toolchain where that specialization
// is unavailable while the older free-function API is already deprecated.
// Centralize the workaround here so callers can keep a single code path until
// the build toolchain can support std::atomic<std::shared_ptr<T>> directly.
template <typename T>
inline std::shared_ptr<T> AtomicSharedPtrLoad(const std::shared_ptr<T>* ptr,
                                              std::memory_order order) {
  ROCKSDB_SUPPRESS_SHARED_PTR_ATOMIC_DEPRECATION_PUSH
  auto result = std::atomic_load_explicit(ptr, order);
  ROCKSDB_SUPPRESS_SHARED_PTR_ATOMIC_DEPRECATION_POP
  return result;
}

template <typename T>
inline void AtomicSharedPtrStore(std::shared_ptr<T>* ptr,
                                 std::shared_ptr<T> desired,
                                 std::memory_order order) {
  ROCKSDB_SUPPRESS_SHARED_PTR_ATOMIC_DEPRECATION_PUSH
  std::atomic_store_explicit(ptr, std::move(desired), order);
  ROCKSDB_SUPPRESS_SHARED_PTR_ATOMIC_DEPRECATION_POP
}

}  // namespace ROCKSDB_NAMESPACE

#undef ROCKSDB_SUPPRESS_SHARED_PTR_ATOMIC_DEPRECATION_PUSH
#undef ROCKSDB_SUPPRESS_SHARED_PTR_ATOMIC_DEPRECATION_POP