XCEngine/engine/src/Resources/Shader/Internal/ShaderRuntimeBuildUtils.cpp

#include "ShaderRuntimeBuildUtils.h"

#include "../ShaderAuthoringParser.h"
#include "../ShaderSourceUtils.h"
#include "ShaderAuthoringInternal.h"
#include "ShaderFileUtils.h"

#include <XCEngine/Core/Asset/ResourceManager.h>
#include <XCEngine/Rendering/Builtin/BuiltinPassLayoutUtils.h>
#include <XCEngine/Resources/BuiltinResources.h>
#include <XCEngine/Resources/Shader/ShaderLoader.h>

#include <filesystem>
#include <memory>
#include <system_error>
#include <unordered_map>
#include <unordered_set>

namespace XCEngine {
namespace Resources {

LoadResult BuildShaderFromIRInternal(
    const Containers::String& path,
    const ShaderIR& shaderIR,
    std::unordered_set<std::string>& activeShaderPathKeys);

namespace {

namespace fs = std::filesystem;

Containers::String ResolveBuiltinShaderPathByAuthoringName(const Containers::String& shaderName) {
    Containers::String builtinShaderPath;
    if (TryGetBuiltinShaderPathByShaderName(shaderName, builtinShaderPath)) {
        return builtinShaderPath;
    }

    return Containers::String();
}

std::string BuildNormalizedShaderPathKey(const Containers::String& shaderPath) {
    if (shaderPath.Empty()) {
        return std::string();
    }

    fs::path normalizedPath(shaderPath.CStr());
    if (!normalizedPath.is_absolute()) {
        const Containers::String& resourceRoot = ResourceManager::Get().GetResourceRoot();
        if (!resourceRoot.Empty()) {
            normalizedPath = fs::path(resourceRoot.CStr()) / normalizedPath;
        } else {
            std::error_code ec;
            normalizedPath = fs::current_path(ec) / normalizedPath;
        }
    }

    return normalizedPath.lexically_normal().generic_string();
}

void AddUniqueSearchRoot(
    const fs::path& candidate,
    std::vector<fs::path>& searchRoots,
    std::unordered_set<std::string>& seenRoots) {
    if (candidate.empty()) {
        return;
    }

    const fs::path normalized = candidate.lexically_normal();
    const std::string key = normalized.generic_string();
    if (key.empty()) {
        return;
    }

    if (seenRoots.insert(key).second) {
        searchRoots.push_back(normalized);
    }
}

bool TryReadDeclaredAuthoringShaderName(
    const Containers::String& shaderPath,
    Containers::String& outShaderName) {
    outShaderName.Clear();

    const Containers::Array<Core::uint8> data = ReadShaderFileData(shaderPath);
    if (data.Empty()) {
        return false;
    }

    const std::string sourceText = ToStdStringFromBytes(data);
    std::vector<std::string> lines;
    Internal::SplitShaderAuthoringLines(sourceText, lines);
    if (lines.empty() || !Internal::StartsWithKeyword(lines.front(), "Shader")) {
        return false;
    }

    std::vector<std::string> tokens;
    if (!Internal::TryTokenizeQuotedArguments(lines.front(), tokens) || tokens.size() < 2u) {
        return false;
    }

    outShaderName = tokens[1].c_str();
    return !outShaderName.Empty();
}

bool TryResolveProjectShaderPathByAuthoringName(
    const Containers::String& currentShaderPath,
    const Containers::String& shaderName,
    Containers::String& outResolvedPath) {
    outResolvedPath.Clear();

    fs::path currentPath(currentShaderPath.CStr());
    if (!currentPath.is_absolute()) {
        const Containers::String& resourceRoot = ResourceManager::Get().GetResourceRoot();
        if (!resourceRoot.Empty()) {
            currentPath = fs::path(resourceRoot.CStr()) / currentPath;
        }
    }

    if (currentPath.empty()) {
        return false;
    }

    std::vector<fs::path> searchRoots;
    std::unordered_set<std::string> seenRoots;

    AddUniqueSearchRoot(currentPath.parent_path(), searchRoots, seenRoots);

    for (fs::path ancestor = currentPath.parent_path();
         !ancestor.empty();
         ancestor = ancestor.parent_path()) {
        std::error_code ec;
        const fs::path assetsRoot = ancestor / "Assets";
        if (fs::exists(assetsRoot, ec) && fs::is_directory(assetsRoot, ec)) {
            AddUniqueSearchRoot(assetsRoot, searchRoots, seenRoots);
        }

        const fs::path parent = ancestor.parent_path();
        if (parent == ancestor) {
            break;
        }
    }

    std::vector<fs::path> matches;
    std::unordered_set<std::string> seenMatches;

    for (const fs::path& root : searchRoots) {
        std::error_code ec;
        if (!fs::exists(root, ec) || !fs::is_directory(root, ec)) {
            continue;
        }

        for (fs::recursive_directory_iterator it(root, ec), end; !ec && it != end; it.increment(ec)) {
            if (!it->is_regular_file()) {
                continue;
            }

            const fs::path candidatePath = it->path();
            if (candidatePath.extension() != ".shader") {
                continue;
            }

            Containers::String declaredShaderName;
            if (!TryReadDeclaredAuthoringShaderName(
                    Containers::String(candidatePath.generic_string().c_str()),
                    declaredShaderName) ||
                declaredShaderName != shaderName) {
                continue;
            }

            const std::string matchKey = candidatePath.lexically_normal().generic_string();
            if (seenMatches.insert(matchKey).second) {
                matches.push_back(candidatePath.lexically_normal());
            }
        }
    }

    if (matches.size() != 1u) {
        return false;
    }

    outResolvedPath = matches.front().generic_string().c_str();
    return true;
}

void ImportConcretePass(Shader& shader, const ShaderPass& sourcePass) {
    if (ShaderPass* existingPass = shader.FindPass(sourcePass.name)) {
        *existingPass = sourcePass;
        return;
    }

    shader.AddPass(sourcePass);
}

ShaderPass BuildConcretePass(
    const ShaderIR& shaderIR,
    const ShaderSubShaderIR& subShader,
    const ShaderPassIR& pass) {
    ShaderPass shaderPass = {};
    shaderPass.name = pass.name;
    shaderPass.hasFixedFunctionState = pass.hasFixedFunctionState;
    shaderPass.fixedFunctionState = pass.fixedFunctionState;

    for (const ShaderTagIR& subShaderTag : subShader.tags) {
        shaderPass.tags.PushBack({ subShaderTag.name, subShaderTag.value });
    }
    for (const ShaderTagIR& passTag : pass.tags) {
        shaderPass.tags.PushBack({ passTag.name, passTag.value });
    }
    for (const ShaderResourceBindingDesc& resourceBinding : pass.resources) {
        shaderPass.resources.PushBack(resourceBinding);
    }
    if (shaderPass.resources.Empty()) {
        Containers::Array<ShaderResourceBindingDesc> defaultBindings;
        if (::XCEngine::Rendering::TryBuildBuiltinPassDefaultResourceBindings(shaderPass, defaultBindings)) {
            for (const ShaderResourceBindingDesc& resourceBinding : defaultBindings) {
                shaderPass.resources.PushBack(resourceBinding);
            }
        }
    }
    for (const ShaderKeywordDeclaration& keywordDeclaration : pass.keywordDeclarations) {
        shaderPass.keywordDeclarations.PushBack(keywordDeclaration);
    }

    if (pass.programSource.Empty()) {
        return shaderPass;
    }

    Containers::String combinedSource;
    AppendAuthoringSourceBlock(combinedSource, shaderIR.sharedProgramSource);
    AppendAuthoringSourceBlock(combinedSource, subShader.sharedProgramSource);
    AppendAuthoringSourceBlock(combinedSource, pass.sharedProgramSource);
    AppendAuthoringSourceBlock(combinedSource, pass.programSource);

    const Containers::String strippedCombinedSource = StripShaderAuthoringPragmas(combinedSource);
    const std::vector<ShaderKeywordSet> keywordSets =
        BuildShaderKeywordVariantSets(pass.keywordDeclarations);

    for (const ShaderKeywordSet& keywordSet : keywordSets) {
        const Containers::String variantSource =
            BuildKeywordVariantSource(strippedCombinedSource, keywordSet);

        ShaderStageVariant vertexVariant = {};
        vertexVariant.stage = ShaderType::Vertex;
        vertexVariant.backend = ShaderBackend::Generic;
        vertexVariant.language = ShaderLanguage::HLSL;
        vertexVariant.requiredKeywords = keywordSet;
        vertexVariant.entryPoint =
            !pass.vertexEntryPoint.Empty()
                ? pass.vertexEntryPoint
                : GetDefaultEntryPoint(ShaderLanguage::HLSL, ShaderType::Vertex);
        vertexVariant.profile = GetDefaultProfile(
            ShaderLanguage::HLSL,
            ShaderBackend::Generic,
            ShaderType::Vertex);
        vertexVariant.sourceCode = variantSource;
        shaderPass.variants.PushBack(vertexVariant);

        ShaderStageVariant fragmentVariant = {};
        fragmentVariant.stage = ShaderType::Fragment;
        fragmentVariant.backend = ShaderBackend::Generic;
        fragmentVariant.language = ShaderLanguage::HLSL;
        fragmentVariant.requiredKeywords = keywordSet;
        fragmentVariant.entryPoint =
            !pass.fragmentEntryPoint.Empty()
                ? pass.fragmentEntryPoint
                : GetDefaultEntryPoint(ShaderLanguage::HLSL, ShaderType::Fragment);
        fragmentVariant.profile = GetDefaultProfile(
            ShaderLanguage::HLSL,
            ShaderBackend::Generic,
            ShaderType::Fragment);
        fragmentVariant.sourceCode = variantSource;
        shaderPass.variants.PushBack(fragmentVariant);
    }

    return shaderPass;
}

bool TryLoadReferencedUsePassShader(
    const Containers::String& referencedShaderPath,
    std::unordered_set<std::string>& activeShaderPathKeys,
    std::unique_ptr<Shader>& outShader,
    Containers::String& outError) {
    Containers::String authoringShaderPath = referencedShaderPath;
    if (IsBuiltinShaderPath(authoringShaderPath) &&
        !TryResolveBuiltinShaderAssetPath(authoringShaderPath, authoringShaderPath)) {
        outError =
            Containers::String("UsePass failed to resolve builtin shader asset: ") +
            referencedShaderPath;
        return false;
    }

    const std::string shaderPathKey = BuildNormalizedShaderPathKey(authoringShaderPath);
    if (shaderPathKey.empty()) {
        outError =
            Containers::String("UsePass could not normalize referenced shader path: ") +
            referencedShaderPath;
        return false;
    }

    if (activeShaderPathKeys.find(shaderPathKey) != activeShaderPathKeys.end()) {
        outError =
            Containers::String("UsePass detected a cyclic shader reference: ") +
            authoringShaderPath;
        return false;
    }

    Containers::String sourceText;
    if (!ReadShaderTextFile(authoringShaderPath, sourceText)) {
        outError =
            Containers::String("UsePass failed to read referenced shader: ") +
            authoringShaderPath;
        return false;
    }

    ShaderIR referencedShaderIR = {};
    Containers::String parseError;
    if (!ParseShaderAuthoring(authoringShaderPath, sourceText.CStr(), referencedShaderIR, &parseError)) {
        outError = parseError;
        return false;
    }

    activeShaderPathKeys.insert(shaderPathKey);
    LoadResult referencedShaderResult =
        BuildShaderFromIRInternal(authoringShaderPath, referencedShaderIR, activeShaderPathKeys);
    activeShaderPathKeys.erase(shaderPathKey);
    if (!referencedShaderResult || referencedShaderResult.resource == nullptr) {
        outError =
            !referencedShaderResult.errorMessage.Empty()
                ? referencedShaderResult.errorMessage
                : Containers::String("UsePass failed to build referenced shader: ") + authoringShaderPath;
        return false;
    }

    outShader.reset(static_cast<Shader*>(referencedShaderResult.resource));
    return true;
}

bool TryResolveUsePass(
    const Containers::String& currentShaderPath,
    const Containers::String& currentShaderName,
    const Containers::String& referencedShaderName,
    const Containers::String& referencedPassName,
    const std::unordered_map<std::string, ShaderPass>& localConcretePasses,
    std::unordered_set<std::string>& activeShaderPathKeys,
    ShaderPass& outPass,
    Containers::String& outError) {
    if (referencedShaderName == currentShaderName) {
        const auto passIt = localConcretePasses.find(ToStdString(referencedPassName));
        if (passIt == localConcretePasses.end()) {
            outError =
                Containers::String("UsePass could not resolve local pass: ") + referencedPassName;
            return false;
        }

        outPass = passIt->second;
        return true;
    }

    Containers::String referencedShaderPath;
    if (!ResolveShaderUsePassPath(
            currentShaderPath,
            currentShaderName,
            referencedShaderName,
            referencedShaderPath)) {
        outError =
            Containers::String("UsePass could not resolve referenced shader: ") +
            referencedShaderName;
        return false;
    }

    std::unique_ptr<Shader> referencedShader;
    if (!TryLoadReferencedUsePassShader(
            referencedShaderPath,
            activeShaderPathKeys,
            referencedShader,
            outError)) {
        return false;
    }

    const ShaderPass* referencedPass = referencedShader->FindPass(referencedPassName);
    if (referencedPass == nullptr) {
        outError =
            Containers::String("UsePass could not find referenced pass: ") + referencedPassName;
        return false;
    }

    outPass = *referencedPass;
    return true;
}

} // namespace

Containers::String GetDefaultEntryPoint(ShaderLanguage language, ShaderType stage) {
    if (language == ShaderLanguage::HLSL) {
        switch (stage) {
        case ShaderType::Vertex: return "VSMain";
        case ShaderType::Fragment: return "PSMain";
        case ShaderType::Geometry: return "GSMain";
        case ShaderType::Compute: return "CSMain";
        case ShaderType::Hull: return "HSMain";
        case ShaderType::Domain: return "DSMain";
        default: return "main";
        }
    }

    return "main";
}

Containers::String GetDefaultProfile(
    ShaderLanguage language,
    ShaderBackend backend,
    ShaderType stage) {
    if (language == ShaderLanguage::HLSL) {
        switch (stage) {
        case ShaderType::Vertex: return "vs_5_0";
        case ShaderType::Fragment: return "ps_5_0";
        case ShaderType::Geometry: return "gs_5_0";
        case ShaderType::Compute: return "cs_5_0";
        case ShaderType::Hull: return "hs_5_0";
        case ShaderType::Domain: return "ds_5_0";
        default: return Containers::String();
        }
    }

    const bool isVulkan = backend == ShaderBackend::Vulkan;
    switch (stage) {
    case ShaderType::Vertex:
        return isVulkan ? "vs_4_50" : "vs_4_30";
    case ShaderType::Fragment:
        return isVulkan ? "fs_4_50" : "fs_4_30";
    case ShaderType::Geometry:
        return isVulkan ? "gs_4_50" : "gs_4_30";
    case ShaderType::Compute:
        return isVulkan ? "cs_4_50" : "cs_4_30";
    case ShaderType::Hull:
        return isVulkan ? "hs_4_50" : "hs_4_30";
    case ShaderType::Domain:
        return isVulkan ? "ds_4_50" : "ds_4_30";
    default:
        return Containers::String();
    }
}

size_t CalculateShaderMemorySize(const Shader& shader) {
    size_t memorySize =
        sizeof(Shader) + shader.GetName().Length() + shader.GetPath().Length() + shader.GetFallback().Length();
    for (const ShaderPropertyDesc& property : shader.GetProperties()) {
        memorySize += property.name.Length();
        memorySize += property.displayName.Length();
        memorySize += property.defaultValue.Length();
        memorySize += property.semantic.Length();
    }
    for (const ShaderPass& pass : shader.GetPasses()) {
        memorySize += pass.name.Length();
        for (const ShaderPassTagEntry& tag : pass.tags) {
            memorySize += tag.name.Length();
            memorySize += tag.value.Length();
        }
        for (const ShaderResourceBindingDesc& binding : pass.resources) {
            memorySize += binding.name.Length();
            memorySize += binding.semantic.Length();
        }
        for (const ShaderKeywordDeclaration& declaration : pass.keywordDeclarations) {
            for (const Containers::String& option : declaration.options) {
                memorySize += option.Length();
            }
        }
        for (const ShaderStageVariant& variant : pass.variants) {
            for (const Containers::String& keyword : variant.requiredKeywords.enabledKeywords) {
                memorySize += keyword.Length();
            }
            memorySize += variant.entryPoint.Length();
            memorySize += variant.profile.Length();
            memorySize += variant.sourceCode.Length();
            memorySize += variant.compiledBinary.Size();
        }
    }

    return memorySize;
}

bool ResolveShaderUsePassPath(
    const Containers::String& currentShaderPath,
    const Containers::String& currentShaderName,
    const Containers::String& targetShaderName,
    Containers::String& outResolvedPath) {
    outResolvedPath.Clear();

    if (targetShaderName.Empty()) {
        return false;
    }

    if (targetShaderName == currentShaderName) {
        outResolvedPath = currentShaderPath;
        return true;
    }

    outResolvedPath = ResolveBuiltinShaderPathByAuthoringName(targetShaderName);
    if (!outResolvedPath.Empty()) {
        return true;
    }

    return TryResolveProjectShaderPathByAuthoringName(
        currentShaderPath,
        targetShaderName,
        outResolvedPath);
}

LoadResult BuildShaderFromIRInternal(
    const Containers::String& path,
    const ShaderIR& shaderIR,
    std::unordered_set<std::string>& activeShaderPathKeys) {
    auto shader = std::make_unique<Shader>();
    IResource::ConstructParams params;
    params.path = path;
    params.guid = ResourceGUID::Generate(path);
    params.name = shaderIR.name;
    shader->Initialize(params);
    shader->SetFallback(shaderIR.fallback);

    for (const ShaderPropertyDesc& property : shaderIR.properties) {
        shader->AddProperty(property);
    }

    std::unordered_map<std::string, ShaderPass> localConcretePasses;
    for (const ShaderSubShaderIR& subShader : shaderIR.subShaders) {
        for (const ShaderPassIR& pass : subShader.passes) {
            if (pass.isUsePass) {
                continue;
            }

            ShaderPass concretePass = BuildConcretePass(shaderIR, subShader, pass);
            const std::string passKey = ToStdString(concretePass.name);
            if (passKey.empty()) {
                return LoadResult("Shader authoring produced a pass with an empty name");
            }

            if (!localConcretePasses.emplace(passKey, std::move(concretePass)).second) {
                return LoadResult(
                    Containers::String("Shader authoring produced duplicate pass name: ") + pass.name);
            }
        }
    }

    for (const ShaderSubShaderIR& subShader : shaderIR.subShaders) {
        for (const ShaderPassIR& pass : subShader.passes) {
            if (pass.isUsePass) {
                ShaderPass importedPass = {};
                Containers::String importError;
                if (!TryResolveUsePass(
                        path,
                        shaderIR.name,
                        pass.usePassShaderName,
                        pass.usePassPassName,
                        localConcretePasses,
                        activeShaderPathKeys,
                        importedPass,
                        importError)) {
                    return LoadResult(importError);
                }

                ImportConcretePass(*shader, importedPass);
                for (const ShaderTagIR& subShaderTag : subShader.tags) {
                    shader->SetPassTag(importedPass.name, subShaderTag.name, subShaderTag.value);
                }
                continue;
            }

            const auto passIt = localConcretePasses.find(ToStdString(pass.name));
            if (passIt == localConcretePasses.end()) {
                return LoadResult(
                    Containers::String("Shader authoring lost concrete pass during build: ") + pass.name);
            }

            ImportConcretePass(*shader, passIt->second);
        }
    }

    shader->m_memorySize = CalculateShaderMemorySize(*shader);
    return LoadResult(shader.release());
}

LoadResult BuildShaderFromIR(
    const Containers::String& path,
    const ShaderIR& shaderIR) {
    std::unordered_set<std::string> activeShaderPathKeys;
    const std::string shaderPathKey = BuildNormalizedShaderPathKey(path);
    if (!shaderPathKey.empty()) {
        activeShaderPathKeys.insert(shaderPathKey);
    }

    return BuildShaderFromIRInternal(path, shaderIR, activeShaderPathKeys);
}

} // namespace Resources
} // namespace XCEngine