KSG-Framework/DX12Shader.cpp

#include "DX12Shader.h"

#include "DX12CommandQueue.h"
#include "DX12SamplerState.h"
#include "DX12Texture.h"
#include "DX12TLAS.h"
#include "Logging.h"

using namespace Framework;

Framework::DX12ShaderSignature::DX12ShaderSignature(bool global)
    : ReferenceCounter(),
      signature(0),
      changed(1),
      shaderBindingTableParamCount(0),
      global(global)
{}

Framework::DX12ShaderSignature::~DX12ShaderSignature()
{
    if (signature)
    {
        signature->Release();
    }
    for (DX12ShaderRegisterUsage* usage : descriptorHeapBindings)
    {
        delete usage;
    }
    for (DX12ShaderRegisterUsage* usage : bindingTableBindings)
    {
        delete usage;
    }
}

int* Framework::DX12ShaderSignature::addRegisterUsageLinkedToShaderBindingTable(
    DX12ShaderRegister registerType, int registerIndex, int spaceIndex)
{
    if (registerType == DX12_SHADER_REGISTER_S_SAMPLER)
    {
        Logging::error() << "Sampler register usage can not be linked to "
                            "shader binding table";
        throw std::logic_error(
            "Sampler register usage can not be linked to shader binding table");
    }
    for (DX12ShaderRegisterUsage* usage : descriptorHeapBindings)
    {
        if (usage->registerType == registerType
            && usage->registerIndex == registerIndex
            && usage->spaceIndex == spaceIndex)
        {
            Logging::error()
                << "Duplicate register usage in root signature: "
                << registerType << " " << registerIndex << " " << spaceIndex;
            throw std::invalid_argument(
                "Duplicate register usage in root signature");
        }
    }
    for (DX12ShaderRegisterUsage* usage : bindingTableBindings)
    {
        if (usage->registerType == registerType
            && usage->registerIndex == registerIndex
            && usage->spaceIndex == spaceIndex)
        {
            Logging::error()
                << "Duplicate register usage in root signature: "
                << registerType << " " << registerIndex << " " << spaceIndex;
            throw std::invalid_argument(
                "Duplicate register usage in root signature");
        }
    }
    DX12ShaderRegisterUsage* usage = new DX12ShaderRegisterUsage{registerType,
        GLOBAL_DESCRIPTOR_HEAP,
        registerIndex,
        spaceIndex,
        -1,
        0,
        0,
        -1};
    bindingTableBindings.add(usage);
    changed = 1;
    return &usage->bindingTableIndex;
}

void Framework::DX12ShaderSignature::addRegisterUsageLinkedToDescriptorHeap(
    int descriptorHeapIndex,
    DX12ShaderRegister registerType,
    int registerIndex,
    int spaceIndex,
    DX12DescriptorHeapType descriptorHeapType,
    bool array,
    int arraySize)
{
    if ((registerType == DX12_SHADER_REGISTER_S_SAMPLER)
        != (descriptorHeapType == SAMPLER_DESCRIPTOR_HEAP))
    {
        Logging::error()
            << "Sampler register usage must be linked to sampler descriptor "
               "heap and vice versa";
        throw std::logic_error(
            "Sampler register usage must be linked to sampler descriptor "
            "heap and vice versa");
    }
    if (descriptorHeapIndex < 0)
    {
        Logging::error() << "descriptorHeapIndex can not be below 0";
        throw std::invalid_argument("descriptorHeapIndex can not be below 0");
    }
    for (DX12ShaderRegisterUsage* usage : descriptorHeapBindings)
    {
        if (usage->registerType == registerType
            && usage->registerIndex == registerIndex
            && usage->spaceIndex == spaceIndex)
        {
            Logging::error()
                << "Duplicate register usage in root signature: "
                << registerType << " " << registerIndex << " " << spaceIndex;
            throw std::invalid_argument(
                "Duplicate register usage in root signature");
        }
    }
    // descriptor heap bindings sould be sorted by descriptorHeapType ->
    // registerType -> spaceIndex -> registerIndex
    ArrayIterator<DX12ShaderRegisterUsage*> it = descriptorHeapBindings.begin();
    bool found = 0;
    while (it)
    {
        if (it->descriptorHeapType > descriptorHeapType)
        {
            found = 1;
            break;
        }
        if (it->descriptorHeapType == descriptorHeapType)
        {
            if (it->registerType > registerType)
            {
                found = 1;
                break;
            }
            if (it->registerType == registerType)
            {
                if (it->spaceIndex > spaceIndex)
                {
                    found = 1;
                    break;
                }
                if (it->spaceIndex == spaceIndex)
                {
                    if (it->registerIndex >= registerIndex)
                    {
                        found = 1;
                        break;
                    }
                }
            }
        }
        ++it;
    }
    if (found)
    {
        if (it->registerIndex == registerIndex && it->spaceIndex == spaceIndex
            && it->registerType == registerType)
        {
            Logging::error()
                << "Duplicate register usage in root signature: "
                << registerType << " " << registerIndex << " " << spaceIndex;
            throw std::invalid_argument(
                "Duplicate register usage in root signature");
        }
        it.addBefore(new DX12ShaderRegisterUsage{registerType,
            descriptorHeapType,
            registerIndex,
            spaceIndex,
            descriptorHeapIndex,
            0,
            array,
            arraySize});
    }
    else
    {
        descriptorHeapBindings.add(new DX12ShaderRegisterUsage{registerType,
            descriptorHeapType,
            registerIndex,
            spaceIndex,
            descriptorHeapIndex,
            0,
            array,
            arraySize});
    }
    changed = 1;
}

void Framework::DX12ShaderSignature::createSignature(ID3D12Device5* zDevice,
    PFN_D3D12_SERIALIZE_ROOT_SIGNATURE pfnD3D12SerializeRootSignature)
{
    if (!changed)
    {
        return;
    }
    changed = 0;
    if (signature)
    {
        signature->Release();
        signature = 0;
    }
    int paramCount = (descriptorHeapBindings.getEntryCount() > 0 ? 2 : 0)
                   + bindingTableBindings.getEntryCount();
    D3D12_ROOT_PARAMETER* descriptorTable
        = new D3D12_ROOT_PARAMETER[paramCount];
    int index = 0;
    D3D12_DESCRIPTOR_RANGE** descriptorRanges = new D3D12_DESCRIPTOR_RANGE*[2];
    int rangeCount = 0;
    ArrayIterator<DX12ShaderRegisterUsage*> it = descriptorHeapBindings.begin();
    DX12DescriptorHeapType currentDescriptorHeapType;
    useGlobalDescriptorHeap = 0;
    useSamplerDescriptorHeap = 0;
    while (it)
    {
        descriptorTable[index].ParameterType
            = D3D12_ROOT_PARAMETER_TYPE_DESCRIPTOR_TABLE;
        descriptorTable[index].ShaderVisibility = D3D12_SHADER_VISIBILITY_ALL;
        descriptorRanges[index]
            = new D3D12_DESCRIPTOR_RANGE[descriptorHeapBindings
                    .getEntryCount()];
        currentDescriptorHeapType = it->descriptorHeapType;
        if (currentDescriptorHeapType == GLOBAL_DESCRIPTOR_HEAP)
        {
            useGlobalDescriptorHeap = true;
        }
        else if (currentDescriptorHeapType == SAMPLER_DESCRIPTOR_HEAP)
        {
            useSamplerDescriptorHeap = true;
        }
        rangeCount = 0;
        while (it && it->descriptorHeapType == currentDescriptorHeapType)
        {
            it->bindingTableIndex = 0;
            D3D12_DESCRIPTOR_RANGE& range = descriptorRanges[index][rangeCount];
            switch (it->registerType)
            {
            case DX12_SHADER_REGISTER_B_CONST_BUFFER:
                range.RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_CBV;
                break;
            case DX12_SHADER_REGISTER_T_SHADER_RESOURCE:
                range.RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_SRV;
                break;
            case DX12_SHADER_REGISTER_U_UNORDERED_ACCESS:
                range.RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_UAV;
                break;
            case DX12_SHADER_REGISTER_S_SAMPLER:
                range.RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_SAMPLER;
                break;
            }
            range.BaseShaderRegister = it->registerIndex;
            range.OffsetInDescriptorsFromTableStart = it->descriptorHeapIndex;
            range.RegisterSpace = it->spaceIndex;
            if (it->array)
            {
                range.NumDescriptors = it->arraySize;
            }
            else
            {
                ArrayIterator<DX12ShaderRegisterUsage*> next = it.next();
                int size = 1;
                while (next && next->registerType == it->registerType
                       && next->spaceIndex == it->spaceIndex
                       && next->registerIndex == it->registerIndex + size
                       && next->descriptorHeapIndex
                              == it->descriptorHeapIndex + size
                       && it->descriptorHeapType == currentDescriptorHeapType)
                {
                    ++size;
                    it = next;
                    it->bindingTableIndex = 0;
                    ++next;
                }
                range.NumDescriptors = size;
            }
            ++it;
            ++rangeCount;
        }
        descriptorTable[index].DescriptorTable.pDescriptorRanges
            = descriptorRanges[index];
        descriptorTable[index].DescriptorTable.NumDescriptorRanges = rangeCount;
        index++;
    }
    for (DX12ShaderRegisterUsage* usage : bindingTableBindings)
    {
        usage->bindingTableIndex = index;
        switch (usage->registerType)
        {
        case DX12_SHADER_REGISTER_B_CONST_BUFFER:
            descriptorTable[index].ParameterType
                = D3D12_ROOT_PARAMETER_TYPE_CBV;
            break;
        case DX12_SHADER_REGISTER_T_SHADER_RESOURCE:
            descriptorTable[index].ParameterType
                = D3D12_ROOT_PARAMETER_TYPE_SRV;
            break;
        case DX12_SHADER_REGISTER_U_UNORDERED_ACCESS:
            descriptorTable[index].ParameterType
                = D3D12_ROOT_PARAMETER_TYPE_UAV;
            break;
        }
        descriptorTable[index].ShaderVisibility = D3D12_SHADER_VISIBILITY_ALL;
        descriptorTable[index].Descriptor.ShaderRegister = usage->registerIndex;
        descriptorTable[index].Descriptor.RegisterSpace = usage->spaceIndex;
        ++index;
    }
    shaderBindingTableParamCount = index;
    D3D12_ROOT_SIGNATURE_DESC rootDesc = {};
    rootDesc.NumParameters = shaderBindingTableParamCount;
    rootDesc.pParameters = descriptorTable;
    rootDesc.Flags = global ? D3D12_ROOT_SIGNATURE_FLAG_NONE
                            : D3D12_ROOT_SIGNATURE_FLAG_LOCAL_ROOT_SIGNATURE;
    ID3DBlob* pSigBlob = 0;
    ID3DBlob* pErrorBlob = 0;
    HRESULT hr = pfnD3D12SerializeRootSignature(
        &rootDesc, D3D_ROOT_SIGNATURE_VERSION_1_0, &pSigBlob, &pErrorBlob);
    if (pSigBlob)
    {
        zDevice->CreateRootSignature(0,
            pSigBlob->GetBufferPointer(),
            pSigBlob->GetBufferSize(),
            __uuidof(ID3D12RootSignature),
            (void**)&signature);
        pSigBlob->Release();
    }
    if (pErrorBlob)
    {
        std::string errorMessage(
            static_cast<const char*>(pErrorBlob->GetBufferPointer()),
            pErrorBlob->GetBufferSize());
        Logging::error() << "Failed to serialize root signature: "
                         << errorMessage;
        pErrorBlob->Release();
    }
    delete[] descriptorRanges;
    delete[] descriptorTable;
}

ID3D12RootSignature* Framework::DX12ShaderSignature::zSignature() const
{
    return signature;
}

const Array<DX12ShaderRegisterUsage*>&
Framework::DX12ShaderSignature::getDescriptorHeapBindings() const
{
    return descriptorHeapBindings;
}

int Framework::DX12ShaderSignature::gerShaderBindingTableParamCount() const
{
    return shaderBindingTableParamCount;
}

bool Framework::DX12ShaderSignature::doesUseGlobalDescriptorHeap() const
{
    return useGlobalDescriptorHeap;
}

bool Framework::DX12ShaderSignature::doesUseSamplerDescriptorHeap() const
{
    return useSamplerDescriptorHeap;
}

Framework::DX12ShaderFunction::DX12ShaderFunction(const Text& functionName,
    DX12ShaderSignature* signature,
    DX12ShaderFunctionType functionType)
    : ReferenceCounter(),
      functionName(functionName),
      signature(signature),
      functionType(functionType),
      exportDesc(new D3D12_EXPORT_DESC())
{
    wchar_t* wc = new wchar_t[functionName.getLength() + 1];
    mbstowcs_s(0,
        wc,
        functionName.getLength() + 1,
        functionName.getText(),
        functionName.getLength() + 1);
    wc[functionName.getLength()] = 0;
    exportDesc->Name = wc;
    exportDesc->ExportToRename = 0;
    exportDesc->Flags = D3D12_EXPORT_FLAG_NONE;
}

Framework::DX12ShaderFunction::~DX12ShaderFunction()
{
    signature->release();
    delete[] exportDesc->Name;
    delete exportDesc;
}

const Text& Framework::DX12ShaderFunction::getFunctionName() const
{
    return functionName;
}

DX12ShaderSignature* Framework::DX12ShaderFunction::zSignature() const
{
    return signature;
}

D3D12_EXPORT_DESC* Framework::DX12ShaderFunction::zExportDesc() const
{
    return exportDesc;
}

DX12ShaderFunctionType Framework::DX12ShaderFunction::getFunctionType() const
{
    return functionType;
}

Framework::DX12Shader::DX12Shader(
    const unsigned char* shaderBytes, int shaderBytesSize)
    : ReferenceCounter(),
      shaderBytes(shaderBytes),
      shaderBytesSize(shaderBytesSize),
      libraryDesc(new D3D12_DXIL_LIBRARY_DESC())
{
    libraryDesc->DXILLibrary.pShaderBytecode = shaderBytes;
    libraryDesc->DXILLibrary.BytecodeLength = shaderBytesSize;
    libraryDesc->NumExports = 0;
    libraryDesc->pExports = 0;
}

Framework::DX12Shader::~DX12Shader()
{
    delete[] libraryDesc->pExports;
    delete libraryDesc;
}

void Framework::DX12Shader::addFunction(DX12ShaderFunction* function)
{
    functions.add(function);
}

int Framework::DX12Shader::getShaderBytesSize() const
{
    return shaderBytesSize;
}

const unsigned char* Framework::DX12Shader::getShaderBytes() const
{
    return shaderBytes;
}

const RCArray<DX12ShaderFunction>& Framework::DX12Shader::getFunctions() const
{
    return functions;
}

D3D12_DXIL_LIBRARY_DESC* Framework::DX12Shader::zLibraryDesc() const
{
    if (libraryDesc->NumExports != functions.getEntryCount())
    {
        delete[] libraryDesc->pExports;
        libraryDesc->NumExports = functions.getEntryCount();
        D3D12_EXPORT_DESC* pExports
            = new D3D12_EXPORT_DESC[functions.getEntryCount()];
        int index = 0;
        for (const auto& function : functions)
        {
            memcpy(pExports + index,
                function->zExportDesc(),
                sizeof(D3D12_EXPORT_DESC));
            index++;
        }
        libraryDesc->pExports = pExports;
    }
    return libraryDesc;
}

Framework::DX12ShaderHitGroup::DX12ShaderHitGroup(const Text name)
    : ReferenceCounter(),
      name(name),
      closestHitShaderFunction(0),
      anyHitShaderFunction(0),
      intersectionShaderFunction(0),
      payloadSize(0),
      attributeSize(0),
      hitGroupDesc(new D3D12_HIT_GROUP_DESC())
{
    wchar_t* wc = new wchar_t[name.getLength() + 1];
    mbstowcs_s(
        0, wc, name.getLength() + 1, name.getText(), name.getLength() + 1);
    wc[name.getLength()] = 0;
    hitGroupDesc->HitGroupExport = wc;
    hitGroupDesc->IntersectionShaderImport = 0;
    hitGroupDesc->AnyHitShaderImport = 0;
    hitGroupDesc->ClosestHitShaderImport = 0;
}

Framework::DX12ShaderHitGroup::~DX12ShaderHitGroup()
{
    delete[] hitGroupDesc->HitGroupExport;
    delete hitGroupDesc;
    if (closestHitShaderFunction)
    {
        closestHitShaderFunction->release();
    }
    if (anyHitShaderFunction)
    {
        anyHitShaderFunction->release();
    }
    if (intersectionShaderFunction)
    {
        intersectionShaderFunction->release();
    }
}

void Framework::DX12ShaderHitGroup::setClosestHitShaderFunction(
    DX12ShaderFunction* zClosestHitShaderFunction)
{
    if (this->closestHitShaderFunction == zClosestHitShaderFunction)
    {
        return;
    }
    if (anyHitShaderFunction
        && anyHitShaderFunction->zSignature()
               != zClosestHitShaderFunction->zSignature())
    {
        Logging::error()
            << "Any-hit shader function and closest-hit shader "
               "function must have the same root signature when they are "
               "combined in the same hit group. HitGroup Name: '"
            << name << "' Any Hit Shader Function: '"
            << anyHitShaderFunction->getFunctionName().getText()
            << "' Closest Hit Shader Function: '"
            << zClosestHitShaderFunction->getFunctionName().getText() << "'";
        throw std::runtime_error("Incompatible root signatures in hit group");
    }
    if (intersectionShaderFunction
        && intersectionShaderFunction->zSignature()
               != zClosestHitShaderFunction->zSignature())
    {
        Logging::error()
            << "Intersection shader function and closest-hit shader "
               "function must have the same root signature when they are "
               "combined in the same hit group. HitGroup Name: '"
            << name << "' Intersection Shader Function: '"
            << intersectionShaderFunction->getFunctionName().getText()
            << "' Closest Hit Shader Function: '"
            << zClosestHitShaderFunction->getFunctionName().getText() << "'";
        throw std::runtime_error("Incompatible root signatures in hit group");
    }

    if (closestHitShaderFunction)
    {
        closestHitShaderFunction->release();
    }
    closestHitShaderFunction = zClosestHitShaderFunction;
    hitGroupDesc->ClosestHitShaderImport = 0;
    if (closestHitShaderFunction)
    {
        hitGroupDesc->ClosestHitShaderImport
            = closestHitShaderFunction->zExportDesc()->Name;
        closestHitShaderFunction->getThis();
    }
}

void Framework::DX12ShaderHitGroup::setAnyHitShaderFunction(
    DX12ShaderFunction* zAnyHitShaderFunction)
{
    if (this->anyHitShaderFunction == zAnyHitShaderFunction)
    {
        return;
    }
    if (closestHitShaderFunction
        && closestHitShaderFunction->zSignature()
               != zAnyHitShaderFunction->zSignature())
    {
        Logging::error()
            << "Any-hit shader function and closest-hit shader "
               "function must have the same root signature when they are "
               "combined in the same hit group. HitGroup Name: '"
            << name << "' Any Hit Shader Function: '"
            << zAnyHitShaderFunction->getFunctionName().getText()
            << "' Closest Hit Shader Function: '"
            << closestHitShaderFunction->getFunctionName().getText() << "'";
        throw std::runtime_error("Incompatible root signatures in hit group");
    }
    if (intersectionShaderFunction
        && intersectionShaderFunction->zSignature()
               != zAnyHitShaderFunction->zSignature())
    {
        Logging::error()
            << "Intersection shader function and any-hit shader "
               "function must have the same root signature when they are "
               "combined in the same hit group. HitGroup Name: '"
            << name << "' Intersection Shader Function: '"
            << intersectionShaderFunction->getFunctionName().getText()
            << "' Any Hit Shader Function: '"
            << zAnyHitShaderFunction->getFunctionName().getText() << "'";
        throw std::runtime_error("Incompatible root signatures in hit group");
    }

    if (anyHitShaderFunction)
    {
        anyHitShaderFunction->release();
    }
    anyHitShaderFunction = zAnyHitShaderFunction;
    hitGroupDesc->AnyHitShaderImport = 0;
    if (anyHitShaderFunction)
    {
        hitGroupDesc->AnyHitShaderImport
            = anyHitShaderFunction->zExportDesc()->Name;
        anyHitShaderFunction->getThis();
    }
}

void Framework::DX12ShaderHitGroup::setIntersectionShaderFunction(
    DX12ShaderFunction* zIntersectionShaderFunction)
{
    if (intersectionShaderFunction == zIntersectionShaderFunction)
    {
        return;
    }
    if (closestHitShaderFunction
        && closestHitShaderFunction->zSignature()
               != zIntersectionShaderFunction->zSignature())
    {
        Logging::error()
            << "Intersection shader function and closest-hit shader "
               "function must have the same root signature when they are "
               "combined in the same hit group. HitGroup Name: '"
            << name << "' Intersection Shader Function: '"
            << zIntersectionShaderFunction->getFunctionName().getText()
            << "' Closest Hit Shader Function: '"
            << closestHitShaderFunction->getFunctionName().getText() << "'";
        throw std::runtime_error("Incompatible root signatures in hit group");
    }
    if (anyHitShaderFunction
        && anyHitShaderFunction->zSignature()
               != zIntersectionShaderFunction->zSignature())
    {
        Logging::error()
            << "Intersection shader function and any-hit shader "
               "function must have the same root signature when they are "
               "combined in the same hit group. HitGroup Name: '"
            << name << "' Intersection Shader Function: '"
            << zIntersectionShaderFunction->getFunctionName().getText()
            << "' Any Hit Shader Function: '"
            << anyHitShaderFunction->getFunctionName().getText() << "'";
        throw std::runtime_error("Incompatible root signatures in hit group");
    }

    if (intersectionShaderFunction)
    {
        intersectionShaderFunction->release();
    }
    intersectionShaderFunction = zIntersectionShaderFunction;
    hitGroupDesc->IntersectionShaderImport = 0;
    if (intersectionShaderFunction)
    {
        hitGroupDesc->IntersectionShaderImport
            = intersectionShaderFunction->zExportDesc()->Name;
        intersectionShaderFunction->getThis();
    }
}

void Framework::DX12ShaderHitGroup::setPayloadSize(int payloadSize)
{
    this->payloadSize = payloadSize;
}

void Framework::DX12ShaderHitGroup::setAttributeSize(int attributeSize)
{
    this->attributeSize = attributeSize;
}

const Text& Framework::DX12ShaderHitGroup::getName() const
{
    return name;
}

DX12ShaderFunction*
Framework::DX12ShaderHitGroup::zClosestHitShaderFunction() const
{
    return closestHitShaderFunction;
}

DX12ShaderFunction* Framework::DX12ShaderHitGroup::zAnyHitShaderFunction() const
{
    return anyHitShaderFunction;
}

DX12ShaderFunction*
Framework::DX12ShaderHitGroup::zIntersectionShaderFunction() const
{
    return intersectionShaderFunction;
}

int Framework::DX12ShaderHitGroup::getPayloadSize() const
{
    return payloadSize;
}

int Framework::DX12ShaderHitGroup::getAttributeSize() const
{
    return attributeSize;
}

D3D12_HIT_GROUP_DESC* Framework::DX12ShaderHitGroup::zHitGroupDesc() const
{
    if (hitGroupDesc->IntersectionShaderImport)
    {
        hitGroupDesc->Type = D3D12_HIT_GROUP_TYPE_PROCEDURAL_PRIMITIVE;
    }
    else
    {
        hitGroupDesc->Type = D3D12_HIT_GROUP_TYPE_TRIANGLES;
    }
    return hitGroupDesc;
}

DX12ShaderSignature* Framework::DX12ShaderHitGroup::zSignature() const
{
    if (closestHitShaderFunction)
    {
        return closestHitShaderFunction->zSignature();
    }
    if (anyHitShaderFunction)
    {
        return anyHitShaderFunction->zSignature();
    }
    if (intersectionShaderFunction)
    {
        return intersectionShaderFunction->zSignature();
    }
    return 0;
}

Framework::DX12Pipeline::DX12Pipeline()
    : ReferenceCounter(),
      globalSignature(new DX12ShaderSignature(true)),
      pipelineState(0),
      maxRecursionDepth(0)
{}

Framework::DX12Pipeline::~DX12Pipeline()
{
    if (globalSignature)
    {
        globalSignature->release();
    }
    if (pipelineState)
    {
        pipelineState->Release();
    }
}

void Framework::DX12Pipeline::addShader(DX12Shader* shader)
{
    shaders.add(shader);
}

void Framework::DX12Pipeline::addHitGroup(DX12ShaderHitGroup* hitGroup)
{
    hitGroups.add(hitGroup);
}

void Framework::DX12Pipeline::setMaxRecursionDepth(int maxRecursionDepth)
{
    this->maxRecursionDepth = maxRecursionDepth;
}

void Framework::DX12Pipeline::createPipelineState(ID3D12Device5* zDevice,
    PFN_D3D12_SERIALIZE_ROOT_SIGNATURE pfnD3D12SerializeRootSignature)
{
    globalSignature->createSignature(zDevice, pfnD3D12SerializeRootSignature);
    unsigned int subobjectCount
        = shaders.getEntryCount() + hitGroups.getEntryCount() + 4;
    Array<DX12ShaderSignature*> distinctSignatures;
    for (const DX12Shader* shader : shaders)
    {
        for (const DX12ShaderFunction* function : shader->getFunctions())
        {
            DX12ShaderSignature* signature = function->zSignature();
            if (distinctSignatures.getValueIndex(signature) < 0)
            {
                distinctSignatures.add(signature);
            }
        }
    }
    subobjectCount += distinctSignatures.getEntryCount()
                    * 2; // Local root signatures for each distinct signature
    D3D12_STATE_SUBOBJECT* subobjects
        = new D3D12_STATE_SUBOBJECT[subobjectCount];
    int index = 0;
    for (const DX12Shader* shader : shaders)
    {
        subobjects[index].Type = D3D12_STATE_SUBOBJECT_TYPE_DXIL_LIBRARY;
        subobjects[index].pDesc = shader->zLibraryDesc();
        index++;
    }
    int maxPayloadSize = 0;
    int maxAttributeSize = 0;
    for (const DX12ShaderHitGroup* hitGroup : hitGroups)
    {
        subobjects[index].Type = D3D12_STATE_SUBOBJECT_TYPE_HIT_GROUP;
        subobjects[index].pDesc = hitGroup->zHitGroupDesc();
        if (hitGroup->getPayloadSize() > maxPayloadSize)
        {
            maxPayloadSize = hitGroup->getPayloadSize();
        }
        if (hitGroup->getAttributeSize() > maxAttributeSize)
        {
            maxAttributeSize = hitGroup->getAttributeSize();
        }
        index++;
    }
    D3D12_RAYTRACING_SHADER_CONFIG shaderDesc = {};
    shaderDesc.MaxPayloadSizeInBytes = maxPayloadSize;
    shaderDesc.MaxAttributeSizeInBytes = maxAttributeSize;

    subobjects[index].Type
        = D3D12_STATE_SUBOBJECT_TYPE_RAYTRACING_SHADER_CONFIG;
    subobjects[index].pDesc = &shaderDesc;
    index++;
    functionsWithoutHitGroups.clear();
    // create export list for hole pipeline with all hitgroups and all functions
    // that are not part of a hit group
    for (const DX12Shader* shader : shaders)
    {
        for (const DX12ShaderFunction* function : shader->getFunctions())
        {
            bool found = 0;
            for (const DX12ShaderHitGroup* hitGroup : hitGroups)
            {
                if (hitGroup->zClosestHitShaderFunction() == function
                    || hitGroup->zAnyHitShaderFunction() == function
                    || hitGroup->zIntersectionShaderFunction() == function)
                {
                    found = 1;
                    break;
                }
            }
            if (!found)
            {
                functionsWithoutHitGroups.add(function);
            }
        }
    }
    const wchar_t** functionAndHitGroupNames
        = new const wchar_t*[functionsWithoutHitGroups.getEntryCount()
                             + hitGroups.getEntryCount()];
    int nameIndex = 0;
    for (const DX12ShaderFunction* function : functionsWithoutHitGroups)
    {
        functionAndHitGroupNames[nameIndex]
            = const_cast<wchar_t*>(function->zExportDesc()->Name);
        nameIndex++;
    }
    for (const DX12ShaderHitGroup* hitGroup : hitGroups)
    {
        functionAndHitGroupNames[nameIndex]
            = const_cast<wchar_t*>(hitGroup->zHitGroupDesc()->HitGroupExport);
        nameIndex++;
    }
    D3D12_SUBOBJECT_TO_EXPORTS_ASSOCIATION shaderPayloadAssociation = {};
    shaderPayloadAssociation.NumExports
        = functionsWithoutHitGroups.getEntryCount() + hitGroups.getEntryCount();
    shaderPayloadAssociation.pExports = functionAndHitGroupNames;
    shaderPayloadAssociation.pSubobjectToAssociate
        = &subobjects[index - 1]; // shader config
    subobjects[index].Type
        = D3D12_STATE_SUBOBJECT_TYPE_SUBOBJECT_TO_EXPORTS_ASSOCIATION;
    subobjects[index].pDesc = &shaderPayloadAssociation;
    index++;
    D3D12_LOCAL_ROOT_SIGNATURE* localRootSignatures
        = new D3D12_LOCAL_ROOT_SIGNATURE[distinctSignatures.getEntryCount()];
    int rootSignatureIndex = 0;
    const wchar_t*** rootSignatureExports
        = new const wchar_t**[distinctSignatures.getEntryCount()];
    D3D12_SUBOBJECT_TO_EXPORTS_ASSOCIATION* localRootAssociations
        = new D3D12_SUBOBJECT_TO_EXPORTS_ASSOCIATION[distinctSignatures
                .getEntryCount()];
    for (DX12ShaderSignature* signature : distinctSignatures)
    {
        signature->createSignature(zDevice, pfnD3D12SerializeRootSignature);
        subobjects[index].Type
            = D3D12_STATE_SUBOBJECT_TYPE_LOCAL_ROOT_SIGNATURE;
        localRootSignatures[rootSignatureIndex].pLocalRootSignature
            = signature->zSignature();
        subobjects[index].pDesc = &localRootSignatures[rootSignatureIndex];
        index++;

        Array<const DX12ShaderFunction*> functionsWithThisSignature;
        for (const DX12Shader* shader : shaders)
        {
            for (const DX12ShaderFunction* function : shader->getFunctions())
            {
                if (function->zSignature() == signature)
                {
                    bool found = 0;
                    for (DX12ShaderHitGroup* hitGroup : hitGroups)
                    {
                        if (hitGroup->zClosestHitShaderFunction() == function
                            || hitGroup->zAnyHitShaderFunction() == function
                            || hitGroup->zIntersectionShaderFunction()
                                   == function)
                        {
                            found = 1;
                            break;
                        }
                    }
                    if (!found)
                    {
                        functionsWithThisSignature.add(function);
                    }
                }
            }
        }
        Array<const DX12ShaderHitGroup*> hitGroupsWithThisSignature;
        for (const DX12ShaderHitGroup* hitGroup : hitGroups)
        {
            if (hitGroup->zClosestHitShaderFunction()
                && hitGroup->zClosestHitShaderFunction()->zSignature()
                       == signature)
            {
                hitGroupsWithThisSignature.add(hitGroup);
            }
            else if (hitGroup->zAnyHitShaderFunction()
                     && hitGroup->zAnyHitShaderFunction()->zSignature()
                            == signature)
            {
                hitGroupsWithThisSignature.add(hitGroup);
            }
            else if (hitGroup->zIntersectionShaderFunction()
                     && hitGroup->zIntersectionShaderFunction()->zSignature()
                            == signature)
            {
                hitGroupsWithThisSignature.add(hitGroup);
            }
        }
        rootSignatureExports[rootSignatureIndex]
            = new const wchar_t*[functionsWithThisSignature.getEntryCount()
                                 + hitGroupsWithThisSignature.getEntryCount()];
        int nameIndex = 0;
        for (const DX12ShaderFunction* function : functionsWithThisSignature)
        {
            rootSignatureExports[rootSignatureIndex][nameIndex]
                = function->zExportDesc()->Name;
            nameIndex++;
        }
        for (const DX12ShaderHitGroup* hitGroup : hitGroupsWithThisSignature)
        {
            rootSignatureExports[rootSignatureIndex][nameIndex]
                = hitGroup->zHitGroupDesc()->HitGroupExport;
            nameIndex++;
        }
        localRootAssociations[rootSignatureIndex].NumExports
            = functionsWithThisSignature.getEntryCount()
            + hitGroupsWithThisSignature.getEntryCount();
        localRootAssociations[rootSignatureIndex].pExports
            = rootSignatureExports[rootSignatureIndex];
        localRootAssociations[rootSignatureIndex].pSubobjectToAssociate
            = &subobjects[index - 1]; // local root signature
        subobjects[index].Type
            = D3D12_STATE_SUBOBJECT_TYPE_SUBOBJECT_TO_EXPORTS_ASSOCIATION;
        subobjects[index].pDesc = &localRootAssociations[rootSignatureIndex];
        index++;

        rootSignatureIndex++;
    }

    D3D12_GLOBAL_ROOT_SIGNATURE globalRootSignature = {};
    globalRootSignature.pGlobalRootSignature = globalSignature->zSignature();
    subobjects[index].Type = D3D12_STATE_SUBOBJECT_TYPE_GLOBAL_ROOT_SIGNATURE;
    subobjects[index].pDesc = &globalRootSignature;
    index++;

    D3D12_RAYTRACING_PIPELINE_CONFIG pipelineConfig = {};
    pipelineConfig.MaxTraceRecursionDepth = maxRecursionDepth;
    subobjects[index].Type
        = D3D12_STATE_SUBOBJECT_TYPE_RAYTRACING_PIPELINE_CONFIG;
    subobjects[index].pDesc = &pipelineConfig;
    index++;

    D3D12_STATE_OBJECT_DESC pipelineDesc = {};
    pipelineDesc.Type = D3D12_STATE_OBJECT_TYPE_RAYTRACING_PIPELINE;
    pipelineDesc.NumSubobjects = subobjectCount;
    pipelineDesc.pSubobjects = subobjects;
    HRESULT hr = zDevice->CreateStateObject(
        &pipelineDesc, __uuidof(ID3D12StateObject), (void**)&pipelineState);
    if (FAILED(hr))
    {
        Logging::error()
            << "Failed to create raytracing pipeline state object: " << std::hex
            << hr;
        throw std::logic_error("Could not create the raytracing state object");
    }

    delete[] functionAndHitGroupNames;
    delete[] localRootSignatures;
    for (int i = 0; i < distinctSignatures.getEntryCount(); i++)
    {
        delete[] rootSignatureExports[i];
    }
    delete[] rootSignatureExports;
    delete[] localRootAssociations;
}

ID3D12StateObject* Framework::DX12Pipeline::zPipelineState() const
{
    return pipelineState;
}

DX12ShaderBindingTable* Framework::DX12Pipeline::createShaderBindingTable()
{
    return new DX12ShaderBindingTable(dynamic_cast<DX12Pipeline*>(getThis()));
}

DX12GlobalDescriptorHeap* Framework::DX12Pipeline::createGlobalDescriptorHeap(
    DX12DescriptorHeapType type)
{
    return new DX12GlobalDescriptorHeap(
        dynamic_cast<DX12Pipeline*>(getThis()), type);
}

const RCArray<DX12Shader>& Framework::DX12Pipeline::getShaders() const
{
    return shaders;
}

const Array<const DX12ShaderFunction*>&
Framework::DX12Pipeline::getFunctionsWithoutHitGroups() const
{
    return functionsWithoutHitGroups;
}

const RCArray<DX12ShaderHitGroup>& Framework::DX12Pipeline::getHitGroups() const
{
    return hitGroups;
}

DX12ShaderSignature* Framework::DX12Pipeline::zGlobalSignature() const
{
    return globalSignature;
}

Framework::DX12GlobalDescriptorHeap::DX12GlobalDescriptorHeap(
    DX12Pipeline* pipeline, DX12DescriptorHeapType type)
    : ReferenceCounter(),
      pipeline(pipeline),
      descriptorHeap(0),
      lastDescriptorHeapSize(0),
      zDevice(0),
      type(type),
      heapChanged(0)
{}

Framework::DX12GlobalDescriptorHeap::~DX12GlobalDescriptorHeap()
{
    pipeline->release();
    if (descriptorHeap)
    {
        descriptorHeap->Release();
    }
    for (const DX12ShaderRegisterInput* input : registerInputs)
    {
        input->inputResource->release();
        delete input;
    }
}

void Framework::DX12GlobalDescriptorHeap::addInput(
    DX12ShaderRegister type, ReferenceCounter* inputResource)
{
    bool found = 0;
    for (DX12Shader* shader : pipeline->getShaders())
    {
        for (DX12ShaderFunction* function : shader->getFunctions())
        {
            for (const DX12ShaderRegisterUsage* usage :
                function->zSignature()->getDescriptorHeapBindings())
            {
                if (usage->descriptorHeapIndex == registerInputs.getEntryCount()
                    && usage->descriptorHeapType == this->type)
                {
                    if (usage->registerType != type)
                    {
                        Logging::error()
                            << "Register type mismatch for register index "
                            << usage->registerIndex << ", space index "
                            << usage->spaceIndex << ". Expected register type: "
                            << usage->registerType
                            << ", given register type: " << type
                            << ". The register type is specified in the "
                               "signature of shader function '"
                            << function->getFunctionName() << "'";
                        throw std::logic_error(
                            "Register type mismatch for shader input");
                    }
                    else
                    {
                        found = 1;
                        break;
                    }
                }
            }
            if (found)
            {
                break;
            }
        }
        if (found)
        {
            break;
        }
    }
    registerInputs.add(new DX12ShaderRegisterInput{
        type, inputResource ? inputResource->getThis() : 0});
}

void Framework::DX12GlobalDescriptorHeap::addTextureInput(
    DX12ShaderRegister type, Texture* zTexture)
{
    addInput(type, zTexture);
}

void Framework::DX12GlobalDescriptorHeap::updateTextureInput(
    int heapIndex, DX12ShaderRegister type, Texture* zTexture)
{
    DX12ShaderRegisterInput* input = registerInputs.get(heapIndex);
    if (registerInputs.get(heapIndex)->inputResource
            != dynamic_cast<ReferenceCounter*>(zTexture)
        || zTexture->hasBufferChanged())
    {
        zTexture->setBufferChanged(0);
        if (input->registerType != type)
        {
            Logging::error()
                << "Register type mismatch for descriptor heap index "
                << heapIndex
                << ". Expected register type: " << input->registerType
                << ", given register type: " << type << ".";
            throw std::logic_error("Register type mismatch in descriptor heap");
        }
        registerInputs.get(heapIndex)->inputResource->release();
        registerInputs.get(heapIndex)->inputResource = zTexture->getThis();
        if (descriptorHeap)
        {
            D3D12_CPU_DESCRIPTOR_HANDLE descriptorHeapHandle
                = descriptorHeap->GetCPUDescriptorHandleForHeapStart();
            descriptorHeapHandle.ptr
                += zDevice->GetDescriptorHandleIncrementSize(
                       D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV)
                 * heapIndex;
            switch (input->registerType)
            {
            case DX12_SHADER_REGISTER_T_SHADER_RESOURCE:
                {
                    D3D12_SHADER_RESOURCE_VIEW_DESC srvDesc;
                    srvDesc.Format = DXGI_FORMAT_UNKNOWN;
                    srvDesc.ViewDimension = D3D12_SRV_DIMENSION_TEXTURE2D;
                    srvDesc.Texture2D.MipLevels = 0;
                    srvDesc.Texture2D.MostDetailedMip = 0;
                    srvDesc.Texture2D.PlaneSlice = 0;
                    srvDesc.Texture2D.ResourceMinLODClamp = 0.0f;
                    srvDesc.Shader4ComponentMapping
                        = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING;
                    zDevice->CreateShaderResourceView(
                        ((DX12Texture*)zTexture)->zResource(),
                        &srvDesc,
                        descriptorHeapHandle);
                    break;
                }
            case DX12_SHADER_REGISTER_U_UNORDERED_ACCESS:
                {
                    D3D12_UNORDERED_ACCESS_VIEW_DESC uavDesc = {};
                    uavDesc.ViewDimension = D3D12_UAV_DIMENSION_TEXTURE2D;
                    uavDesc.Format = DXGI_FORMAT_UNKNOWN;
                    uavDesc.Texture2D.MipSlice = 0;
                    uavDesc.Texture2D.PlaneSlice = 0;
                    zDevice->CreateUnorderedAccessView(
                        ((DX12Texture*)zTexture)->zResource(),
                        0,
                        &uavDesc,
                        descriptorHeapHandle);
                    break;
                }
            }
        }
    }
}

void Framework::DX12GlobalDescriptorHeap::addBufferInput(
    DX12ShaderRegister type, DXBuffer* zBuffer)
{
    addInput(type, zBuffer);
}

void Framework::DX12GlobalDescriptorHeap::addTLASInput(
    DX12ShaderRegister type, DX12TLAS* zTLAS)
{
    addInput(type, zTLAS);
}

void Framework::DX12GlobalDescriptorHeap::updateTLASInput(
    int heapIndex, DX12ShaderRegister type, DX12TLAS* zTLAS)
{
    DX12ShaderRegisterInput* input = registerInputs.get(heapIndex);
    if (registerInputs.get(heapIndex)->inputResource
            != dynamic_cast<ReferenceCounter*>(zTLAS)
        || zTLAS->hasBufferChanged())
    {
        zTLAS->setBufferChanged(0);
        if (input->registerType != type)
        {
            Logging::error()
                << "Register type mismatch for descriptor heap index "
                << heapIndex
                << ". Expected register type: " << input->registerType
                << ", given register type: " << type << ".";
            throw std::logic_error("Register type mismatch in descriptor heap");
        }
        if (registerInputs.get(heapIndex)->inputResource)
        {
            registerInputs.get(heapIndex)->inputResource->release();
        }
        registerInputs.get(heapIndex)->inputResource = zTLAS->getThis();
        if (descriptorHeap)
        {
            D3D12_CPU_DESCRIPTOR_HANDLE descriptorHeapHandle
                = descriptorHeap->GetCPUDescriptorHandleForHeapStart();
            descriptorHeapHandle.ptr
                += zDevice->GetDescriptorHandleIncrementSize(
                       D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV)
                 * heapIndex;
            switch (input->registerType)
            {
            case DX12_SHADER_REGISTER_T_SHADER_RESOURCE:
                {
                    D3D12_SHADER_RESOURCE_VIEW_DESC srvDesc;
                    srvDesc.Format = DXGI_FORMAT_UNKNOWN;
                    srvDesc.ViewDimension
                        = D3D12_SRV_DIMENSION_RAYTRACING_ACCELERATION_STRUCTURE;
                    srvDesc.RaytracingAccelerationStructure.Location
                        = zTLAS->zResultBuffer()
                              ->zBuffer()
                              ->GetGPUVirtualAddress();
                    srvDesc.Shader4ComponentMapping
                        = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING;
                    zDevice->CreateShaderResourceView(
                        0, &srvDesc, descriptorHeapHandle);
                    break;
                }
            }
        }
    }
}

void Framework::DX12GlobalDescriptorHeap::addSamplerInput(
    DX12SamplerState* zSampler)
{
    addInput(DX12_SHADER_REGISTER_S_SAMPLER, zSampler);
}

void Framework::DX12GlobalDescriptorHeap::updateDescriptorHeap(
    ID3D12Device5* zDevice)
{
    this->zDevice = zDevice;
    if (!descriptorHeap
        || lastDescriptorHeapSize != registerInputs.getEntryCount())
    {
        if (descriptorHeap)
        {
            descriptorHeap->Release();
        }
        D3D12_DESCRIPTOR_HEAP_DESC desc = {};
        desc.NumDescriptors = registerInputs.getEntryCount();
        desc.Type = type == SAMPLER_DESCRIPTOR_HEAP
                      ? D3D12_DESCRIPTOR_HEAP_TYPE_SAMPLER
                      : D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV;
        desc.Flags = D3D12_DESCRIPTOR_HEAP_FLAG_SHADER_VISIBLE;
        desc.NodeMask = 0;

        HRESULT r = zDevice->CreateDescriptorHeap(
            &desc, __uuidof(ID3D12DescriptorHeap), (void**)&descriptorHeap);
        lastDescriptorHeapSize = registerInputs.getEntryCount();
        heapChanged = 1;
    }

    D3D12_CPU_DESCRIPTOR_HANDLE descriptorHeapHandle
        = descriptorHeap->GetCPUDescriptorHandleForHeapStart();
    for (const DX12ShaderRegisterInput* input : registerInputs)
    {
        DX12TLAS* zTLAS = dynamic_cast<DX12TLAS*>(input->inputResource);
        DX12Texture* zTexture
            = dynamic_cast<DX12Texture*>(input->inputResource);
        if (zTexture && !zTexture->zResource())
        {
            zTexture = 0;
        }
        DX12Buffer* zBuffer = dynamic_cast<DX12Buffer*>(input->inputResource);
        DX12SamplerState* zSampler
            = dynamic_cast<DX12SamplerState*>(input->inputResource);
        switch (input->registerType)
        {
        case DX12_SHADER_REGISTER_B_CONST_BUFFER:
            {
                D3D12_CONSTANT_BUFFER_VIEW_DESC cbvDesc = {};
                if (!zBuffer)
                {
                    Logging::error()
                        << "Expected a buffer resource for register type "
                        << input->registerType;
                    throw std::logic_error(
                        "Expected a buffer resource for register type "
                        + std::to_string(input->registerType));
                }
                cbvDesc.BufferLocation
                    = zBuffer->zBuffer()->GetGPUVirtualAddress();
                cbvDesc.SizeInBytes = (unsigned)zBuffer->getElementCount()
                                    * zBuffer->getElementLength();
                zDevice->CreateConstantBufferView(
                    &cbvDesc, descriptorHeapHandle);
                break;
            }
        case DX12_SHADER_REGISTER_T_SHADER_RESOURCE:
            {
                D3D12_SHADER_RESOURCE_VIEW_DESC srvDesc;
                srvDesc.Format = DXGI_FORMAT_UNKNOWN;
                bool doNothing = 0;
                if (zTLAS)
                {
                    srvDesc.ViewDimension
                        = D3D12_SRV_DIMENSION_RAYTRACING_ACCELERATION_STRUCTURE;
                    srvDesc.RaytracingAccelerationStructure.Location
                        = zTLAS->zResultBuffer()
                              ->zBuffer()
                              ->GetGPUVirtualAddress();
                }
                else if (zTexture)
                {
                    srvDesc.ViewDimension = D3D12_SRV_DIMENSION_TEXTURE2D;
                    srvDesc.Texture2D.MipLevels = 1;
                    srvDesc.Texture2D.MostDetailedMip = 0;
                    srvDesc.Texture2D.PlaneSlice = 0;
                    srvDesc.Texture2D.ResourceMinLODClamp = 0.0f;
                }
                else if (zBuffer)
                {
                    srvDesc.ViewDimension = D3D12_SRV_DIMENSION_BUFFER;
                    srvDesc.Buffer.FirstElement = 0;
                    srvDesc.Buffer.NumElements
                        = (unsigned)zBuffer->getElementCount();
                    srvDesc.Buffer.StructureByteStride
                        = zBuffer->getElementLength();
                    srvDesc.Buffer.Flags = D3D12_BUFFER_SRV_FLAG_NONE;
                }
                else if (zSampler)
                {
                    Logging::error()
                        << "Expected a texture or buffer or TLAS resource for "
                           "register type "
                        << input->registerType;
                    throw std::logic_error(
                        "Expected a texture or buffer or TLAS resource for "
                        "register type "
                        + std::to_string(input->registerType));
                }
                else
                {
                    doNothing = 1;
                }
                if (!doNothing)
                {
                    srvDesc.Shader4ComponentMapping
                        = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING;
                    zDevice->CreateShaderResourceView(
                        zTexture ? zTexture->zResource()
                                 : (zBuffer ? zBuffer->zBuffer() : 0),
                        &srvDesc,
                        descriptorHeapHandle);
                }
                break;
            }
        case DX12_SHADER_REGISTER_U_UNORDERED_ACCESS:
            {
                D3D12_UNORDERED_ACCESS_VIEW_DESC uavDesc = {};
                bool doNothing = 0;
                if (zTexture)
                {
                    uavDesc.ViewDimension = D3D12_UAV_DIMENSION_TEXTURE2D;
                    uavDesc.Format = DXGI_FORMAT_UNKNOWN;
                    uavDesc.Texture2D.MipSlice = 0;
                    uavDesc.Texture2D.PlaneSlice = 0;
                }
                else if (zBuffer)
                {
                    uavDesc.ViewDimension = D3D12_UAV_DIMENSION_BUFFER;
                    uavDesc.Buffer.FirstElement = 0;
                    uavDesc.Buffer.NumElements
                        = (unsigned)zBuffer->getElementCount();
                    uavDesc.Buffer.StructureByteStride
                        = zBuffer->getElementLength();
                    uavDesc.Buffer.CounterOffsetInBytes = 0;
                    uavDesc.Buffer.Flags = D3D12_BUFFER_UAV_FLAG_NONE;
                }
                else if (zTLAS || zSampler)
                {
                    Logging::error()
                        << "Expected a texture or buffer resource for "
                           "register type "
                        << input->registerType;
                    throw std::logic_error(
                        "Expected a texture or buffer resource for "
                        "register "
                        "type "
                        + std::to_string(input->registerType));
                }
                else
                {
                    doNothing = 1;
                }
                if (!doNothing)
                {
                    zDevice->CreateUnorderedAccessView(
                        zTexture ? zTexture->zResource() : zBuffer->zBuffer(),
                        0,
                        &uavDesc,
                        descriptorHeapHandle);
                }
                break;
            }
        case DX12_SHADER_REGISTER_S_SAMPLER:
            if (zSampler)
            {
                zDevice->CreateSampler(
                    zSampler->zSamplerDesc(), descriptorHeapHandle);
            }
            else
            {
                Logging::error()
                    << "Expected a sampler resource for register type "
                    << input->registerType;
                throw std::logic_error(
                    "Expected a sampler resource for register type "
                    + std::to_string(input->registerType));
            }
        }
        descriptorHeapHandle.ptr += zDevice->GetDescriptorHandleIncrementSize(
            type == SAMPLER_DESCRIPTOR_HEAP
                ? D3D12_DESCRIPTOR_HEAP_TYPE_SAMPLER
                : D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV);
    }
}

DX12Pipeline* Framework::DX12GlobalDescriptorHeap::zPipeline() const
{
    return pipeline;
}

ID3D12DescriptorHeap*
Framework::DX12GlobalDescriptorHeap::zDescriptorHeap() const
{
    return descriptorHeap;
}

void Framework::DX12GlobalDescriptorHeap::setHeapChanged(bool changed)
{
    heapChanged = changed;
}

bool Framework::DX12GlobalDescriptorHeap::wasHeapChanged() const
{
    return heapChanged;
}

int Framework::DX12GlobalDescriptorHeap::getRegisterInputCount() const
{
    return registerInputs.getEntryCount();
}

Framework::DX12ShaderBindingTable::DX12ShaderBindingTable(
    DX12Pipeline* pipeline)
    : ReferenceCounter(),
      pipeline(pipeline),
      shaderBindingTableBuffer(0),
      globalDescriptorHeap(0),
      samplerDescriptorHeap(0),
      rayGenRecordSize(0),
      rayGenCount(0),
      missRecordSize(0),
      missCount(0),
      callableRecordSize(0),
      callableCount(0),
      hitGroupRecordSize(0),
      hitGroupCount(0),
      tableBuffer(0),
      tableBufferSize(0),
      nextHitGroupOffset(0),
      stateObjectProperties(0)
{
    pipeline->zPipelineState()->QueryInterface(
        __uuidof(ID3D12StateObjectProperties), (void**)&stateObjectProperties);
}

Framework::DX12ShaderBindingTable::~DX12ShaderBindingTable()
{
    stateObjectProperties->Release();
    if (pipeline)
    {
        pipeline->release();
    }
    if (shaderBindingTableBuffer)
    {
        shaderBindingTableBuffer->release();
    }
    if (globalDescriptorHeap)
    {
        globalDescriptorHeap->release();
    }
    if (samplerDescriptorHeap)
    {
        samplerDescriptorHeap->release();
    }
    for (const char* buffer : tempBuffers)
    {
        delete[] buffer;
    }
}

void Framework::DX12ShaderBindingTable::setGlobalDescriptorHeap(
    DX12GlobalDescriptorHeap* zGlobalDescriptorHeap)
{
    if (this->globalDescriptorHeap != zGlobalDescriptorHeap)
    {
        if (this->globalDescriptorHeap)
        {
            this->globalDescriptorHeap->release();
        }
        this->globalDescriptorHeap = zGlobalDescriptorHeap;
        if (this->globalDescriptorHeap)
        {
            this->globalDescriptorHeap->getThis();
        }
    }
}

void Framework::DX12ShaderBindingTable::setSamplerDescriptorHeap(
    DX12GlobalDescriptorHeap* zSamplerDescriptorHeap)
{
    if (this->samplerDescriptorHeap != zSamplerDescriptorHeap)
    {
        if (this->samplerDescriptorHeap)
        {
            this->samplerDescriptorHeap->release();
        }
        this->samplerDescriptorHeap = zSamplerDescriptorHeap;
        if (this->samplerDescriptorHeap)
        {
            this->samplerDescriptorHeap->getThis();
        }
    }
}

void Framework::DX12ShaderBindingTable::startUpdate()
{
    if (shaderBindingTableBuffer)
    {
        shaderBindingTableBuffer->zBuffer()->Map(0, 0, (void**)&tableBuffer);
        tableBufferSize = (int)shaderBindingTableBuffer->getElementCount()
                        * shaderBindingTableBuffer->getElementLength();
    }
    else
    {
        tableBuffer = 0;
        tableBufferSize = 0;
    }
    rayGenRecordSize = 0;
    rayGenCount = 0;
    missRecordSize = 0;
    missCount = 0;
    callableRecordSize = 0;
    callableCount = 0;
    for (const DX12ShaderFunction* function :
        pipeline->getFunctionsWithoutHitGroups())
    {
        DX12ShaderSignature* signature = function->zSignature();
        if (function->getFunctionType() == DX12_SHADER_FUNCTION_TYPE_RAY_GEN)
        {
            if (rayGenRecordSize < signature->gerShaderBindingTableParamCount())
            {
                rayGenRecordSize = signature->gerShaderBindingTableParamCount();
            }
            rayGenCount++;
        }
        else if (function->getFunctionType() == DX12_SHADER_FUNCTION_TYPE_MISS)
        {
            if (missRecordSize < signature->gerShaderBindingTableParamCount())
            {
                missRecordSize = signature->gerShaderBindingTableParamCount();
            }
            missCount++;
        }
        else if (function->getFunctionType()
                 == DX12_SHADER_FUNCTION_TYPE_CALLABLE)
        {
            if (callableRecordSize
                < signature->gerShaderBindingTableParamCount())
            {
                callableRecordSize
                    = signature->gerShaderBindingTableParamCount();
            }
            callableCount++;
        }
    }
    rayGenRecordSize = ROUND_UP_POWER_OF_2(
        D3D12_RAYTRACING_SHADER_RECORD_BYTE_ALIGNMENT + rayGenRecordSize * 8,
        D3D12_RAYTRACING_SHADER_RECORD_BYTE_ALIGNMENT);
    missRecordSize = ROUND_UP_POWER_OF_2(
        D3D12_RAYTRACING_SHADER_RECORD_BYTE_ALIGNMENT + missRecordSize * 8,
        D3D12_RAYTRACING_SHADER_RECORD_BYTE_ALIGNMENT);
    callableRecordSize = ROUND_UP_POWER_OF_2(
        D3D12_RAYTRACING_SHADER_RECORD_BYTE_ALIGNMENT + callableRecordSize * 8,
        D3D12_RAYTRACING_SHADER_RECORD_BYTE_ALIGNMENT);
    hitGroupCount = 0;
    hitGroupRecordSize = 0;
    for (DX12ShaderHitGroup* hitGroup : pipeline->getHitGroups())
    {
        DX12ShaderSignature* signature = hitGroup->zSignature();
        if (hitGroupRecordSize < signature->gerShaderBindingTableParamCount())
        {
            hitGroupRecordSize = signature->gerShaderBindingTableParamCount();
        }
    }
    hitGroupRecordSize = ROUND_UP_POWER_OF_2(
        D3D12_RAYTRACING_SHADER_RECORD_BYTE_ALIGNMENT + hitGroupRecordSize * 8,
        D3D12_RAYTRACING_SHADER_RECORD_BYTE_ALIGNMENT);
    nextHitGroupOffset
        = ROUND_UP_POWER_OF_2(rayGenRecordSize * rayGenCount, 64)
        + ROUND_UP_POWER_OF_2(missRecordSize * missCount, 64)
        + (callableCount > 0
                ? ROUND_UP_POWER_OF_2(callableRecordSize * callableCount, 64)
                : 0);
}

void Framework::DX12ShaderBindingTable::set(int index, void* data, int size)
{
    if (tableBufferSize >= index + size)
    {
        memcpy(tableBuffer + index, data, size);
    }
    else
    {
        index -= tableBufferSize;
        ArrayIterator<char*> it = tempBuffers.begin();
        while (it && 2048 < index + size)
        {
            index -= 2048;
            it++;
        }
        while (2048 < index + size)
        {
            char* newBuffer = new char[2048];
            memset(newBuffer, 0, 2048);
            tempBuffers.add(newBuffer);
            index -= 2048;
        }
        // because the buffer size is allways rounded up to multiples
        // of 32 and the size written at once is always 8 or 32 it
        // should never be possible to have a negative index here
        assert(index >= 0);
        if (!it)
        {
            char* newBuffer = new char[2048];
            memset(newBuffer, 0, 2048);
            tempBuffers.add(newBuffer);
            memcpy(newBuffer + index, data, size);
        }
        else
        {
            memcpy(it.val() + index, data, size);
        }
    }
}

void Framework::DX12ShaderBindingTable::setShaderInput(
    DX12ShaderFunction* zFunction, int* offsetPointer, __int64 gpuAddress)
{
    int index = 0;
    for (const DX12ShaderFunction* pf :
        pipeline->getFunctionsWithoutHitGroups())
    {
        if (pf == zFunction)
        {
            break;
        }
        if (pf->getFunctionType() == zFunction->getFunctionType())
        {
            index++;
        }
    }
    int offset = 0;
    if (zFunction->getFunctionType() == DX12_SHADER_FUNCTION_TYPE_RAY_GEN)
    {
        offset = index * rayGenRecordSize;
    }
    else if (zFunction->getFunctionType() == DX12_SHADER_FUNCTION_TYPE_MISS)
    {
        offset = ROUND_UP_POWER_OF_2(rayGenRecordSize * rayGenCount, 64)
               + index * missRecordSize;
    }
    else if (zFunction->getFunctionType() == DX12_SHADER_FUNCTION_TYPE_CALLABLE)
    {
        offset = ROUND_UP_POWER_OF_2(rayGenRecordSize * rayGenCount, 64)
               + ROUND_UP_POWER_OF_2(missRecordSize * missCount, 64)
               + index * callableRecordSize;
    }
    else
    {
        throw std::logic_error("setShaderInput can only be used for ray "
                               "generation, miss, and callable "
                               "shader functions");
    }
    offset += D3D12_RAYTRACING_SHADER_RECORD_BYTE_ALIGNMENT
            + sizeof(__int64) * *offsetPointer;
    set(offset, &gpuAddress, sizeof(__int64));
}

int Framework::DX12ShaderBindingTable::addHitGroup(
    DX12ShaderHitGroup* zHitGroup, int lastIndex)
{
    int index = nextHitGroupOffset;
    bool changed
        = index + D3D12_RAYTRACING_SHADER_RECORD_BYTE_ALIGNMENT != lastIndex;
    if (changed)
    {
        set(index,
            stateObjectProperties->GetShaderIdentifier(
                zHitGroup->zHitGroupDesc()->HitGroupExport),
            D3D12_RAYTRACING_SHADER_RECORD_BYTE_ALIGNMENT);
    }
    index += D3D12_RAYTRACING_SHADER_RECORD_BYTE_ALIGNMENT;
    if (zHitGroup->zSignature()->doesUseGlobalDescriptorHeap())
    {
        if ((changed || globalDescriptorHeap->wasHeapChanged()))
        {
            D3D12_GPU_DESCRIPTOR_HANDLE gpuAddress
                = globalDescriptorHeap->zDescriptorHeap()
                      ->GetGPUDescriptorHandleForHeapStart();
            set(index, &gpuAddress.ptr, sizeof(__int64));
        }
        index += sizeof(__int64);
    }
    if (zHitGroup->zSignature()->doesUseSamplerDescriptorHeap())
    {
        if ((changed || samplerDescriptorHeap->wasHeapChanged()))
        {
            D3D12_GPU_DESCRIPTOR_HANDLE gpuAddress
                = samplerDescriptorHeap->zDescriptorHeap()
                      ->GetGPUDescriptorHandleForHeapStart();
            set(index, &gpuAddress.ptr, sizeof(__int64));
        }
        index += sizeof(__int64);
    }
    hitGroupCount++;
    nextHitGroupOffset += hitGroupRecordSize;
    return nextHitGroupOffset - hitGroupRecordSize
         + D3D12_RAYTRACING_SHADER_RECORD_BYTE_ALIGNMENT;
}

void Framework::DX12ShaderBindingTable::setHitGroupShaderInput(
    int hitGroupOffset, int* offsetPointer, __int64 gpuAddress)
{
    set(hitGroupOffset + sizeof(__int64) * *offsetPointer,
        &gpuAddress,
        sizeof(__int64));
}

void Framework::DX12ShaderBindingTable::endUpdate(
    ID3D12Device5* zDevice, DX12CommandQueue* zQueue)
{
    int rayTracingIndex = 0;
    int missIndex = 0;
    int callableIndex = 0;
    for (const DX12ShaderFunction* function :
        pipeline->getFunctionsWithoutHitGroups())
    {
        int offset = -1;
        DX12ShaderSignature* signature = function->zSignature();
        if (function->getFunctionType() == DX12_SHADER_FUNCTION_TYPE_RAY_GEN)
        {
            offset = rayTracingIndex * rayGenRecordSize;
            rayTracingIndex++;
        }
        else if (function->getFunctionType() == DX12_SHADER_FUNCTION_TYPE_MISS)
        {
            offset = ROUND_UP_POWER_OF_2(rayGenRecordSize * rayGenCount, 64)
                   + missIndex * missRecordSize;
            missIndex++;
        }
        else if (function->getFunctionType()
                 == DX12_SHADER_FUNCTION_TYPE_CALLABLE)
        {
            offset = ROUND_UP_POWER_OF_2(rayGenRecordSize * rayGenCount, 64)
                   + ROUND_UP_POWER_OF_2(missRecordSize * missCount, 64)
                   + callableIndex * callableRecordSize;
            callableIndex++;
        }
        if (offset >= 0)
        {
            set(offset,
                stateObjectProperties->GetShaderIdentifier(
                    function->zExportDesc()->Name),
                D3D12_RAYTRACING_SHADER_RECORD_BYTE_ALIGNMENT);
            offset += D3D12_RAYTRACING_SHADER_RECORD_BYTE_ALIGNMENT;
            if (function->zSignature()->doesUseGlobalDescriptorHeap())
            {
                D3D12_GPU_DESCRIPTOR_HANDLE gpuAddress
                    = globalDescriptorHeap->zDescriptorHeap()
                          ->GetGPUDescriptorHandleForHeapStart();
                set(offset, &gpuAddress.ptr, sizeof(__int64));
                offset += sizeof(__int64);
            }
            if (function->zSignature()->doesUseSamplerDescriptorHeap())
            {
                D3D12_GPU_DESCRIPTOR_HANDLE gpuAddress
                    = samplerDescriptorHeap->zDescriptorHeap()
                          ->GetGPUDescriptorHandleForHeapStart();
                set(offset, &gpuAddress.ptr, sizeof(__int64));
                offset += sizeof(__int64);
            }
        }
    }
    if (nextHitGroupOffset > tableBufferSize)
    {
        DX12Buffer* newBuffer = new DX12Buffer(1,
            zDevice,
            dynamic_cast<DX12CommandQueue*>(zQueue->getThis()),
            D3D12_RESOURCE_FLAG_NONE);
        newBuffer->setLength(ROUND_UP_POWER_OF_2(nextHitGroupOffset, 256));
        newBuffer->createBufferWithoutData(
            D3D12_RESOURCE_STATE_GENERIC_READ, D3D12_HEAP_TYPE_UPLOAD);
        void* newTableBuffer = 0;
        newBuffer->zBuffer()->Map(0, 0, (void**)&newTableBuffer);
        int sizeToCopy = nextHitGroupOffset;
        int index = 0;
        if (tableBufferSize > 0)
        {
            memcpy(newTableBuffer, tableBuffer, tableBufferSize);
            sizeToCopy -= tableBufferSize;
            index = tableBufferSize;
        }
        for (char* buffer : tempBuffers)
        {
            int bytesToCopy = sizeToCopy < 2048 ? sizeToCopy : 2048;
            memcpy((char*)newTableBuffer + index, buffer, bytesToCopy);
            sizeToCopy -= bytesToCopy;
            index += bytesToCopy;
            if (sizeToCopy == 0)
            {
                break;
            }
        }
        assert(sizeToCopy == 0);
        newBuffer->zBuffer()->Unmap(0, 0);
        if (shaderBindingTableBuffer)
        {
            shaderBindingTableBuffer->zBuffer()->Unmap(0, 0);
            shaderBindingTableBuffer->release();
        }
        shaderBindingTableBuffer = newBuffer;
    }
    else
    {
        shaderBindingTableBuffer->zBuffer()->Unmap(0, 0);
    }
}

void Framework::DX12ShaderBindingTable::fillDispatchRaysDesc(
    D3D12_DISPATCH_RAYS_DESC* dispatchRaysDesc)
{
    dispatchRaysDesc->RayGenerationShaderRecord.SizeInBytes
        = rayGenCount * rayGenRecordSize;
    dispatchRaysDesc->RayGenerationShaderRecord.StartAddress
        = shaderBindingTableBuffer->zBuffer()->GetGPUVirtualAddress();

    dispatchRaysDesc->MissShaderTable.SizeInBytes = missCount * missRecordSize;
    dispatchRaysDesc->MissShaderTable.StartAddress
        = shaderBindingTableBuffer->zBuffer()->GetGPUVirtualAddress()
        + ROUND_UP_POWER_OF_2(rayGenCount * rayGenRecordSize, 64);
    dispatchRaysDesc->MissShaderTable.StrideInBytes = missRecordSize;

    dispatchRaysDesc->HitGroupTable.SizeInBytes
        = hitGroupCount * hitGroupRecordSize;
    dispatchRaysDesc->HitGroupTable.StartAddress
        = shaderBindingTableBuffer->zBuffer()->GetGPUVirtualAddress()
        + ROUND_UP_POWER_OF_2(rayGenCount * rayGenRecordSize, 64)
        + ROUND_UP_POWER_OF_2(missCount * missRecordSize, 64)
        + (callableCount > 0
                ? ROUND_UP_POWER_OF_2(callableCount * callableRecordSize, 64)
                : 0);
    dispatchRaysDesc->HitGroupTable.StrideInBytes = hitGroupRecordSize;

    dispatchRaysDesc->CallableShaderTable.SizeInBytes
        = callableCount * callableRecordSize;
    dispatchRaysDesc->CallableShaderTable.StartAddress
        = callableCount > 0
            ? shaderBindingTableBuffer->zBuffer()->GetGPUVirtualAddress()
                  + ROUND_UP_POWER_OF_2(rayGenCount * rayGenRecordSize, 64)
                  + ROUND_UP_POWER_OF_2(missCount * missRecordSize, 64)
            : 0;
    dispatchRaysDesc->CallableShaderTable.StrideInBytes
        = callableCount > 0 ? callableRecordSize : 0;
}

DX12Pipeline* Framework::DX12ShaderBindingTable::zPipeline() const
{
    return pipeline;
}

int Framework::DX12ShaderBindingTable::getHitGoupCount() const
{
    return hitGroupCount;
}