///////////////////////////////////////////////////////////////////////////////
//                                                                           //
// microcom.h                                                                //
// Copyright (C) Microsoft Corporation. All rights reserved.                 //
// This file is distributed under the University of Illinois Open Source     //
// License. See LICENSE.TXT for details.                                     //
//                                                                           //
// Provides support for basic COM-like constructs.                           //
//                                                                           //
///////////////////////////////////////////////////////////////////////////////

#ifndef __DXC_MICROCOM__
#define __DXC_MICROCOM__

#include "dxc/Support/WinIncludes.h"
#include "llvm/Support/Atomic.h"
#include <atomic>

template <typename TIface> class CComInterfaceArray {
private:
  TIface **m_pData;
  unsigned m_length;

public:
  CComInterfaceArray() : m_pData(nullptr), m_length(0) {}
  ~CComInterfaceArray() { clear(); }
  bool empty() const { return m_length == 0; }
  unsigned size() const { return m_length; }
  TIface ***data_ref() { return &m_pData; }
  unsigned *size_ref() { return &m_length; }
  TIface **begin() { return m_pData; }
  TIface **end() { return m_pData + m_length; }
  void clear() {
    if (m_length) {
      for (unsigned i = 0; i < m_length; ++i) {
        if (m_pData[i] != nullptr) {
          m_pData[i]->Release();
          m_pData[i] = nullptr;
        }
      }
      m_length = 0;
    }
    if (m_pData) {
      CoTaskMemFree(m_pData);
      m_pData = nullptr;
    }
  }
  HRESULT alloc(unsigned count) {
    clear();
    m_pData = (TIface **)CoTaskMemAlloc(sizeof(TIface *) * count);
    if (m_pData == nullptr)
      return E_OUTOFMEMORY;
    m_length = count;
    ZeroMemory(m_pData, sizeof(TIface *) * count);
    return S_OK;
  }
  TIface **get_address_of(unsigned index) { return &(m_pData[index]); }
  TIface **release() {
    TIface **result = m_pData;
    m_pData = nullptr;
    m_length = 0;
    return result;
  }
  void release(TIface ***pValues, unsigned *length) {
    *pValues = m_pData;
    m_pData = nullptr;
    *length = m_length;
    m_length = 0;
  }
};

template <typename T> void DxcCallDestructor(T *obj) { obj->T::~T(); }

#define DXC_MICROCOM_REF_FIELD(m_dwRef)                                        \
  volatile std::atomic<llvm::sys::cas_flag> m_dwRef = {0};
#define DXC_MICROCOM_ADDREF_IMPL(m_dwRef)                                      \
  ULONG STDMETHODCALLTYPE AddRef() noexcept override {                         \
    return (ULONG)++m_dwRef;                                                   \
  }
#define DXC_MICROCOM_ADDREF_RELEASE_IMPL(m_dwRef)                              \
  DXC_MICROCOM_ADDREF_IMPL(m_dwRef)                                            \
  ULONG STDMETHODCALLTYPE Release() override {                                 \
    ULONG result = (ULONG)--m_dwRef;                                           \
    if (result == 0) {                                                         \
      DxcCallDestructor(this);                                                 \
      operator delete(this);                                                   \
    }                                                                          \
    return result;                                                             \
  }

template <typename T, typename... Args>
inline T *CreateOnMalloc(IMalloc *pMalloc, Args &&...args) {
  void *P = pMalloc->Alloc(sizeof(T));
  try {
    if (P)
      new (P) T(pMalloc, std::forward<Args>(args)...);
  } catch (...) {
    pMalloc->Free(P);
    throw;
  }
  return (T *)P;
}

// The "TM" version keep an IMalloc field that, if not null, indicate
// ownership of 'this' and of any allocations used during release.
#define DXC_MICROCOM_TM_REF_FIELDS()                                           \
  volatile std::atomic<llvm::sys::cas_flag> m_dwRef = {0};                     \
  CComPtr<IMalloc> m_pMalloc;

#define DXC_MICROCOM_TM_ADDREF_RELEASE_IMPL()                                  \
  DXC_MICROCOM_ADDREF_IMPL(m_dwRef)                                            \
  ULONG STDMETHODCALLTYPE Release() noexcept override {                        \
    ULONG result = (ULONG)--m_dwRef;                                           \
    if (result == 0) {                                                         \
      CComPtr<IMalloc> pTmp(m_pMalloc);                                        \
      DxcThreadMalloc M(pTmp);                                                 \
      DxcCallDestructor(this);                                                 \
      pTmp->Free(this);                                                        \
    }                                                                          \
    return result;                                                             \
  }

#define DXC_MICROCOM_TM_CTOR(T)                                                \
  DXC_MICROCOM_TM_CTOR_ONLY(T)                                                 \
  DXC_MICROCOM_TM_ALLOC(T)
#define DXC_MICROCOM_TM_CTOR_ONLY(T)                                           \
  T(IMalloc *pMalloc) : m_dwRef(0), m_pMalloc(pMalloc) {}
#define DXC_MICROCOM_TM_ALLOC(T)                                               \
  template <typename... Args>                                                  \
  static T *Alloc(IMalloc *pMalloc, Args &&...args) {                          \
    void *P = pMalloc->Alloc(sizeof(T));                                       \
    try {                                                                      \
      if (P)                                                                   \
        new (P) T(pMalloc, std::forward<Args>(args)...);                       \
    } catch (...) {                                                            \
      pMalloc->Free(P);                                                        \
      throw;                                                                   \
    }                                                                          \
    return (T *)P;                                                             \
  }

/// <summary>
/// Provides a QueryInterface implementation for a class that supports
/// any number of interfaces in addition to IUnknown.
/// </summary>
/// <remarks>
/// This implementation will also report the instance as not supporting
/// marshaling. This will help catch marshaling problems early or avoid
/// them altogether.
/// </remarks>
template <typename TObject>
HRESULT DoBasicQueryInterface_recurse(TObject *self, REFIID iid,
                                      void **ppvObject) {
  return E_NOINTERFACE;
}
template <typename TObject, typename TInterface, typename... Ts>
HRESULT DoBasicQueryInterface_recurse(TObject *self, REFIID iid,
                                      void **ppvObject) {
  if (ppvObject == nullptr)
    return E_POINTER;
  if (IsEqualIID(iid, __uuidof(TInterface))) {
    *(TInterface **)ppvObject = self;
    self->AddRef();
    return S_OK;
  }
  return DoBasicQueryInterface_recurse<TObject, Ts...>(self, iid, ppvObject);
}
template <typename... Ts, typename TObject>
HRESULT DoBasicQueryInterface(TObject *self, REFIID iid, void **ppvObject) {
  if (ppvObject == nullptr)
    return E_POINTER;

  // Support INoMarshal to void GIT shenanigans.
  if (IsEqualIID(iid, __uuidof(IUnknown)) ||
      IsEqualIID(iid, __uuidof(INoMarshal))) {
    *ppvObject = reinterpret_cast<IUnknown *>(self);
    reinterpret_cast<IUnknown *>(self)->AddRef();
    return S_OK;
  }

  return DoBasicQueryInterface_recurse<TObject, Ts...>(self, iid, ppvObject);
}

template <typename T> HRESULT AssignToOut(T value, T *pResult) {
  if (pResult == nullptr)
    return E_POINTER;
  *pResult = value;
  return S_OK;
}
template <typename T> HRESULT AssignToOut(std::nullptr_t value, T *pResult) {
  if (pResult == nullptr)
    return E_POINTER;
  *pResult = value;
  return S_OK;
}
template <typename T> HRESULT ZeroMemoryToOut(T *pResult) {
  if (pResult == nullptr)
    return E_POINTER;
  ZeroMemory(pResult, sizeof(*pResult));
  return S_OK;
}

template <typename T> void AssignToOutOpt(T value, T *pResult) {
  if (pResult != nullptr)
    *pResult = value;
}
template <typename T> void AssignToOutOpt(std::nullptr_t value, T *pResult) {
  if (pResult != nullptr)
    *pResult = value;
}

#endif
