// Copyright 2016 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef FBL_REF_COUNTED_INTERNAL_H_
#define FBL_REF_COUNTED_INTERNAL_H_

#include <zircon/assert.h>
#include <zircon/compiler.h>

#include <atomic>

namespace fbl {
namespace internal {

// Adoption validation will help to catch:
// - Double-adoptions
// - AddRef/Release without adopting first
// - Re-wrapping raw pointers to destroyed objects
//
// It also provides some limited defense against
// - Wrapping bad pointers
template <bool EnableAdoptionValidator>
class RefCountedBase {
 protected:
  constexpr RefCountedBase() : ref_count_(kPreAdoptSentinel) {}

  ~RefCountedBase() {
    if constexpr (EnableAdoptionValidator) {
      // Reset the ref-count back to the pre-adopt sentinel value so that we
      // have the best chance of catching a use-after-free situation, even if
      // we have a messed up mix of debug/release translation units being
      // linked together.
      ref_count_.store(kPreAdoptSentinel, std::memory_order_release);
    }
  }

  void AddRef() const {
    const int32_t rc = ref_count_.fetch_add(1, std::memory_order_relaxed);

    // This assertion will fire if either of the following occur.
    //
    // 1) someone calls AddRef() before the object has been properly
    // Adopted.
    //
    // 2) someone calls AddRef() on a ref-counted object that has
    // reached ref_count_ == 0 but has not been destroyed yet. This
    // could happen by manually calling AddRef(), or re-wrapping such a
    // pointer with RefPtr<T>(T*) (which calls AddRef()).
    //
    // Note: leave the ASSERT on in all builds.  The constant
    // EnableAdoptionValidator check above should cause this code path to be
    // pruned in release builds, but leaving this as an always on ASSERT
    // will mean that the tests continue to function even when built as
    // release.
    if constexpr (EnableAdoptionValidator) {
      ZX_ASSERT_MSG(rc >= 1, "count %d(0x%08x) < 1\n", rc, static_cast<uint32_t>(rc));
    }
  }

  // Returns true if the object should self-delete.
  bool Release() const __WARN_UNUSED_RESULT {
    const int32_t rc = ref_count_.fetch_sub(1, std::memory_order_release);

    // This assertion will fire if someone manually calls Release()
    // on a ref-counted object too many times, or if Release is called
    // before an object has been Adopted.
    //
    // Note: leave the ASSERT on in all builds.  The constant
    // EnableAdoptionValidator check above should cause this code path to be
    // pruned in release builds, but leaving this as an always on ASSERT
    // will mean that the tests continue to function even when built as
    // release.
    if constexpr (EnableAdoptionValidator) {
      ZX_ASSERT_MSG(rc >= 1, "count %d(0x%08x) < 1\n", rc, static_cast<uint32_t>(rc));
    }

    if (rc == 1) {
      atomic_thread_fence(std::memory_order_acquire);
      return true;
    }

    return false;
  }

  bool IsLastReference() const __WARN_UNUSED_RESULT {
    return ref_count_.load(std::memory_order_seq_cst) == 1;
  }

  void Adopt() const {
    if constexpr (EnableAdoptionValidator) {
      int32_t expected = kPreAdoptSentinel;
      bool res = ref_count_.compare_exchange_strong(expected, 1, std::memory_order_acq_rel,
                                                    std::memory_order_acquire);
      // Note: leave the ASSERT on in all builds.  The constant
      // EnableAdoptionValidator check above should cause this code path
      // to be pruned in release builds, but leaving this as an always on
      // ASSERT will mean that the tests continue to function even when
      // built as release.
      ZX_ASSERT_MSG(res, "count(0x%08x) != sentinel(0x%08x)\n", static_cast<uint32_t>(expected),
                    static_cast<uint32_t>(kPreAdoptSentinel));
    } else {
      ref_count_.store(1, std::memory_order_release);
    }
  }

  // Current ref count. Only to be used for debugging purposes.
  int ref_count_debug() const { return ref_count_.load(std::memory_order_relaxed); }

  // Note:
  //
  // The PreAdoptSentinel value is chosen specifically to be negative when
  // stored as an int32_t, and as far away from becoming positive (via either
  // addition or subtraction) as possible.  These properties allow us to
  // combine the debug-build adopt sanity checks and the lifecycle sanity
  // checks into a single debug assert.
  //
  // If a user creates an object, but never adopts it, they would need to
  // perform 0x4000000 (about 1 billion) unchecked AddRef or Release
  // operations before making the internal ref_count become positive again.
  // At this point, even a checked AddRef or Release operation would fail to
  // detect the bad state of the system fails to detect the problem.
  //
  static constexpr int32_t kPreAdoptSentinel = static_cast<int32_t>(0xC0000000);
  mutable std::atomic_int32_t ref_count_;
};

}  // namespace internal
}  // namespace fbl

#endif  // FBL_REF_COUNTED_INTERNAL_H_