/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
/* vim: set ts=8 sts=2 et sw=2 tw=80: */
/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

/* A thread-safe weak pointer */

/**
 * Derive from SupportsThreadSafeWeakPtr to allow thread-safe weak pointers to
 * an atomically refcounted derived class. These thread-safe weak pointers may
 * be safely accessed and converted to strong pointers on multiple threads.
 *
 * Note that SupportsThreadSafeWeakPtr necessarily already inherits from
 * AtomicRefCounted, so you should not separately inherit from AtomicRefCounted.
 *
 * ThreadSafeWeakPtr and its implementation is distinct from the normal WeakPtr
 * which is not thread-safe. The interface discipline and implementation details
 * are different enough that these two implementations are separated for now for
 * efficiency reasons. If you don't actually need to use weak pointers on
 * multiple threads, you can just use WeakPtr instead.
 *
 * When deriving from SupportsThreadSafeWeakPtr, you should add
 * MOZ_DECLARE_THREADSAFEWEAKREFERENCE_TYPENAME(ClassName) and
 * MOZ_DECLARE_REFCOUNTED_TYPENAME(ClassName) to the public section of your
 * class, where ClassName is the name of your class.
 *
 * Example usage:
 *
 *   class C : public SupportsThreadSafeWeakPtr<C>
 *   {
 *   public:
 *     MOZ_DECLARE_THREADSAFEWEAKREFERENCE_TYPENAME(C)
 *     MOZ_DECLARE_REFCOUNTED_TYPENAME(C)
 *     void doStuff();
 *   };
 *
 *   ThreadSafeWeakPtr<C> weak;
 *   {
 *     RefPtr<C> strong = new C;
 *     if (strong) {
 *       strong->doStuff();
 *     }
 *     // Make a new weak reference to the object from the strong reference.
 *     weak = strong;
 *   }
 *   MOZ_ASSERT(!bool(weak), "Weak pointers are cleared after all "
 *                           "strong references are released.");
 *
 *   // Convert the weak reference to a strong reference for usage.
 *   RefPtr<C> other(weak);
 *   if (other) {
 *     other->doStuff();
 *   }
 */

#ifndef mozilla_ThreadSafeWeakPtr_h
#define mozilla_ThreadSafeWeakPtr_h

#include "mozilla/Assertions.h"
#include "mozilla/Atomics.h"
#include "mozilla/RefCounted.h"
#include "mozilla/RefPtr.h"
#include "mozilla/TypeTraits.h"
#include "mozilla/Unused.h"

#include <limits>

namespace mozilla {

template <typename T>
class ThreadSafeWeakPtr;
template <typename T>
class SupportsThreadSafeWeakPtr;

#ifdef MOZ_REFCOUNTED_LEAK_CHECKING
#  define MOZ_DECLARE_THREADSAFEWEAKREFERENCE_TYPENAME(T)  \
    static const char* threadSafeWeakReferenceTypeName() { \
      return "ThreadSafeWeakReference<" #T ">";            \
    }
#else
#  define MOZ_DECLARE_THREADSAFEWEAKREFERENCE_TYPENAME(T)
#endif

namespace detail {

// A multiple reader, single writer spin-lock.
// This lock maintains an atomic counter which is incremented every time the
// lock is acquired reading. So long as the counter remains positive, it may be
// incremented for reading multiple times. When acquiring the lock for writing,
// we must ensure the counter is 0 (no readers), and if so, set it to a negative
// value to indicate that no new readers may take the lock.
class ReadWriteSpinLock {
  // Only need a type large enough to represent the number of simultaneously
  // accessing threads.
  typedef int32_t CounterType;

 public:
  // Try to increment the counter for reading, so long as it is positive.
  void readLock() {
    for (;;) {
      CounterType oldCounter =
          mCounter & std::numeric_limits<CounterType>::max();
      CounterType newCounter = oldCounter + 1;
      if (mCounter.compareExchange(oldCounter, newCounter)) {
        break;
      }
    }
  }

  // Decrement the counter to remove a read lock.
  void readUnlock() { mCounter--; }

  // Spins until a write lock is acquired. This can only occur if there are no
  // readers or writers. Once it is acquired, the counter is set to a negative
  // value.
  void writeLock() {
    while (true) {
      if (mCounter.compareExchange(0,
                                   std::numeric_limits<CounterType>::min())) {
        return;
      }
    }
  }

  // Reset the counter to 0.
  void writeUnlock() { mCounter = 0; }

 private:
  Atomic<CounterType> mCounter;
};

// A shared weak reference that is used to track a SupportsThreadSafeWeakPtr
// object. It guards access to that object via a read-write spinlock.
template <typename T>
class ThreadSafeWeakReference
    : public external::AtomicRefCounted<ThreadSafeWeakReference<T>> {
 public:
  typedef T ElementType;

  explicit ThreadSafeWeakReference(T* aPtr) { mPtr = aPtr; }

#ifdef MOZ_REFCOUNTED_LEAK_CHECKING
  const char* typeName() const {
    // The first time this is called mPtr is null, so don't
    // invoke any methods on mPtr.
    return T::threadSafeWeakReferenceTypeName();
  }
  size_t typeSize() const { return sizeof(*this); }
#endif

 private:
  friend class mozilla::SupportsThreadSafeWeakPtr<T>;
  template <typename U>
  friend class mozilla::ThreadSafeWeakPtr;

  // Does an unsafe read of the raw weak pointer.
  T* get() const { return mPtr; }

  // Creates a new RefPtr to the tracked object.
  // We need to acquire the read lock while we do this, as we need to atomically
  // both read the pointer and then increment the refcount on it within the
  // scope of the lock. This guards against the object being destroyed while in
  // the middle of creating the new RefPtr.
  already_AddRefed<T> getRefPtr() {
    mLock.readLock();
    RefPtr<T> result(get());
    mLock.readUnlock();
    return result.forget();
  }

  // Try to detach the weak reference from the tracked object.
  // We need to acquire the write lock while we do this, to ensure that no
  // RefPtr is created to this while detaching. Once acquired, it is safe
  // to check the refcount and verify that this is the last reference to
  // the tracked object, so the weak reference can be safely detached.
  void tryDetach(const SupportsThreadSafeWeakPtr<T>* aOwner) {
    mLock.writeLock();
    if (aOwner->hasOneRef()) {
      mPtr = nullptr;
    }
    mLock.writeUnlock();
  }

  ReadWriteSpinLock mLock;
  Atomic<T*> mPtr;
};

}  // namespace detail

template <typename T>
class SupportsThreadSafeWeakPtr : public external::AtomicRefCounted<T> {
 protected:
  typedef external::AtomicRefCounted<T> AtomicRefCounted;
  typedef detail::ThreadSafeWeakReference<T> ThreadSafeWeakReference;

 public:
  ~SupportsThreadSafeWeakPtr() {
    // Clean up the shared weak reference if one exists.
    if (ThreadSafeWeakReference* ptr = mRef) {
      ptr->Release();
    }
  }

  void Release() const {
    // If there is only one remaining reference to the object when trying to
    // release, then attempt to detach it from its weak reference. New
    // references could possibly be created to the object while this happens, so
    // take care to do this atomically inside tryDetach.
    if (AtomicRefCounted::hasOneRef()) {
      if (ThreadSafeWeakReference* ptr = mRef) {
        ptr->tryDetach(this);
      }
    }

    // Once possibly detached, it is now safe to continue to decrement the
    // refcount.
    AtomicRefCounted::Release();
  }

 private:
  template <typename U>
  friend class ThreadSafeWeakPtr;

  // Creates a shared weak reference for the object if one does not exist. Note
  // that the object may be of an actual derived type U, but the weak reference
  // is created for the supplied type T of SupportsThreadSafeWeakPtr<T>.
  already_AddRefed<ThreadSafeWeakReference> getThreadSafeWeakReference() {
    static_assert(std::is_base_of<SupportsThreadSafeWeakPtr<T>, T>::value,
                  "T must derive from SupportsThreadSafeWeakPtr<T>");

    if (!mRef) {
      RefPtr<ThreadSafeWeakReference> ptr(
          new ThreadSafeWeakReference(static_cast<T*>(this)));
      // Only set the new weak reference if one does not exist (== nullptr).
      // If there is already a weak reference, just let this superflous weak
      // reference get destroyed when it goes out of scope.
      if (mRef.compareExchange(nullptr, ptr)) {
        // If successful, forget the refcount so that the weak reference stays
        // alive.
        Unused << ptr.forget();
      }
    }

    // Create a new RefPtr to weak reference.
    RefPtr<ThreadSafeWeakReference> ptr(mRef);
    return ptr.forget();
  }

  Atomic<ThreadSafeWeakReference*> mRef;
};

// A thread-safe variant of a weak pointer
template <typename T>
class ThreadSafeWeakPtr {
  // Be careful to use the weak reference type T in the
  // SupportsThreadSafeWeakPtr<T> definition.
  typedef typename T::ThreadSafeWeakReference ThreadSafeWeakReference;

 public:
  ThreadSafeWeakPtr() = default;

  ThreadSafeWeakPtr& operator=(const ThreadSafeWeakPtr& aOther) = default;
  ThreadSafeWeakPtr(const ThreadSafeWeakPtr& aOther) = default;

  ThreadSafeWeakPtr& operator=(ThreadSafeWeakPtr&& aOther) = default;
  ThreadSafeWeakPtr(ThreadSafeWeakPtr&& aOther) = default;

  ThreadSafeWeakPtr& operator=(const RefPtr<T>& aOther) {
    if (aOther) {
      // Get the underlying shared weak reference to the object, creating one if
      // necessary.
      mRef = aOther->getThreadSafeWeakReference();
    } else {
      mRef = nullptr;
    }
    return *this;
  }

  explicit ThreadSafeWeakPtr(const RefPtr<T>& aOther) { *this = aOther; }

  ThreadSafeWeakPtr& operator=(decltype(nullptr)) {
    mRef = nullptr;
    return *this;
  }

  explicit ThreadSafeWeakPtr(decltype(nullptr)) {}

  explicit operator bool() const { return !!get(); }

  bool operator==(const ThreadSafeWeakPtr& aOther) const {
    return get() == aOther.get();
  }

  bool operator==(const RefPtr<T>& aOther) const {
    return get() == aOther.get();
  }

  bool operator==(const T* aOther) const { return get() == aOther; }

  template <typename U>
  bool operator!=(const U& aOther) const {
    return !(*this == aOther);
  }

  // Convert the weak pointer to a strong RefPtr.
  explicit operator RefPtr<T>() const { return getRefPtr(); }

 private:
  // Gets a new strong reference of the proper type T to the tracked object.
  already_AddRefed<T> getRefPtr() const {
    static_assert(std::is_base_of<typename ThreadSafeWeakReference::ElementType,
                                  T>::value,
                  "T must derive from ThreadSafeWeakReference::ElementType");
    return mRef ? mRef->getRefPtr().template downcast<T>() : nullptr;
  }

  // Get a pointer to the tracked object, downcasting to the proper type T.
  // Note that this operation is unsafe as it may cause races if downwind
  // code depends on the value not to change after reading.
  T* get() const {
    static_assert(std::is_base_of<typename ThreadSafeWeakReference::ElementType,
                                  T>::value,
                  "T must derive from ThreadSafeWeakReference::ElementType");
    return mRef ? static_cast<T*>(mRef->get()) : nullptr;
  }

  // A shared weak reference to an object. Note that this may be null so as to
  // save memory (at the slight cost of an extra null check) if no object is
  // being tracked.
  RefPtr<ThreadSafeWeakReference> mRef;
};

}  // namespace mozilla

template <typename T>
inline already_AddRefed<T> do_AddRef(
    const mozilla::ThreadSafeWeakPtr<T>& aObj) {
  RefPtr<T> ref(aObj);
  return ref.forget();
}

#endif /* mozilla_ThreadSafeWeakPtr_h */