/* -*- 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/. */

#ifndef mozilla_mscom_Ptr_h
#define mozilla_mscom_Ptr_h

#include "mozilla/Assertions.h"
#include "mozilla/DebugOnly.h"
#include "mozilla/mscom/EnsureMTA.h"
#include "mozilla/SystemGroup.h"
#include "mozilla/UniquePtr.h"
#include "nsError.h"
#include "nsThreadUtils.h"
#include "nsXULAppAPI.h"

/**
 * The glue code in mozilla::mscom often needs to pass around interface pointers
 * belonging to a different apartment from the current one. We must not touch
 * the reference counts of those objects on the wrong apartment. By using these
 * UniquePtr specializations, we may ensure that the reference counts are always
 * handled correctly.
 */

namespace mozilla {
namespace mscom {

namespace detail {

template <typename T>
struct MainThreadRelease
{
  void operator()(T* aPtr)
  {
    if (!aPtr) {
      return;
    }
    if (NS_IsMainThread()) {
      aPtr->Release();
      return;
    }
    DebugOnly<nsresult> rv =
      SystemGroup::Dispatch(TaskCategory::Other,
                            NewNonOwningRunnableMethod("mscom::MainThreadRelease",
                                                       aPtr, &T::Release));
    MOZ_ASSERT(NS_SUCCEEDED(rv));
  }
};

template <typename T>
struct MTADelete
{
  void operator()(T* aPtr)
  {
    if (!aPtr) {
      return;
    }

    EnsureMTA::AsyncOperation([aPtr]() -> void {
      delete aPtr;
    });
  }
};

template <typename T>
struct MTARelease
{
  void operator()(T* aPtr)
  {
    if (!aPtr) {
      return;
    }

    // Static analysis doesn't recognize that, even though aPtr escapes the
    // current scope, we are in effect moving our strong ref into the lambda.
    void* ptr = aPtr;
    EnsureMTA::AsyncOperation([ptr]() -> void {
      reinterpret_cast<T*>(ptr)->Release();
    });
  }
};

template <typename T>
struct MTAReleaseInChildProcess
{
  void operator()(T* aPtr)
  {
    if (!aPtr) {
      return;
    }

    if (XRE_IsParentProcess()) {
      MOZ_ASSERT(NS_IsMainThread());
      aPtr->Release();
      return;
    }

    // Static analysis doesn't recognize that, even though aPtr escapes the
    // current scope, we are in effect moving our strong ref into the lambda.
    void* ptr = aPtr;
    EnsureMTA::AsyncOperation([ptr]() -> void {
      reinterpret_cast<T*>(ptr)->Release();
    });
  }
};

struct InterceptorTargetDeleter
{
  void operator()(IUnknown* aPtr)
  {
    // We intentionally do not touch the refcounts of interceptor targets!
  }
};

struct PreservedStreamDeleter
{
  void operator()(IStream* aPtr)
  {
    if (!aPtr) {
      return;
    }

    // Static analysis doesn't recognize that, even though aPtr escapes the
    // current scope, we are in effect moving our strong ref into the lambda.
    void* ptr = aPtr;
    auto cleanup = [ptr]() -> void {
      DebugOnly<HRESULT> hr =
        ::CoReleaseMarshalData(reinterpret_cast<LPSTREAM>(ptr));
      MOZ_ASSERT(SUCCEEDED(hr));
      reinterpret_cast<LPSTREAM>(ptr)->Release();
    };

    if (XRE_IsParentProcess()) {
      MOZ_ASSERT(NS_IsMainThread());
      cleanup();
      return;
    }

    EnsureMTA::AsyncOperation(cleanup);
  }
};

} // namespace detail

template <typename T>
using STAUniquePtr = mozilla::UniquePtr<T, detail::MainThreadRelease<T>>;

template <typename T>
using MTAUniquePtr = mozilla::UniquePtr<T, detail::MTARelease<T>>;

template <typename T>
using MTADeletePtr = mozilla::UniquePtr<T, detail::MTADelete<T>>;

template <typename T>
using ProxyUniquePtr = mozilla::UniquePtr<T, detail::MTAReleaseInChildProcess<T>>;

template <typename T>
using InterceptorTargetPtr =
  mozilla::UniquePtr<T, detail::InterceptorTargetDeleter>;

using PreservedStreamPtr =
  mozilla::UniquePtr<IStream, detail::PreservedStreamDeleter>;

namespace detail {

// We don't have direct access to UniquePtr's storage, so we use mPtrStorage
// to receive the pointer and then set the target inside the destructor.
template <typename T, typename Deleter>
class UniquePtrGetterAddRefs
{
public:
  explicit UniquePtrGetterAddRefs(UniquePtr<T, Deleter>& aSmartPtr)
    : mTargetSmartPtr(aSmartPtr)
    , mPtrStorage(nullptr)
  {
  }

  ~UniquePtrGetterAddRefs()
  {
    mTargetSmartPtr.reset(mPtrStorage);
  }

  operator void**()
  {
    return reinterpret_cast<void**>(&mPtrStorage);
  }

  operator T**()
  {
    return &mPtrStorage;
  }

  T*& operator*()
  {
    return mPtrStorage;
  }

private:
  UniquePtr<T, Deleter>& mTargetSmartPtr;
  T*                     mPtrStorage;
};

} // namespace detail

template <typename T>
inline STAUniquePtr<T>
ToSTAUniquePtr(RefPtr<T>&& aRefPtr)
{
  return STAUniquePtr<T>(aRefPtr.forget().take());
}

template <typename T>
inline STAUniquePtr<T>
ToSTAUniquePtr(const RefPtr<T>& aRefPtr)
{
  MOZ_ASSERT(NS_IsMainThread());
  return STAUniquePtr<T>(do_AddRef(aRefPtr).take());
}

template <typename T>
inline STAUniquePtr<T>
ToSTAUniquePtr(T* aRawPtr)
{
  MOZ_ASSERT(NS_IsMainThread());
  if (aRawPtr) {
    aRawPtr->AddRef();
  }
  return STAUniquePtr<T>(aRawPtr);
}

template <typename T, typename U>
inline STAUniquePtr<T>
ToSTAUniquePtr(const InterceptorTargetPtr<U>& aTarget)
{
  MOZ_ASSERT(NS_IsMainThread());
  RefPtr<T> newRef(static_cast<T*>(aTarget.get()));
  return ToSTAUniquePtr(std::move(newRef));
}

template <typename T>
inline MTAUniquePtr<T>
ToMTAUniquePtr(RefPtr<T>&& aRefPtr)
{
  return MTAUniquePtr<T>(aRefPtr.forget().take());
}

template <typename T>
inline MTAUniquePtr<T>
ToMTAUniquePtr(const RefPtr<T>& aRefPtr)
{
  MOZ_ASSERT(IsCurrentThreadMTA());
  return MTAUniquePtr<T>(do_AddRef(aRefPtr).take());
}

template <typename T>
inline MTAUniquePtr<T>
ToMTAUniquePtr(T* aRawPtr)
{
  MOZ_ASSERT(IsCurrentThreadMTA());
  if (aRawPtr) {
    aRawPtr->AddRef();
  }
  return MTAUniquePtr<T>(aRawPtr);
}

template <typename T>
inline ProxyUniquePtr<T>
ToProxyUniquePtr(RefPtr<T>&& aRefPtr)
{
  return ProxyUniquePtr<T>(aRefPtr.forget().take());
}

template <typename T>
inline ProxyUniquePtr<T>
ToProxyUniquePtr(const RefPtr<T>& aRefPtr)
{
  MOZ_ASSERT(IsProxy(aRefPtr));
  MOZ_ASSERT((XRE_IsParentProcess() && NS_IsMainThread()) ||
             (XRE_IsContentProcess() && IsCurrentThreadMTA()));

  return ProxyUniquePtr<T>(do_AddRef(aRefPtr).take());
}

template <typename T>
inline ProxyUniquePtr<T>
ToProxyUniquePtr(T* aRawPtr)
{
  MOZ_ASSERT(IsProxy(aRawPtr));
  MOZ_ASSERT((XRE_IsParentProcess() && NS_IsMainThread()) ||
             (XRE_IsContentProcess() && IsCurrentThreadMTA()));

  if (aRawPtr) {
    aRawPtr->AddRef();
  }
  return ProxyUniquePtr<T>(aRawPtr);
}

template <typename T, typename Deleter>
inline InterceptorTargetPtr<T>
ToInterceptorTargetPtr(const UniquePtr<T, Deleter>& aTargetPtr)
{
  return InterceptorTargetPtr<T>(aTargetPtr.get());
}

inline PreservedStreamPtr
ToPreservedStreamPtr(RefPtr<IStream>&& aStream)
{
  return PreservedStreamPtr(aStream.forget().take());
}

inline PreservedStreamPtr
ToPreservedStreamPtr(already_AddRefed<IStream>& aStream)
{
  return PreservedStreamPtr(aStream.take());
}

template <typename T, typename Deleter>
inline detail::UniquePtrGetterAddRefs<T, Deleter>
getter_AddRefs(UniquePtr<T, Deleter>& aSmartPtr)
{
  return detail::UniquePtrGetterAddRefs<T, Deleter>(aSmartPtr);
}

} // namespace mscom
} // namespace mozilla

// This block makes it possible for these smart pointers to be correctly
// applied in NewRunnableMethod and friends
namespace detail {

template<typename T>
struct SmartPointerStorageClass<mozilla::mscom::STAUniquePtr<T>>
{
  typedef StoreCopyPassByRRef<mozilla::mscom::STAUniquePtr<T>> Type;
};

template<typename T>
struct SmartPointerStorageClass<mozilla::mscom::MTAUniquePtr<T>>
{
  typedef StoreCopyPassByRRef<mozilla::mscom::MTAUniquePtr<T>> Type;
};

template<typename T>
struct SmartPointerStorageClass<mozilla::mscom::ProxyUniquePtr<T>>
{
  typedef StoreCopyPassByRRef<mozilla::mscom::ProxyUniquePtr<T>> Type;
};

template<typename T>
struct SmartPointerStorageClass<mozilla::mscom::InterceptorTargetPtr<T>>
{
  typedef StoreCopyPassByRRef<mozilla::mscom::InterceptorTargetPtr<T>> Type;
};

template<>
struct SmartPointerStorageClass<mozilla::mscom::PreservedStreamPtr>
{
  typedef StoreCopyPassByRRef<mozilla::mscom::PreservedStreamPtr> Type;
};

} // namespace detail

#endif // mozilla_mscom_Ptr_h