#include "DX12Shader.h"

#include "Logging.h"

using namespace Framework;

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

Framework::DX12ShaderSignature::~DX12ShaderSignature()
{
    if (signature)
    {
        signature->Release();
    }
}

void Framework::DX12ShaderSignature::addRegisterUsage(
    DX12ShaderRegister registerType, int registerIndex)
{
    addRegisterUsage(registerType, registerIndex, 0);
}

void Framework::DX12ShaderSignature::addRegisterUsage(
    DX12ShaderRegister registerType, int registerIndex, int spaceIndex)
{
    registerUsages.add({registerType, registerIndex, spaceIndex});
    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;
    }
    D3D12_ROOT_PARAMETER descriptorTable;
    descriptorTable.ParameterType = D3D12_ROOT_PARAMETER_TYPE_DESCRIPTOR_TABLE;
    descriptorTable.ShaderVisibility = D3D12_SHADER_VISIBILITY_ALL;
    descriptorTable.DescriptorTable.NumDescriptorRanges
        = registerUsages.getEntryCount();
    D3D12_DESCRIPTOR_RANGE* descriptorRanges
        = new D3D12_DESCRIPTOR_RANGE[registerUsages.getEntryCount()];
    int index = 0;
    for (const auto& usage : registerUsages)
    {
        D3D12_DESCRIPTOR_RANGE* range = &descriptorRanges[index];
        switch (usage.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;
        default:
            Logging::error() << "Unknown register type for root signature: "
                             << usage.registerType;
            range->RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_SRV;
        }
        range->NumDescriptors
            = 1; // TODO: optimize when multiple descriptors are used in a row
                 // e.g. u0, u1, u2 with the same range then NumDescriptors can
                 // be 3 for u0 and u1 and u2 would not be needed in this array
        range->BaseShaderRegister = usage.registerIndex;
        range->RegisterSpace = usage.spaceIndex;
        range->OffsetInDescriptorsFromTableStart = index;
        index++;
    }
    descriptorTable.DescriptorTable.pDescriptorRanges = descriptorRanges;
    D3D12_ROOT_SIGNATURE_DESC rootDesc = {};
    rootDesc.NumParameters = 1;
    rootDesc.pParameters = &descriptorTable;
    rootDesc.Flags = 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;
}

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

Framework::DX12ShaderFunction::DX12ShaderFunction(
    const Text& functionName, DX12ShaderSignature* signature)
    : ReferenceCounter(),
      functionName(functionName),
      signature(signature),
      exportDesc(new D3D12_EXPORT_DESC())
{
    wchar_t* wc = new wchar_t[functionName.getLength() + 1];
    mbtowc(wc, 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;
}

Framework::DX12Shader::DX12Shader(const 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 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];
    mbtowc(wc, 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* closestHitShaderFunction)
{
    if (anyHitShaderFunction
        && anyHitShaderFunction->zSignature()
               != closestHitShaderFunction->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: '"
            << closestHitShaderFunction->getFunctionName().getText() << "'";
        throw std::runtime_error("Incompatible root signatures in hit group");
    }
    if (intersectionShaderFunction
        && intersectionShaderFunction->zSignature()
               != closestHitShaderFunction->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: '"
            << closestHitShaderFunction->getFunctionName().getText() << "'";
        throw std::runtime_error("Incompatible root signatures in hit group");
    }

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

void Framework::DX12ShaderHitGroup::setAnyHitShaderFunction(
    DX12ShaderFunction* anyHitShaderFunction)
{
    if (closestHitShaderFunction
        && closestHitShaderFunction->zSignature()
               != anyHitShaderFunction->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: '"
            << closestHitShaderFunction->getFunctionName().getText() << "'";
        throw std::runtime_error("Incompatible root signatures in hit group");
    }
    if (intersectionShaderFunction
        && intersectionShaderFunction->zSignature()
               != anyHitShaderFunction->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: '"
            << anyHitShaderFunction->getFunctionName().getText() << "'";
        throw std::runtime_error("Incompatible root signatures in hit group");
    }

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

void Framework::DX12ShaderHitGroup::setIntersectionShaderFunction(
    DX12ShaderFunction* intersectionShaderFunction)
{
    if (closestHitShaderFunction
        && closestHitShaderFunction->zSignature()
               != intersectionShaderFunction->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: '"
            << closestHitShaderFunction->getFunctionName().getText() << "'";
        throw std::runtime_error("Incompatible root signatures in hit group");
    }
    if (anyHitShaderFunction
        && anyHitShaderFunction->zSignature()
               != intersectionShaderFunction->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: '"
            << anyHitShaderFunction->getFunctionName().getText() << "'";
        throw std::runtime_error("Incompatible root signatures in hit group");
    }

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

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;
}

Framework::DX12Pipeline::DX12Pipeline()
    : ReferenceCounter(),
      emptyGlobalRootSignature(0),
      emptyLocalRootSignature(0),
      pipelineState(0),
      maxRecursionDepth(0)
{}

Framework::DX12Pipeline::~DX12Pipeline()
{
    if (emptyGlobalRootSignature)
    {
        emptyGlobalRootSignature->Release();
    }
    if (emptyLocalRootSignature)
    {
        emptyLocalRootSignature->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)
{
    if (!emptyGlobalRootSignature)
    {
        D3D12_ROOT_SIGNATURE_DESC rootDesc = {};
        rootDesc.NumParameters = 0;
        rootDesc.pParameters = 0;
        rootDesc.Flags = D3D12_ROOT_SIGNATURE_FLAG_NONE;
        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**)&emptyGlobalRootSignature);
            pSigBlob->Release();
        }
        if (pErrorBlob)
        {
            std::string errorMessage(
                static_cast<const char*>(pErrorBlob->GetBufferPointer()),
                pErrorBlob->GetBufferSize());
            Logging::error()
                << "Failed to serialize empty root signature: " << errorMessage;
            pErrorBlob->Release();
        }
    }
    if (!emptyLocalRootSignature)
    {
        D3D12_ROOT_SIGNATURE_DESC rootDesc = {};
        rootDesc.NumParameters = 0;
        rootDesc.pParameters = 0;
        rootDesc.Flags = 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**)&emptyGlobalRootSignature);
            pSigBlob->Release();
        }
        if (pErrorBlob)
        {
            std::string errorMessage(
                static_cast<const char*>(pErrorBlob->GetBufferPointer()),
                pErrorBlob->GetBufferSize());
            Logging::error()
                << "Failed to serialize empty root signature: " << errorMessage;
            pErrorBlob->Release();
        }
    }
    unsigned int subobjectCount
        = shaders.getEntryCount() + hitGroups.getEntryCount() + 5;
    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++;
    // create export list for hole pipeline with all hitgroups and all functions
    // that are not part of a hit group
    Array<const DX12ShaderFunction*> functionsWithoutHitGroups;
    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);
            }
        }
    }
    wchar_t** functionAndHitGroupNames
        = new 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++;
    }

    subobjects[index].Type = D3D12_STATE_SUBOBJECT_TYPE_GLOBAL_ROOT_SIGNATURE;
    subobjects[index].pDesc = &emptyGlobalRootSignature;
    index++;

    subobjects[index].Type = D3D12_STATE_SUBOBJECT_TYPE_LOCAL_ROOT_SIGNATURE;
    subobjects[index].pDesc = &emptyLocalRootSignature;
    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;
}