Add support for custom allocators

include/dxc/Support/DxcLangExtensionsHelper.h

@@ -221,21 +221,27 @@ class DxcLangExtensionsHelper : public DxcLangExtensionsHelperApply {
 // Note that QueryInterface still needs to return the vtable.
 #define DXC_LANGEXTENSIONS_HELPER_IMPL(_helper_field_) \
   __override HRESULT STDMETHODCALLTYPE RegisterIntrinsicTable(_In_ IDxcIntrinsicTable *pTable) { \
+    DxcThreadMalloc TM(m_pMalloc); \
     return (_helper_field_).RegisterIntrinsicTable(pTable); \
   } \
   __override HRESULT STDMETHODCALLTYPE RegisterSemanticDefine(LPCWSTR name) { \
+    DxcThreadMalloc TM(m_pMalloc); \
     return (_helper_field_).RegisterSemanticDefine(name); \
   } \
   __override HRESULT STDMETHODCALLTYPE RegisterSemanticDefineExclusion(LPCWSTR name) { \
+    DxcThreadMalloc TM(m_pMalloc); \
     return (_helper_field_).RegisterSemanticDefineExclusion(name); \
   } \
   __override HRESULT STDMETHODCALLTYPE RegisterDefine(LPCWSTR name) { \
+    DxcThreadMalloc TM(m_pMalloc); \
     return (_helper_field_).RegisterDefine(name); \
   } \
   __override HRESULT STDMETHODCALLTYPE SetSemanticDefineValidator(_In_ IDxcSemanticDefineValidator* pValidator) { \
+    DxcThreadMalloc TM(m_pMalloc); \
     return (_helper_field_).SetSemanticDefineValidator(pValidator); \
   } \
   __override HRESULT STDMETHODCALLTYPE SetSemanticDefineMetaDataName(LPCSTR name) { \
+    DxcThreadMalloc TM(m_pMalloc); \
     return (_helper_field_).SetSemanticDefineMetaDataName(name); \
   } \
 

include/dxc/Support/FileIOHelper.h

@@ -17,6 +17,99 @@ struct IDxcBlobEncoding;
 
 namespace hlsl {
 
+IMalloc *GetGlobalHeapMalloc() throw();
+
+class CDxcThreadMallocAllocator {
+public:
+  _Ret_maybenull_ _Post_writable_byte_size_(nBytes) _ATL_DECLSPEC_ALLOCATOR
+  static void *Reallocate(_In_ void *p, _In_ size_t nBytes) throw() {
+    return DxcGetThreadMallocNoRef()->Realloc(p, nBytes);
+  }
+
+  _Ret_maybenull_ _Post_writable_byte_size_(nBytes) _ATL_DECLSPEC_ALLOCATOR
+  static void *Allocate(_In_ size_t nBytes) throw() {
+    return DxcGetThreadMallocNoRef()->Alloc(nBytes);
+  }
+
+  static void Free(_In_ void *p) throw() {
+    return DxcGetThreadMallocNoRef()->Free(p);
+  }
+};
+
+// Like CComHeapPtr, but with CDxcThreadMallocAllocator.
+template <typename T>
+class CDxcTMHeapPtr :
+  public CHeapPtr<T, CDxcThreadMallocAllocator>
+{
+public:
+  CDxcTMHeapPtr() throw()
+  {
+  }
+
+  explicit CDxcTMHeapPtr(_In_ T* pData) throw() :
+    CDxcTMHeapPtr<T, CDxcThreadMallocAllocator>(pData)
+  {
+  }
+};
+
+// Like CComHeapPtr, but with a stateful allocator.
+template <typename T>
+class CDxcMallocHeapPtr
+{
+private:
+  CComPtr<IMalloc> m_pMalloc;
+public:
+  T *m_pData;
+
+  CDxcMallocHeapPtr(IMalloc *pMalloc) throw()
+      : m_pMalloc(pMalloc), m_pData(nullptr) {}
+
+  ~CDxcMallocHeapPtr() {
+    if (m_pData)
+      m_pMalloc->Free(m_pData);
+  }
+
+  operator T *() const throw() { return m_pData; }
+
+  bool Allocate(_In_ SIZE_T ElementCount) throw() {
+    ATLASSERT(m_pData == NULL);
+    SIZE_T nBytes = ElementCount * sizeof(T);
+    m_pData = static_cast<T *>(m_pMalloc->Alloc(nBytes));
+    if (m_pData == NULL)
+      return false;
+    return true;
+  }
+
+  void AllocateBytes(_In_ SIZE_T ByteCount) throw() {
+    if (m_pData)
+      m_pMalloc->Free(m_pData);
+    m_pData = static_cast<T *>(m_pMalloc->Alloc(ByteCount));
+  }
+
+  // Attach to an existing pointer (takes ownership)
+  void Attach(_In_ T *pData) throw() {
+    m_pMalloc->Free(m_pData);
+    m_pData = pData;
+  }
+
+  // Detach the pointer (releases ownership)
+  T *Detach() throw() {
+    T *pTemp = m_pData;
+    m_pData = NULL;
+    return pTemp;
+  }
+
+  // Free the memory pointed to, and set the pointer to NULL
+  void Free() throw() {
+    m_pMalloc->Free(m_pData);
+    m_pData = NULL;
+  }
+};
+
+void ReadBinaryFile(_In_opt_ IMalloc *pMalloc,
+                    _In_z_ LPCWSTR pFileName,
+                    _Outptr_result_bytebuffer_(*pDataSize) void **ppData,
+                    _Out_ DWORD *pDataSize);
 void ReadBinaryFile(_In_z_ LPCWSTR pFileName,
                     _Outptr_result_bytebuffer_(*pDataSize) void **ppData,
                     _Out_ DWORD *pDataSize);
@@ -30,8 +123,13 @@ void WriteBinaryFile(_In_z_ LPCWSTR pFileName,
 UINT32 DxcCodePageFromBytes(_In_count_(byteLen) const char *bytes,
                             size_t byteLen) throw();
 
+HRESULT
+DxcCreateBlobFromFile(_In_opt_ IMalloc *pMalloc, LPCWSTR pFileName,
+                      _In_opt_ UINT32 *pCodePage,
+                      _COM_Outptr_ IDxcBlobEncoding **pBlobEncoding) throw();
+
 HRESULT DxcCreateBlobFromFile(LPCWSTR pFileName, _In_opt_ UINT32 *pCodePage,
-                              _COM_Outptr_ IDxcBlobEncoding **pBlobEncoding) throw();
+                              _COM_Outptr_ IDxcBlobEncoding **ppBlobEncoding) throw();
 
 // Given a blob, creates a subrange view.
 HRESULT DxcCreateBlobFromBlob(_In_ IDxcBlob *pBlob, UINT32 offset,
@@ -49,7 +147,11 @@ DxcCreateBlobOnHeapCopy(_In_bytecount_(size) LPCVOID pData, UINT32 size,
 // Given a blob, creates a new instance with a specific code page set.
 HRESULT
 DxcCreateBlobWithEncodingSet(_In_ IDxcBlob *pBlob, UINT32 codePage,
-                             _COM_Outptr_ IDxcBlobEncoding **pBlobEncoding) throw();
+                             _COM_Outptr_ IDxcBlobEncoding **ppBlobEncoding) throw();
+HRESULT
+DxcCreateBlobWithEncodingSet(
+    _In_ IMalloc *pMalloc, _In_ IDxcBlob *pBlob, UINT32 codePage,
+    _COM_Outptr_ IDxcBlobEncoding **ppBlobEncoding) throw();
 
 HRESULT DxcCreateBlobWithEncodingFromPinned(
     _In_bytecount_(size) LPCVOID pText, UINT32 size, UINT32 codePage,
@@ -65,6 +167,7 @@ DxcCreateBlobWithEncodingOnHeap(_In_bytecount_(size) LPCVOID pText, UINT32 size,
                                 UINT32 codePage,
                                 _COM_Outptr_ IDxcBlobEncoding **pBlobEncoding) throw();
 
+// Should rename this 'OnHeap' to be 'OnMalloc', change callers to pass arg. Using TLS.
 HRESULT
 DxcCreateBlobWithEncodingOnHeapCopy(
     _In_bytecount_(size) LPCVOID pText, UINT32 size, UINT32 codePage,
@@ -75,15 +178,21 @@ DxcCreateBlobWithEncodingOnMalloc(
   _In_bytecount_(size) LPCVOID pText, IMalloc *pIMalloc, UINT32 size, UINT32 codePage,
   _COM_Outptr_ IDxcBlobEncoding **pBlobEncoding) throw();
 
+HRESULT
+DxcCreateBlobWithEncodingOnMallocCopy(
+  _In_ IMalloc *pIMalloc, _In_bytecount_(size) LPCVOID pText, UINT32 size, UINT32 codePage,
+  _COM_Outptr_ IDxcBlobEncoding **pBlobEncoding) throw();
+
 HRESULT DxcGetBlobAsUtf8(_In_ IDxcBlob *pBlob,
                          _COM_Outptr_ IDxcBlobEncoding **pBlobEncoding) throw();
 HRESULT
 DxcGetBlobAsUtf8NullTerm(
     _In_ IDxcBlob *pBlob,
     _COM_Outptr_ IDxcBlobEncoding **ppBlobEncoding) throw();
 
-HRESULT DxcGetBlobAsUtf16(_In_ IDxcBlob *pBlob,
-                          _COM_Outptr_ IDxcBlobEncoding **pBlobEncoding) throw();
+HRESULT
+DxcGetBlobAsUtf16(_In_ IDxcBlob *pBlob, _In_ IMalloc *pMalloc,
+                  _COM_Outptr_ IDxcBlobEncoding **pBlobEncoding) throw();
 
 bool IsBlobNullOrEmpty(_In_opt_ IDxcBlob *pBlob) throw();
 

include/dxc/Support/Global.h

@@ -5,7 +5,7 @@
 // This file is distributed under the University of Illinois Open Source     //
 // License. See LICENSE.TXT for details.                                     //
 //                                                                           //
-// Provides important declarations global to all DX Compiler code.          //
+// Provides important declarations global to all DX Compiler code.           //
 //                                                                           //
 ///////////////////////////////////////////////////////////////////////////////
 
@@ -25,6 +25,53 @@ typedef _Return_type_success_(return >= 0) long HRESULT;
 #include <stdarg.h>
 #include "dxc/Support/exception.h"
 
+///////////////////////////////////////////////////////////////////////////////
+// Memory allocation support.
+//
+// This mechanism ties into the C++ new and delete operators.
+//
+// Other allocators may be used in specific situations, eg sub-allocators or
+// the COM allocator for interop. This is the preferred allocator in general,
+// however, as it eventually allows the library user to specify their own.
+//
+
+struct IMalloc;
+
+// Used by DllMain to set up and tear down per-thread tracking.
+HRESULT DxcInitThreadMalloc() throw();
+void DxcCleanupThreadMalloc() throw();
+
+// Used by APIs entry points to set up per-thread/invocation allocator.
+// Setting the IMalloc on the thread increases the reference count,
+// clearing it decreases it.
+void DxcClearThreadMalloc() throw();
+void DxcSetThreadMalloc(IMalloc *pMalloc) throw();
+void DxcSetThreadMallocOrDefault(IMalloc *pMalloc) throw();
+
+// Swapping does not AddRef or Release new or prior. The pattern is to keep both alive,
+// either in TLS, or on the stack to restore later. The returned value is the effective
+// IMalloc also available in TLS.
+IMalloc *DxcSwapThreadMalloc(IMalloc *pMalloc, IMalloc **ppPrior) throw();
+IMalloc *DxcSwapThreadMallocOrDefault(IMalloc *pMalloc, IMalloc **ppPrior) throw();
+
+// Used to retrieve the current invocation's allocator or perform an alloc/free/realloc.
+IMalloc *DxcGetThreadMallocNoRef() throw();
+_Ret_maybenull_ _Post_writable_byte_size_(nBytes) void *DxcThreadAlloc(size_t nBytes) throw();
+void DxcThreadFree(void *) throw();
+
+struct DxcThreadMalloc {
+  DxcThreadMalloc(IMalloc *pMallocOrNull) throw() {
+    p = DxcSwapThreadMallocOrDefault(pMallocOrNull, &pPrior);
+  }
+  ~DxcThreadMalloc() {
+    DxcSwapThreadMalloc(pPrior, nullptr);
+  }
+  IMalloc *p;
+  IMalloc *pPrior;
+};
+
+///////////////////////////////////////////////////////////////////////////////
+// Error handling support.
 namespace std { class error_code; }
 void CheckLLVMErrorCode(const std::error_code &ec);
 

include/dxc/Support/HLSLOptions.h

@@ -52,6 +52,8 @@ enum ID {
   };
 
 const llvm::opt::OptTable *getHlslOptTable();
+std::error_code initHlslOptTable();
+void cleanupHlslOptTable();
 
 ///////////////////////////////////////////////////////////////////////////////
 // Helper classes to deal with options.

include/dxc/Support/WinIncludes.h

@@ -32,6 +32,7 @@
 #include <atlbase.h> // atlbase.h needs to come before strsafe.h
 #include <strsafe.h>
 #include <intsafe.h>
+#include <ObjIdl.h>
 
 /// Swap two ComPtr classes.
 template <class T> void swap(CComHeapPtr<T> &a, CComHeapPtr<T> &b) {

include/dxc/Support/dxcapi.impl.h

@@ -34,15 +34,18 @@ class raw_stream_ostream : public llvm::raw_ostream {
 
 class DxcOperationResult : public IDxcOperationResult {
 private:
-  DXC_MICROCOM_REF_FIELD(m_dwRef)
+  DXC_MICROCOM_TM_REF_FIELDS()
 
-  DxcOperationResult(_In_opt_ IDxcBlob *pResultBlob,
-    _In_opt_ IDxcBlobEncoding *pErrorBlob, HRESULT status)
-    : m_dwRef(0), m_status(status), m_result(pResultBlob),
-    m_errors(pErrorBlob) {}
+  void Init(_In_opt_ IDxcBlob *pResultBlob,
+            _In_opt_ IDxcBlobEncoding *pErrorBlob, HRESULT status) {
+    m_status = status;
+    m_result = pResultBlob;
+    m_errors = pErrorBlob;
+  }
 
 public:
-  DXC_MICROCOM_ADDREF_RELEASE_IMPL(m_dwRef)
+  DXC_MICROCOM_TM_ADDREF_RELEASE_IMPL()
+  DXC_MICROCOM_TM_CTOR(DxcOperationResult)
 
   HRESULT m_status;
   CComPtr<IDxcBlob> m_result;
@@ -57,8 +60,9 @@ class DxcOperationResult : public IDxcOperationResult {
                                              HRESULT status,
                                              _COM_Outptr_ IDxcOperationResult **ppResult) {
     *ppResult = nullptr;
-    CComPtr<DxcOperationResult> result = new (std::nothrow) DxcOperationResult(pResultBlob, pErrorBlob, status);
-    if (result.p == nullptr) return E_OUTOFMEMORY;
+    CComPtr<DxcOperationResult> result = DxcOperationResult::Alloc(DxcGetThreadMallocNoRef());
+    IFROOM(result.p);
+    result->Init(pResultBlob, pErrorBlob, status);
     *ppResult = result.Detach();
     return S_OK;
   }

include/dxc/Support/dxcapi.use.h

@@ -21,6 +21,7 @@ class DxcDllSupport {
 protected:
   HMODULE m_dll;
   DxcCreateInstanceProc m_createFn;
+  DxcCreateInstance2Proc m_createFn2;
 
   HRESULT InitializeInternal(LPCWSTR dllName, LPCSTR fnName) {
     if (m_dll != nullptr) return S_OK;
@@ -36,16 +37,28 @@ class DxcDllSupport {
       return hr;
     }
 
+    // Only basic functions used to avoid requiring additional headers.
+    m_createFn2 = nullptr;
+    char fnName2[128];
+    size_t s = strlen(fnName);
+    if (s < sizeof(fnName2) - 2) {
+      memcpy(fnName2, fnName, s);
+      fnName2[s] = '2';
+      fnName2[s + 1] = '\0';
+      m_createFn2 = (DxcCreateInstance2Proc)GetProcAddress(m_dll, fnName2);
+    }
+
     return S_OK;
   }
 
 public:
-  DxcDllSupport() : m_dll(nullptr), m_createFn(nullptr) {
+  DxcDllSupport() : m_dll(nullptr), m_createFn(nullptr), m_createFn2(nullptr) {
   }
 
   DxcDllSupport(DxcDllSupport&& other) {
     m_dll = other.m_dll; other.m_dll = nullptr;
-    m_createFn = other.m_createFn; other.m_dll = nullptr;
+    m_createFn = other.m_createFn; other.m_createFn = nullptr;
+    m_createFn2 = other.m_createFn2; other.m_createFn2 = nullptr;
   }
 
   ~DxcDllSupport() {
@@ -72,13 +85,31 @@ class DxcDllSupport {
     return hr;
   }
 
+  template <typename TInterface>
+  HRESULT CreateInstance2(IMalloc *pMalloc, REFCLSID clsid, _Outptr_ TInterface** pResult) {
+    return CreateInstance2(pMalloc, clsid, __uuidof(TInterface), (IUnknown**)pResult);
+  }
+
+  HRESULT CreateInstance2(IMalloc *pMalloc, REFCLSID clsid, REFIID riid, _Outptr_ IUnknown **pResult) {
+    if (pResult == nullptr) return E_POINTER;
+    if (m_dll == nullptr) return E_FAIL;
+    if (m_createFn2 == nullptr) return E_FAIL;
+    HRESULT hr = m_createFn2(pMalloc, clsid, riid, (LPVOID*)pResult);
+    return hr;
+  }
+
+  bool HasCreateWithMalloc() const {
+    return m_createFn2 != nullptr;
+  }
+
   bool IsEnabled() const {
     return m_dll != nullptr;
   }
 
   void Cleanup() {
     if (m_dll != nullptr) {
       m_createFn = nullptr;
+      m_createFn2 = nullptr;
       FreeLibrary(m_dll);
       m_dll = nullptr;
     }

include/dxc/Support/dxcfilesystem.h

@@ -42,9 +42,8 @@ class DxcArgsFileSystem : public ::llvm::sys::fs::MSFileSystem {
   virtual HRESULT RegisterOutputStream(LPCWSTR pName, IStream *pStream) = 0;
 };
 
-HRESULT
+DxcArgsFileSystem *
 CreateDxcArgsFileSystem(_In_ IDxcBlob *pSource, _In_ LPCWSTR pSourceName,
-                        _In_opt_ IDxcIncludeHandler *pIncludeHandler,
-                        _Outptr_ DxcArgsFileSystem **ppResult) throw();
+                        _In_opt_ IDxcIncludeHandler *pIncludeHandler);
 
 } // namespace dxcutil