text: //*********************************************************
  //
  // Copyright (c) Microsoft. All rights reserved.
  // This code is licensed under the MIT License (MIT).
  // THIS CODE IS PROVIDED *AS IS* WITHOUT WARRANTY OF
  // ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING ANY
  // IMPLIED WARRANTIES OF FITNESS FOR A PARTICULAR
  // PURPOSE, MERCHANTABILITY, OR NON-INFRINGEMENT.
  //
  //*********************************************************
  
  #pragma once
  
  #include <stdexcept>
  #include <unordered_map>
  
  #include <windows.h>
  #include <windowsx.h>
  #include <wrl/client.h>
  #include <d3d12.h>
  #include <dxgi1_4.h>
  #include <DirectXMath.h>
  #include <memory>
  
  #include <atlbase.h>
  #include <wrl.h>
  using Microsoft::WRL::ComPtr;
  
  #include <sstream>
  #include <wchar.h>
  
  #define SizeOfInUint32(obj) ((sizeof(obj) - 1) / sizeof(UINT32) + 1)
  
  inline UINT Align(UINT size, UINT alignment)
  {
      return (size + (alignment - 1)) & ~(alignment - 1);
  }
  
  class HrException : public std::runtime_error
  {
      inline std::string HrToString(HRESULT hr)
      {
          char s_str[64] = {};
          sprintf_s(s_str, "HRESULT of 0x%08X", static_cast<UINT>(hr));
          return std::string(s_str);
      }
  public:
      HrException(HRESULT hr) : std::runtime_error(HrToString(hr)), m_hr(hr) {}
      HRESULT Error() const { return m_hr; }
  private:
      const HRESULT m_hr;
  };
  
  #define SAFE_RELEASE(p) if (p) (p)->Release()
  
  inline void ThrowIfFailed(HRESULT hr)
  {
      if (FAILED(hr))
      {
          throw HrException(hr);
      }
  }
  
  inline void ThrowIfFailed(HRESULT hr, const char* msg)
  {
      if (FAILED(hr))
      {
          OutputDebugString(msg);
          throw HrException(hr);
      }
  }
  
  inline void ThrowIfFalse(bool value)
  {
      ThrowIfFailed(value ? S_OK : E_FAIL);
  }
  
  inline void ThrowIfFalse(bool value, const char* msg)
  {
      ThrowIfFailed(value ? S_OK : E_FAIL, msg);
  }
  
  // Assign a name to the object to aid with debugging.
  #if defined(_DEBUG) || defined(DBG)
  inline void SetName(ID3D12Object* pObject, const char* name)
  {
      WCHAR fullName[50];
      if (swprintf_s(fullName, L"%s", name) > 0)
      {
          pObject->SetName(fullName);
      }
  }
  inline void SetNameIndexed(ID3D12Object* pObject, const char* name, UINT index)
  {
      WCHAR fullName[50];
      if (swprintf_s(fullName, L"%s[%u]", name, index) > 0)
      {
          pObject->SetName(fullName);
      }
  }
  #else
  inline void SetName(ID3D12Object*, const char*)
  {
  }
  inline void SetNameIndexed(ID3D12Object*, const char*, UINT)
  {
  }
  #endif
  
  // Naming helper for ComPtr<T>.
  // Assigns the name of the variable as the name of the object.
  // The indexed variant will include the index in the name of the object.
  #define NAME_D3D12_OBJECT(x) SetName((x).Get(), #x)
  #define NAME_D3D12_OBJECT_INDEXED(x, n) SetNameIndexed((x)[n].Get(), #x, n)
  
  struct AccelerationStructureBuffers
  {
      ComPtr<ID3D12Resource> scratch;
      ComPtr<ID3D12Resource> accelerationStructure;
      ComPtr<ID3D12Resource> instanceDesc;    // Used only for top-level AS
      UINT64                 ResultDataMaxSizeInBytes;
  };
  
  struct D3D12Buffer
  {
      ComPtr<ID3D12Resource> resource;
      D3D12_CPU_DESCRIPTOR_HANDLE cpuDescriptorHandle;
      D3D12_GPU_DESCRIPTOR_HANDLE gpuDescriptorHandle;
  };
  
  class D3D12Texture
  {
  public:
      ComPtr<ID3D12Resource> resource;
      D3D12_GPU_DESCRIPTOR_HANDLE uavGPUHandle;
      D3D12_GPU_DESCRIPTOR_HANDLE srvGPUHandle;
      UINT uavHeapIndex;
      UINT srvHeapIndex;
  
      D3D12Texture() 
      {
          uavHeapIndex = UINT_MAX;
          srvHeapIndex = UINT_MAX;
      }
  
      void Reset()
      {
          uavHeapIndex = UINT_MAX;
          srvHeapIndex = UINT_MAX;
      }
  };
  
  class GpuUploadBuffer
  {
  public:
      ComPtr<ID3D12Resource> GetResource() { return m_resource; }
  
  protected:
      ComPtr<ID3D12Resource> m_resource;
  
      GpuUploadBuffer() {}
      ~GpuUploadBuffer()
      {
          if (m_resource.Get())
          {
              m_resource->Unmap(0, nullptr);
          }
      }
  
      void Allocate(ID3D12Device* device, UINT bufferSize, LPCWSTR resourceName = nullptr)
      {
          auto uploadHeapProperties = CD3DX12_HEAP_PROPERTIES(D3D12_HEAP_TYPE_UPLOAD);
  
          auto bufferDesc = CD3DX12_RESOURCE_DESC::Buffer(bufferSize);
          ThrowIfFailed(device->CreateCommittedResource(
              &uploadHeapProperties,
              D3D12_HEAP_FLAG_NONE,
              &bufferDesc,
              D3D12_RESOURCE_STATE_GENERIC_READ,
              nullptr,
              IID_PPV_ARGS(&m_resource)));
          m_resource->SetName(resourceName);
      }
  
      uint8_t* MapCpuWriteOnly()
      {
          uint8_t* mappedData;
          // We don't unmap this until the app closes. Keeping buffer mapped for the lifetime of the resource is okay.
          CD3DX12_RANGE readRange(0, 0);        // We do not intend to read from this resource on the CPU.
          ThrowIfFailed(m_resource->Map(0, &readRange, reinterpret_cast<void**>(&mappedData)));
          return mappedData;
      }
  };
  
  // Shader record = {{Shader ID}, {RootArguments}}
  class ShaderRecord
  {
  public:
      ShaderRecord(void* pShaderIdentifier, UINT shaderIdentifierSize) :
          shaderIdentifier(pShaderIdentifier, shaderIdentifierSize)
      {
      }
  
      ShaderRecord(void* pShaderIdentifier, UINT shaderIdentifierSize, void* pLocalRootArguments, UINT localRootArgumentsSize) :
          shaderIdentifier(pShaderIdentifier, shaderIdentifierSize),
          localRootArguments(pLocalRootArguments, localRootArgumentsSize)
      {
      }
  
      void CopyTo(void* dest) const
      {
          uint8_t* byteDest = static_cast<uint8_t*>(dest);
          memcpy(byteDest, shaderIdentifier.ptr, shaderIdentifier.size);
          if (localRootArguments.ptr)
          {
              memcpy(byteDest + shaderIdentifier.size, localRootArguments.ptr, localRootArguments.size);
          }
      }
  
      struct PointerWithSize {
          void *ptr;
          UINT size;
  
          PointerWithSize() : ptr(nullptr), size(0) {}
          PointerWithSize(void* _ptr, UINT _size) : ptr(_ptr), size(_size) {};
      };
      PointerWithSize shaderIdentifier;
      PointerWithSize localRootArguments;
  };
  
  // Shader table = {{ ShaderRecord 1}, {ShaderRecord 2}, ...}
  class ShaderTable : public GpuUploadBuffer
  {
      uint8_t* m_mappedShaderRecords;
      UINT m_shaderRecordSize;
  
      // Debug support
      std::wstring m_name;
      std::vector<ShaderRecord> m_shaderRecords;
  
      ShaderTable() {}
  public:
      ShaderTable(ID3D12Device* device, UINT numShaderRecords, UINT shaderRecordSize, LPCWSTR resourceName = nullptr) 
          : m_name(resourceName)
      {
          m_shaderRecordSize = Align(shaderRecordSize, D3D12_RAYTRACING_SHADER_RECORD_BYTE_ALIGNMENT);
          m_shaderRecords.reserve(numShaderRecords);
          UINT bufferSize = numShaderRecords * m_shaderRecordSize;
          Allocate(device, bufferSize, resourceName);
          m_mappedShaderRecords = MapCpuWriteOnly();
      }
      
      void push_back(const ShaderRecord& shaderRecord)
      {
          ThrowIfFalse(m_shaderRecords.size() < m_shaderRecords.capacity());
          m_shaderRecords.push_back(shaderRecord);
          shaderRecord.CopyTo(m_mappedShaderRecords);
          m_mappedShaderRecords += m_shaderRecordSize;
      }
  
      UINT GetShaderRecordSize() { return m_shaderRecordSize; }
  
      // Pretty-print the shader records.
      void DebugPrint(std::unordered_map<void*, std::string> shaderIdToStringMap)
      {
          std::stringstream str;
          str << "|--------------------------------------------------------------------\n";
          str << "|Shader table - " << m_name.c_str() << L": " 
               << m_shaderRecordSize << L" | "
               << m_shaderRecords.size() * m_shaderRecordSize << L" bytes\n";
  
          for (UINT i = 0; i < m_shaderRecords.size(); i++)
          {
              str << "| [" << i << L"]: ";
              str << shaderIdToStringMap[m_shaderRecords[i].shaderIdentifier.ptr] << L", ";
              str << m_shaderRecords[i].shaderIdentifier.size << L" + " << m_shaderRecords[i].localRootArguments.size << L" bytes \n";
          }
          str << "|--------------------------------------------------------------------\n";
          str << "\n";
          OutputDebugString(str.str().c_str());
      }
  };
  
  inline void AllocateUAVBuffer(ID3D12Device* pDevice, UINT64 bufferSize, ID3D12Resource **ppResource, D3D12_RESOURCE_STATES initialResourceState = D3D12_RESOURCE_STATE_COMMON, const wchar_t* resourceName = nullptr)
  {
      auto uploadHeapProperties = CD3DX12_HEAP_PROPERTIES(D3D12_HEAP_TYPE_DEFAULT);
      auto bufferDesc = CD3DX12_RESOURCE_DESC::Buffer(bufferSize, D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS);
      ThrowIfFailed(pDevice->CreateCommittedResource(
          &uploadHeapProperties,
          D3D12_HEAP_FLAG_NONE,
          &bufferDesc,
          initialResourceState,
          nullptr,
          IID_PPV_ARGS(ppResource)));
      if (resourceName)
      {
          (*ppResource)->SetName(resourceName);
      }
  }
  
  inline void AllocateUploadBuffer(ID3D12Device* pDevice, void *pData, UINT64 datasize, ID3D12Resource **ppResource, const wchar_t* resourceName = nullptr)
  {
      auto uploadHeapProperties = CD3DX12_HEAP_PROPERTIES(D3D12_HEAP_TYPE_UPLOAD);
      auto bufferDesc = CD3DX12_RESOURCE_DESC::Buffer(datasize);
      ThrowIfFailed(pDevice->CreateCommittedResource(
          &uploadHeapProperties,
          D3D12_HEAP_FLAG_NONE,
          &bufferDesc,
          D3D12_RESOURCE_STATE_GENERIC_READ,
          nullptr,
          IID_PPV_ARGS(ppResource)));
      if (resourceName)
      {
          (*ppResource)->SetName(resourceName);
      }
      void *pMappedData;
      (*ppResource)->Map(0, nullptr, &pMappedData);
      memcpy(pMappedData, pData, datasize);
      (*ppResource)->Unmap(0, nullptr);
  }
  
  // Pretty-print a state object tree.
  inline void PrintStateObjectDesc(const D3D12_STATE_OBJECT_DESC* desc)
  {
      std::wstringstream wstr;
      wstr << L"\n";
      wstr << L"--------------------------------------------------------------------\n";
      wstr << L"| D3D12 State Object 0x" << static_cast<const void*>(desc) << L": ";
      if (desc->Type == D3D12_STATE_OBJECT_TYPE_COLLECTION) wstr << L"Collection\n";
      if (desc->Type == D3D12_STATE_OBJECT_TYPE_RAYTRACING_PIPELINE) wstr << L"Raytracing Pipeline\n";
  
      auto ExportTree = [](UINT depth, UINT numExports, const D3D12_EXPORT_DESC* exports)
      {
          std::wostringstream woss;
          for (UINT i = 0; i < numExports; i++)
          {
              woss << L"|";
              if (depth > 0)
              {
                  for (UINT j = 0; j < 2 * depth - 1; j++) woss << L" ";
              }
              woss << L" [" << i << L"]: ";
              if (exports[i].ExportToRename) woss << exports[i].ExportToRename << L" --> ";
              woss << exports[i].Name << L"\n";
          }
          return woss.str();
      };
  
      for (UINT i = 0; i < desc->NumSubobjects; i++)
      {
          wstr << L"| [" << i << L"]: ";
          switch (desc->pSubobjects[i].Type)
          {
          case D3D12_STATE_SUBOBJECT_TYPE_GLOBAL_ROOT_SIGNATURE:
              wstr << L"Global Root Signature 0x" << desc->pSubobjects[i].pDesc << L"\n";
              break;
          case D3D12_STATE_SUBOBJECT_TYPE_LOCAL_ROOT_SIGNATURE:
              wstr << L"Local Root Signature 0x" << desc->pSubobjects[i].pDesc << L"\n";
              break;
          case D3D12_STATE_SUBOBJECT_TYPE_NODE_MASK:
              //wstr << L"Node Mask: 0x" << std::hex << std::setfill(L'0') << std::setw(8) << *static_cast<const UINT*>(desc->pSubobjects[i].pDesc) << std::setw(0) << std::dec << L"\n";
              break;
          case D3D12_STATE_SUBOBJECT_TYPE_DXIL_LIBRARY:
          {
              wstr << L"DXIL Library 0x";
              auto lib = static_cast<const D3D12_DXIL_LIBRARY_DESC*>(desc->pSubobjects[i].pDesc);
              wstr << lib->DXILLibrary.pShaderBytecode << L", " << lib->DXILLibrary.BytecodeLength << L" bytes\n";
              wstr << ExportTree(1, lib->NumExports, lib->pExports);
              break;
          }
          case D3D12_STATE_SUBOBJECT_TYPE_EXISTING_COLLECTION:
          {
              wstr << L"Existing Library 0x";
              auto collection = static_cast<const D3D12_EXISTING_COLLECTION_DESC*>(desc->pSubobjects[i].pDesc);
              wstr << collection->pExistingCollection << L"\n";
              wstr << ExportTree(1, collection->NumExports, collection->pExports);
              break;
          }
          case D3D12_STATE_SUBOBJECT_TYPE_SUBOBJECT_TO_EXPORTS_ASSOCIATION:
          {
              wstr << L"Subobject to Exports Association (Subobject [";
              auto association = static_cast<const D3D12_SUBOBJECT_TO_EXPORTS_ASSOCIATION*>(desc->pSubobjects[i].pDesc);
              UINT index = static_cast<UINT>(association->pSubobjectToAssociate - desc->pSubobjects);
              wstr << index << L"])\n";
              for (UINT j = 0; j < association->NumExports; j++)
              {
                  wstr << L"|  [" << j << L"]: " << association->pExports[j] << L"\n";
              }
              break;
          }
          case D3D12_STATE_SUBOBJECT_TYPE_DXIL_SUBOBJECT_TO_EXPORTS_ASSOCIATION:
          {
              wstr << L"DXIL Subobjects to Exports Association (";
              auto association = static_cast<const D3D12_DXIL_SUBOBJECT_TO_EXPORTS_ASSOCIATION*>(desc->pSubobjects[i].pDesc);
              wstr << association->SubobjectToAssociate << L")\n";
              for (UINT j = 0; j < association->NumExports; j++)
              {
                  wstr << L"|  [" << j << L"]: " << association->pExports[j] << L"\n";
              }
              break;
          }
          case D3D12_STATE_SUBOBJECT_TYPE_RAYTRACING_SHADER_CONFIG:
          {
              wstr << L"Raytracing Shader Config\n";
              auto config = static_cast<const D3D12_RAYTRACING_SHADER_CONFIG*>(desc->pSubobjects[i].pDesc);
              wstr << L"|  [0]: Max Payload Size: " << config->MaxPayloadSizeInBytes << L" bytes\n";
              wstr << L"|  [1]: Max Attribute Size: " << config->MaxAttributeSizeInBytes << L" bytes\n";
              break;
          }
          case D3D12_STATE_SUBOBJECT_TYPE_RAYTRACING_PIPELINE_CONFIG:
          {
              wstr << L"Raytracing Pipeline Config\n";
              auto config = static_cast<const D3D12_RAYTRACING_PIPELINE_CONFIG*>(desc->pSubobjects[i].pDesc);
              wstr << L"|  [0]: Max Recursion Depth: " << config->MaxTraceRecursionDepth << L"\n";
              break;
          }
          case D3D12_STATE_SUBOBJECT_TYPE_HIT_GROUP:
          {
              wstr << L"Hit Group (";
              auto hitGroup = static_cast<const D3D12_HIT_GROUP_DESC*>(desc->pSubobjects[i].pDesc);
              wstr << (hitGroup->HitGroupExport ? hitGroup->HitGroupExport : L"[none]") << L")\n";
              wstr << L"|  [0]: Any Hit Import: " << (hitGroup->AnyHitShaderImport ? hitGroup->AnyHitShaderImport : L"[none]") << L"\n";
              wstr << L"|  [1]: Closest Hit Import: " << (hitGroup->ClosestHitShaderImport ? hitGroup->ClosestHitShaderImport : L"[none]") << L"\n";
              wstr << L"|  [2]: Intersection Import: " << (hitGroup->IntersectionShaderImport ? hitGroup->IntersectionShaderImport : L"[none]") << L"\n";
              break;
          }
          }
          wstr << L"|--------------------------------------------------------------------\n";
      }
      wstr << L"\n";
      OutputDebugStringW(wstr.str().c_str());
  }