XCEngine/engine/src/Rendering/Pipelines/BuiltinForwardPipeline.cpp

#include "Rendering/Pipelines/BuiltinForwardPipeline.h"

#include "Components/MeshFilterComponent.h"
#include "Components/MeshRendererComponent.h"
#include "RHI/RHICommandList.h"
#include "Rendering/RenderMaterialUtility.h"
#include "Rendering/RenderSurface.h"
#include "Resources/Material/Material.h"
#include "Resources/Texture/Texture.h"

#include <cstddef>
#include <cstring>

namespace XCEngine {
namespace Rendering {
namespace Pipelines {

namespace {

constexpr uint32_t kDescriptorFirstSet = 1;
constexpr uint32_t kDescriptorSetCount = 4;

const char kBuiltinForwardHlsl[] = R"(
Texture2D gBaseColorTexture : register(t1);
SamplerState gLinearSampler : register(s1);

cbuffer PerObjectConstants : register(b1) {
    float4x4 gProjectionMatrix;
    float4x4 gViewMatrix;
    float4x4 gModelMatrix;
};

struct VSInput {
    float3 position : POSITION;
    float2 texcoord : TEXCOORD0;
};

struct PSInput {
    float4 position : SV_POSITION;
    float2 texcoord : TEXCOORD0;
};

PSInput MainVS(VSInput input) {
    PSInput output;
    float4 positionWS = mul(gModelMatrix, float4(input.position, 1.0f));
    float4 positionVS = mul(gViewMatrix, positionWS);
    output.position = mul(gProjectionMatrix, positionVS);
    output.texcoord = input.texcoord;
    return output;
}

float4 MainPS(PSInput input) : SV_TARGET {
    return gBaseColorTexture.Sample(gLinearSampler, input.texcoord);
}
)";

const char kBuiltinForwardVertexShader[] = R"(#version 430
layout(location = 0) in vec3 aPosition;
layout(location = 1) in vec2 aTexCoord;

layout(std140, binding = 1) uniform PerObjectConstants {
    mat4 gProjectionMatrix;
    mat4 gViewMatrix;
    mat4 gModelMatrix;
};

out vec2 vTexCoord;

void main() {
    vec4 positionWS = gModelMatrix * vec4(aPosition, 1.0);
    vec4 positionVS = gViewMatrix * positionWS;
    gl_Position = gProjectionMatrix * positionVS;
    vTexCoord = aTexCoord;
}
)";

const char kBuiltinForwardFragmentShader[] = R"(#version 430
layout(binding = 1) uniform sampler2D uBaseColorTexture;

in vec2 vTexCoord;

layout(location = 0) out vec4 fragColor;

void main() {
    fragColor = texture(uBaseColorTexture, vTexCoord);
}
)";

RHI::GraphicsPipelineDesc CreatePipelineDesc(RHI::RHIType backendType, RHI::RHIPipelineLayout* pipelineLayout) {
    RHI::GraphicsPipelineDesc pipelineDesc = {};
    pipelineDesc.pipelineLayout = pipelineLayout;
    pipelineDesc.topologyType = static_cast<uint32_t>(RHI::PrimitiveTopologyType::Triangle);
    pipelineDesc.renderTargetCount = 1;
    pipelineDesc.renderTargetFormats[0] = static_cast<uint32_t>(RHI::Format::R8G8B8A8_UNorm);
    pipelineDesc.depthStencilFormat = static_cast<uint32_t>(RHI::Format::D24_UNorm_S8_UInt);
    pipelineDesc.sampleCount = 1;

    pipelineDesc.rasterizerState.fillMode = static_cast<uint32_t>(RHI::FillMode::Solid);
    pipelineDesc.rasterizerState.cullMode = static_cast<uint32_t>(RHI::CullMode::None);
    pipelineDesc.rasterizerState.frontFace = static_cast<uint32_t>(RHI::FrontFace::CounterClockwise);
    pipelineDesc.rasterizerState.depthClipEnable = true;

    pipelineDesc.depthStencilState.depthTestEnable = true;
    pipelineDesc.depthStencilState.depthWriteEnable = true;
    pipelineDesc.depthStencilState.depthFunc = static_cast<uint32_t>(RHI::ComparisonFunc::Less);
    pipelineDesc.depthStencilState.stencilEnable = false;

    RHI::InputElementDesc position = {};
    position.semanticName = "POSITION";
    position.semanticIndex = 0;
    position.format = static_cast<uint32_t>(RHI::Format::R32G32B32A32_Float);
    position.inputSlot = 0;
    position.alignedByteOffset = 0;
    pipelineDesc.inputLayout.elements.push_back(position);

    RHI::InputElementDesc texcoord = {};
    texcoord.semanticName = "TEXCOORD";
    texcoord.semanticIndex = 0;
    texcoord.format = static_cast<uint32_t>(RHI::Format::R32G32_Float);
    texcoord.inputSlot = 0;
    texcoord.alignedByteOffset = static_cast<uint32_t>(offsetof(Resources::StaticMeshVertex, uv0));
    pipelineDesc.inputLayout.elements.push_back(texcoord);

    if (backendType == RHI::RHIType::D3D12) {
        pipelineDesc.vertexShader.source.assign(
            kBuiltinForwardHlsl,
            kBuiltinForwardHlsl + std::strlen(kBuiltinForwardHlsl));
        pipelineDesc.vertexShader.sourceLanguage = RHI::ShaderLanguage::HLSL;
        pipelineDesc.vertexShader.entryPoint = L"MainVS";
        pipelineDesc.vertexShader.profile = L"vs_5_0";

        pipelineDesc.fragmentShader.source.assign(
            kBuiltinForwardHlsl,
            kBuiltinForwardHlsl + std::strlen(kBuiltinForwardHlsl));
        pipelineDesc.fragmentShader.sourceLanguage = RHI::ShaderLanguage::HLSL;
        pipelineDesc.fragmentShader.entryPoint = L"MainPS";
        pipelineDesc.fragmentShader.profile = L"ps_5_0";
    } else {
        pipelineDesc.vertexShader.source.assign(
            kBuiltinForwardVertexShader,
            kBuiltinForwardVertexShader + std::strlen(kBuiltinForwardVertexShader));
        pipelineDesc.vertexShader.sourceLanguage = RHI::ShaderLanguage::GLSL;
        pipelineDesc.vertexShader.profile = L"vs_4_30";

        pipelineDesc.fragmentShader.source.assign(
            kBuiltinForwardFragmentShader,
            kBuiltinForwardFragmentShader + std::strlen(kBuiltinForwardFragmentShader));
        pipelineDesc.fragmentShader.sourceLanguage = RHI::ShaderLanguage::GLSL;
        pipelineDesc.fragmentShader.profile = L"fs_4_30";
    }

    return pipelineDesc;
}

const Resources::Texture* FindMaterialTexture(const Resources::Material& material) {
    static const char* kTextureNames[] = {
        "baseColorTexture",
        "_BaseColorTexture",
        "_MainTex",
        "albedoTexture",
        "mainTexture",
        "texture"
    };

    for (const char* textureName : kTextureNames) {
        const Resources::ResourceHandle<Resources::Texture> textureHandle = material.GetTexture(textureName);
        if (textureHandle.Get() != nullptr && textureHandle->IsValid()) {
            return textureHandle.Get();
        }
    }

    return nullptr;
}

} // namespace

BuiltinForwardPipeline::~BuiltinForwardPipeline() {
    Shutdown();
}

bool BuiltinForwardPipeline::Initialize(const RenderContext& context) {
    return EnsureInitialized(context);
}

void BuiltinForwardPipeline::Shutdown() {
    DestroyPipelineResources();
}

bool BuiltinForwardPipeline::Render(
    const RenderContext& context,
    const RenderSurface& surface,
    const RenderSceneData& sceneData) {
    if (!EnsureInitialized(context)) {
        return false;
    }

    const std::vector<RHI::RHIResourceView*>& colorAttachments = surface.GetColorAttachments();
    if (colorAttachments.empty()) {
        return false;
    }

    std::vector<RHI::RHIResourceView*> renderTargets = colorAttachments;
    RHI::RHICommandList* commandList = context.commandList;

    if (surface.IsAutoTransitionEnabled()) {
        for (RHI::RHIResourceView* renderTarget : renderTargets) {
            if (renderTarget != nullptr) {
                commandList->TransitionBarrier(
                    renderTarget,
                    surface.GetColorStateBefore(),
                    RHI::ResourceStates::RenderTarget);
            }
        }
    }

    commandList->SetRenderTargets(
        static_cast<uint32_t>(renderTargets.size()),
        renderTargets.data(),
        surface.GetDepthAttachment());

    const RHI::Viewport viewport = {
        0.0f,
        0.0f,
        static_cast<float>(surface.GetWidth()),
        static_cast<float>(surface.GetHeight()),
        0.0f,
        1.0f
    };
    const RHI::Rect scissorRect = {
        0,
        0,
        static_cast<int32_t>(surface.GetWidth()),
        static_cast<int32_t>(surface.GetHeight())
    };
    commandList->SetViewport(viewport);
    commandList->SetScissorRect(scissorRect);

    const Math::Color clearColor = surface.HasClearColorOverride()
        ? surface.GetClearColorOverride()
        : sceneData.cameraData.clearColor;
    const float clearValues[4] = { clearColor.r, clearColor.g, clearColor.b, clearColor.a };
    for (RHI::RHIResourceView* renderTarget : renderTargets) {
        if (renderTarget != nullptr) {
            commandList->ClearRenderTarget(renderTarget, clearValues);
        }
    }
    if (surface.GetDepthAttachment() != nullptr) {
        commandList->ClearDepthStencil(surface.GetDepthAttachment(), 1.0f, 0);
    }

    commandList->SetPipelineState(m_pipelineState);
    commandList->SetPrimitiveTopology(RHI::PrimitiveTopology::TriangleList);

    for (const VisibleRenderItem& visibleItem : sceneData.visibleItems) {
        if (!MatchesBuiltinPass(ResolveMaterial(visibleItem), BuiltinMaterialPass::Forward)) {
            continue;
        }

        DrawVisibleItem(context, sceneData, visibleItem);
    }

    if (surface.IsAutoTransitionEnabled()) {
        for (RHI::RHIResourceView* renderTarget : renderTargets) {
            if (renderTarget != nullptr) {
                commandList->TransitionBarrier(
                    renderTarget,
                    RHI::ResourceStates::RenderTarget,
                    surface.GetColorStateAfter());
            }
        }
    }

    return true;
}

bool BuiltinForwardPipeline::EnsureInitialized(const RenderContext& context) {
    if (!context.IsValid()) {
        return false;
    }

    if (m_initialized &&
        m_device == context.device &&
        m_backendType == context.backendType) {
        return true;
    }

    DestroyPipelineResources();
    m_device = context.device;
    m_backendType = context.backendType;
    m_initialized = CreatePipelineResources(context);
    return m_initialized;
}

bool BuiltinForwardPipeline::CreatePipelineResources(const RenderContext& context) {
    RHI::DescriptorPoolDesc constantPoolDesc = {};
    constantPoolDesc.type = RHI::DescriptorHeapType::CBV_SRV_UAV;
    constantPoolDesc.descriptorCount = 1;
    constantPoolDesc.shaderVisible = false;
    m_constantPool = context.device->CreateDescriptorPool(constantPoolDesc);
    if (m_constantPool == nullptr) {
        return false;
    }

    RHI::DescriptorSetLayoutBinding constantBinding = {};
    constantBinding.binding = 0;
    constantBinding.type = static_cast<uint32_t>(RHI::DescriptorType::CBV);
    constantBinding.count = 1;

    RHI::DescriptorSetLayoutDesc constantLayout = {};
    constantLayout.bindings = &constantBinding;
    constantLayout.bindingCount = 1;
    m_constantSet = m_constantPool->AllocateSet(constantLayout);
    if (m_constantSet == nullptr) {
        return false;
    }

    RHI::DescriptorPoolDesc texturePoolDesc = {};
    texturePoolDesc.type = RHI::DescriptorHeapType::CBV_SRV_UAV;
    texturePoolDesc.descriptorCount = 1;
    texturePoolDesc.shaderVisible = true;
    m_texturePool = context.device->CreateDescriptorPool(texturePoolDesc);
    if (m_texturePool == nullptr) {
        return false;
    }

    RHI::DescriptorSetLayoutBinding textureBinding = {};
    textureBinding.binding = 0;
    textureBinding.type = static_cast<uint32_t>(RHI::DescriptorType::SRV);
    textureBinding.count = 1;

    RHI::DescriptorSetLayoutDesc textureLayout = {};
    textureLayout.bindings = &textureBinding;
    textureLayout.bindingCount = 1;
    m_textureSet = m_texturePool->AllocateSet(textureLayout);
    if (m_textureSet == nullptr) {
        return false;
    }

    RHI::DescriptorPoolDesc samplerPoolDesc = {};
    samplerPoolDesc.type = RHI::DescriptorHeapType::Sampler;
    samplerPoolDesc.descriptorCount = 1;
    samplerPoolDesc.shaderVisible = true;
    m_samplerPool = context.device->CreateDescriptorPool(samplerPoolDesc);
    if (m_samplerPool == nullptr) {
        return false;
    }

    RHI::DescriptorSetLayoutBinding samplerBinding = {};
    samplerBinding.binding = 0;
    samplerBinding.type = static_cast<uint32_t>(RHI::DescriptorType::Sampler);
    samplerBinding.count = 1;

    RHI::DescriptorSetLayoutDesc samplerLayout = {};
    samplerLayout.bindings = &samplerBinding;
    samplerLayout.bindingCount = 1;
    m_samplerSet = m_samplerPool->AllocateSet(samplerLayout);
    if (m_samplerSet == nullptr) {
        return false;
    }

    RHI::DescriptorSetLayoutBinding reservedBindings[3] = {};
    reservedBindings[0].binding = 0;
    reservedBindings[0].type = static_cast<uint32_t>(RHI::DescriptorType::CBV);
    reservedBindings[0].count = 1;
    reservedBindings[1].binding = 0;
    reservedBindings[1].type = static_cast<uint32_t>(RHI::DescriptorType::SRV);
    reservedBindings[1].count = 1;
    reservedBindings[2].binding = 0;
    reservedBindings[2].type = static_cast<uint32_t>(RHI::DescriptorType::Sampler);
    reservedBindings[2].count = 1;

    RHI::DescriptorSetLayoutDesc reservedLayout = {};
    reservedLayout.bindings = reservedBindings;
    reservedLayout.bindingCount = 3;

    RHI::DescriptorSetLayoutDesc setLayouts[kDescriptorSetCount] = {};
    setLayouts[0] = reservedLayout;
    setLayouts[1] = constantLayout;
    setLayouts[2] = textureLayout;
    setLayouts[3] = samplerLayout;

    RHI::RHIPipelineLayoutDesc pipelineLayoutDesc = {};
    pipelineLayoutDesc.setLayouts = setLayouts;
    pipelineLayoutDesc.setLayoutCount = kDescriptorSetCount;
    m_pipelineLayout = context.device->CreatePipelineLayout(pipelineLayoutDesc);
    if (m_pipelineLayout == nullptr) {
        return false;
    }

    const RHI::GraphicsPipelineDesc pipelineDesc = CreatePipelineDesc(context.backendType, m_pipelineLayout);
    m_pipelineState = context.device->CreatePipelineState(pipelineDesc);
    if (m_pipelineState == nullptr || !m_pipelineState->IsValid()) {
        return false;
    }

    RHI::SamplerDesc samplerDesc = {};
    samplerDesc.filter = static_cast<uint32_t>(RHI::FilterMode::Linear);
    samplerDesc.addressU = static_cast<uint32_t>(RHI::TextureAddressMode::Clamp);
    samplerDesc.addressV = static_cast<uint32_t>(RHI::TextureAddressMode::Clamp);
    samplerDesc.addressW = static_cast<uint32_t>(RHI::TextureAddressMode::Clamp);
    samplerDesc.mipLodBias = 0.0f;
    samplerDesc.maxAnisotropy = 1;
    samplerDesc.comparisonFunc = static_cast<uint32_t>(RHI::ComparisonFunc::Always);
    samplerDesc.minLod = 0.0f;
    samplerDesc.maxLod = 1000.0f;
    m_sampler = context.device->CreateSampler(samplerDesc);
    if (m_sampler == nullptr) {
        return false;
    }
    m_samplerSet->UpdateSampler(0, m_sampler);

    const unsigned char whitePixel[4] = { 255, 255, 255, 255 };
    RHI::TextureDesc textureDesc = {};
    textureDesc.width = 1;
    textureDesc.height = 1;
    textureDesc.depth = 1;
    textureDesc.mipLevels = 1;
    textureDesc.arraySize = 1;
    textureDesc.format = static_cast<uint32_t>(RHI::Format::R8G8B8A8_UNorm);
    textureDesc.textureType = static_cast<uint32_t>(RHI::TextureType::Texture2D);
    textureDesc.sampleCount = 1;
    textureDesc.sampleQuality = 0;
    textureDesc.flags = 0;
    m_fallbackTexture = context.device->CreateTexture(textureDesc, whitePixel, sizeof(whitePixel), 4);
    if (m_fallbackTexture == nullptr) {
        return false;
    }

    RHI::ResourceViewDesc textureViewDesc = {};
    textureViewDesc.format = static_cast<uint32_t>(RHI::Format::R8G8B8A8_UNorm);
    textureViewDesc.dimension = RHI::ResourceViewDimension::Texture2D;
    textureViewDesc.mipLevel = 0;
    m_fallbackTextureView = context.device->CreateShaderResourceView(m_fallbackTexture, textureViewDesc);
    if (m_fallbackTextureView == nullptr) {
        return false;
    }

    return true;
}

void BuiltinForwardPipeline::DestroyPipelineResources() {
    m_resourceCache.Shutdown();

    if (m_fallbackTextureView != nullptr) {
        m_fallbackTextureView->Shutdown();
        delete m_fallbackTextureView;
        m_fallbackTextureView = nullptr;
    }

    if (m_fallbackTexture != nullptr) {
        m_fallbackTexture->Shutdown();
        delete m_fallbackTexture;
        m_fallbackTexture = nullptr;
    }

    if (m_sampler != nullptr) {
        m_sampler->Shutdown();
        delete m_sampler;
        m_sampler = nullptr;
    }

    if (m_pipelineState != nullptr) {
        m_pipelineState->Shutdown();
        delete m_pipelineState;
        m_pipelineState = nullptr;
    }

    if (m_pipelineLayout != nullptr) {
        m_pipelineLayout->Shutdown();
        delete m_pipelineLayout;
        m_pipelineLayout = nullptr;
    }

    if (m_constantSet != nullptr) {
        m_constantSet->Shutdown();
        delete m_constantSet;
        m_constantSet = nullptr;
    }

    if (m_textureSet != nullptr) {
        m_textureSet->Shutdown();
        delete m_textureSet;
        m_textureSet = nullptr;
    }

    if (m_samplerSet != nullptr) {
        m_samplerSet->Shutdown();
        delete m_samplerSet;
        m_samplerSet = nullptr;
    }

    if (m_constantPool != nullptr) {
        m_constantPool->Shutdown();
        delete m_constantPool;
        m_constantPool = nullptr;
    }

    if (m_texturePool != nullptr) {
        m_texturePool->Shutdown();
        delete m_texturePool;
        m_texturePool = nullptr;
    }

    if (m_samplerPool != nullptr) {
        m_samplerPool->Shutdown();
        delete m_samplerPool;
        m_samplerPool = nullptr;
    }

    m_device = nullptr;
    m_initialized = false;
}

const Resources::Texture* BuiltinForwardPipeline::ResolveTexture(const Resources::Material* material) const {
    return material != nullptr ? FindMaterialTexture(*material) : nullptr;
}

RHI::RHIResourceView* BuiltinForwardPipeline::ResolveTextureView(
    const VisibleRenderItem& visibleItem) {
    const Resources::Material* material = ResolveMaterial(visibleItem);
    const Resources::Texture* texture = ResolveTexture(material);
    if (texture != nullptr) {
        const RenderResourceCache::CachedTexture* cachedTexture = m_resourceCache.GetOrCreateTexture(m_device, texture);
        if (cachedTexture != nullptr && cachedTexture->shaderResourceView != nullptr) {
            return cachedTexture->shaderResourceView;
        }
    }

    return m_fallbackTextureView;
}

bool BuiltinForwardPipeline::DrawVisibleItem(
    const RenderContext& context,
    const RenderSceneData& sceneData,
    const VisibleRenderItem& visibleItem) {
    const RenderResourceCache::CachedMesh* cachedMesh = m_resourceCache.GetOrCreateMesh(m_device, visibleItem.mesh);
    if (cachedMesh == nullptr || cachedMesh->vertexBufferView == nullptr) {
        return false;
    }

    RHI::RHICommandList* commandList = context.commandList;

    RHI::RHIResourceView* vertexBuffers[] = { cachedMesh->vertexBufferView };
    const uint64_t offsets[] = { 0 };
    const uint32_t strides[] = { cachedMesh->vertexStride };
    commandList->SetVertexBuffers(0, 1, vertexBuffers, offsets, strides);
    if (cachedMesh->indexBufferView != nullptr) {
        commandList->SetIndexBuffer(cachedMesh->indexBufferView, 0);
    }

    const PerObjectConstants constants = {
        sceneData.cameraData.projection,
        sceneData.cameraData.view,
        visibleItem.localToWorld.Transpose()
    };

    RHI::RHIResourceView* textureView = ResolveTextureView(visibleItem);
    if (textureView == nullptr) {
        return false;
    }

    m_constantSet->WriteConstant(0, &constants, sizeof(constants));
    m_textureSet->Update(0, textureView);
    RHI::RHIDescriptorSet* descriptorSets[] = { m_constantSet, m_textureSet, m_samplerSet };
    commandList->SetGraphicsDescriptorSets(kDescriptorFirstSet, 3, descriptorSets, m_pipelineLayout);

    if (visibleItem.hasSection) {
        const Containers::Array<Resources::MeshSection>& sections = visibleItem.mesh->GetSections();
        if (visibleItem.sectionIndex >= sections.Size()) {
            return false;
        }

        const Resources::MeshSection& section = sections[visibleItem.sectionIndex];
        if (cachedMesh->indexBufferView != nullptr && section.indexCount > 0) {
            // MeshLoader flattens section indices into a single global index buffer.
            commandList->DrawIndexed(section.indexCount, 1, section.startIndex, 0, 0);
        } else if (section.vertexCount > 0) {
            commandList->Draw(section.vertexCount, 1, section.baseVertex, 0);
        }

        return true;
    }

    if (cachedMesh->indexBufferView != nullptr && cachedMesh->indexCount > 0) {
        commandList->DrawIndexed(cachedMesh->indexCount, 1, 0, 0, 0);
    } else if (cachedMesh->vertexCount > 0) {
        commandList->Draw(cachedMesh->vertexCount, 1, 0, 0);
    }

    return true;
}

} // namespace Pipelines
} // namespace Rendering
} // namespace XCEngine