XCEngine/engine/src/Resources/Shader/ShaderLoader.cpp

#include <XCEngine/Resources/Shader/ShaderLoader.h>

#include <XCEngine/Core/Asset/ResourceManager.h>
#include <XCEngine/Core/Asset/ResourceTypes.h>
#include <XCEngine/Resources/BuiltinResources.h>

#include <cctype>
#include <filesystem>
#include <fstream>
#include <functional>
#include <memory>
#include <string>
#include <vector>

namespace XCEngine {
namespace Resources {

namespace {

std::string ToStdString(const Containers::Array<Core::uint8>& data) {
    return std::string(reinterpret_cast<const char*>(data.Data()), data.Size());
}

Containers::Array<Core::uint8> TryReadFileData(
    const std::filesystem::path& filePath,
    bool& opened) {
    Containers::Array<Core::uint8> data;

    std::ifstream file(filePath, std::ios::binary | std::ios::ate);
    if (!file.is_open()) {
        opened = false;
        return data;
    }

    opened = true;
    const std::streamsize size = file.tellg();
    if (size <= 0) {
        return data;
    }

    file.seekg(0, std::ios::beg);
    data.Resize(static_cast<size_t>(size));
    if (!file.read(reinterpret_cast<char*>(data.Data()), size)) {
        data.Clear();
    }

    return data;
}

Containers::Array<Core::uint8> ReadShaderFileData(const Containers::String& path) {
    bool opened = false;
    const std::filesystem::path inputPath(path.CStr());
    Containers::Array<Core::uint8> data = TryReadFileData(inputPath, opened);
    if (opened || path.Empty() || inputPath.is_absolute()) {
        return data;
    }

    const Containers::String& resourceRoot = ResourceManager::Get().GetResourceRoot();
    if (resourceRoot.Empty()) {
        return data;
    }

    return TryReadFileData(std::filesystem::path(resourceRoot.CStr()) / inputPath, opened);
}

Containers::String NormalizePathString(const std::filesystem::path& path) {
    return Containers::String(path.lexically_normal().generic_string().c_str());
}

Containers::String GetPathExtension(const Containers::String& path) {
    size_t dotPos = Containers::String::npos;
    for (size_t i = path.Length(); i > 0; --i) {
        if (path[i - 1] == '.') {
            dotPos = i - 1;
            break;
        }
    }

    if (dotPos == Containers::String::npos) {
        return Containers::String();
    }

    return path.Substring(dotPos + 1);
}

size_t SkipWhitespace(const std::string& text, size_t pos) {
    while (pos < text.size() && std::isspace(static_cast<unsigned char>(text[pos])) != 0) {
        ++pos;
    }
    return pos;
}

std::string TrimCopy(const std::string& text) {
    const size_t first = SkipWhitespace(text, 0);
    if (first >= text.size()) {
        return std::string();
    }

    size_t last = text.size();
    while (last > first && std::isspace(static_cast<unsigned char>(text[last - 1])) != 0) {
        --last;
    }

    return text.substr(first, last - first);
}

bool FindValueStart(const std::string& json, const char* key, size_t& valuePos) {
    const std::string token = std::string("\"") + key + "\"";
    const size_t keyPos = json.find(token);
    if (keyPos == std::string::npos) {
        return false;
    }

    const size_t colonPos = json.find(':', keyPos + token.length());
    if (colonPos == std::string::npos) {
        return false;
    }

    valuePos = SkipWhitespace(json, colonPos + 1);
    return valuePos < json.size();
}

bool ParseQuotedString(
    const std::string& text,
    size_t quotePos,
    Containers::String& outValue,
    size_t* nextPos = nullptr) {
    if (quotePos >= text.size() || text[quotePos] != '"') {
        return false;
    }

    std::string parsed;
    ++quotePos;

    while (quotePos < text.size()) {
        const char ch = text[quotePos];
        if (ch == '\\') {
            if (quotePos + 1 >= text.size()) {
                return false;
            }

            parsed.push_back(text[quotePos + 1]);
            quotePos += 2;
            continue;
        }

        if (ch == '"') {
            outValue = parsed.c_str();
            if (nextPos != nullptr) {
                *nextPos = quotePos + 1;
            }
            return true;
        }

        parsed.push_back(ch);
        ++quotePos;
    }

    return false;
}

bool TryParseStringValue(const std::string& json, const char* key, Containers::String& outValue) {
    size_t valuePos = 0;
    if (!FindValueStart(json, key, valuePos)) {
        return false;
    }

    return ParseQuotedString(json, valuePos, outValue);
}

bool TryExtractDelimitedValue(
    const std::string& json,
    const char* key,
    char openChar,
    char closeChar,
    std::string& outValue) {
    size_t valuePos = 0;
    if (!FindValueStart(json, key, valuePos) || valuePos >= json.size() || json[valuePos] != openChar) {
        return false;
    }

    bool inString = false;
    bool escaped = false;
    int depth = 0;

    for (size_t pos = valuePos; pos < json.size(); ++pos) {
        const char ch = json[pos];
        if (escaped) {
            escaped = false;
            continue;
        }

        if (ch == '\\') {
            escaped = true;
            continue;
        }

        if (ch == '"') {
            inString = !inString;
            continue;
        }

        if (inString) {
            continue;
        }

        if (ch == openChar) {
            ++depth;
        } else if (ch == closeChar) {
            --depth;
            if (depth == 0) {
                outValue = json.substr(valuePos, pos - valuePos + 1);
                return true;
            }
        }
    }

    return false;
}

bool TryExtractObject(const std::string& json, const char* key, std::string& outObject) {
    return TryExtractDelimitedValue(json, key, '{', '}', outObject);
}

bool TryExtractArray(const std::string& json, const char* key, std::string& outArray) {
    return TryExtractDelimitedValue(json, key, '[', ']', outArray);
}

bool TryParseStringMapObject(
    const std::string& objectText,
    const std::function<void(const Containers::String&, const Containers::String&)>& onEntry) {
    if (!onEntry || objectText.empty() || objectText.front() != '{' || objectText.back() != '}') {
        return false;
    }

    size_t pos = 1;
    while (pos < objectText.size()) {
        pos = SkipWhitespace(objectText, pos);
        if (pos >= objectText.size()) {
            return false;
        }

        if (objectText[pos] == '}') {
            return true;
        }

        Containers::String key;
        if (!ParseQuotedString(objectText, pos, key, &pos)) {
            return false;
        }

        pos = SkipWhitespace(objectText, pos);
        if (pos >= objectText.size() || objectText[pos] != ':') {
            return false;
        }

        pos = SkipWhitespace(objectText, pos + 1);
        Containers::String value;
        if (!ParseQuotedString(objectText, pos, value, &pos)) {
            return false;
        }

        onEntry(key, value);

        pos = SkipWhitespace(objectText, pos);
        if (pos >= objectText.size()) {
            return false;
        }

        if (objectText[pos] == ',') {
            ++pos;
            continue;
        }

        if (objectText[pos] == '}') {
            return true;
        }

        return false;
    }

    return false;
}

bool SplitTopLevelArrayElements(const std::string& arrayText, std::vector<std::string>& outElements) {
    outElements.clear();
    if (arrayText.size() < 2 || arrayText.front() != '[' || arrayText.back() != ']') {
        return false;
    }

    bool inString = false;
    bool escaped = false;
    int objectDepth = 0;
    int arrayDepth = 0;
    size_t elementStart = std::string::npos;

    for (size_t pos = 1; pos + 1 < arrayText.size(); ++pos) {
        const char ch = arrayText[pos];
        if (escaped) {
            escaped = false;
            continue;
        }

        if (ch == '\\') {
            escaped = true;
            continue;
        }

        if (ch == '"') {
            if (elementStart == std::string::npos) {
                elementStart = pos;
            }
            inString = !inString;
            continue;
        }

        if (inString) {
            continue;
        }

        if (std::isspace(static_cast<unsigned char>(ch)) != 0) {
            continue;
        }

        if (elementStart == std::string::npos) {
            elementStart = pos;
        }

        if (ch == '{') {
            ++objectDepth;
            continue;
        }
        if (ch == '[') {
            ++arrayDepth;
            continue;
        }
        if (ch == '}') {
            --objectDepth;
            continue;
        }
        if (ch == ']') {
            --arrayDepth;
            continue;
        }

        if (ch == ',' && objectDepth == 0 && arrayDepth == 0) {
            if (elementStart != std::string::npos && pos > elementStart) {
                outElements.push_back(TrimCopy(arrayText.substr(elementStart, pos - elementStart)));
            }
            elementStart = std::string::npos;
        }
    }

    if (elementStart != std::string::npos) {
        const std::string tail = TrimCopy(arrayText.substr(elementStart, arrayText.size() - 1 - elementStart));
        if (!tail.empty()) {
            outElements.push_back(tail);
        }
    }

    return true;
}

bool TryParseShaderType(const Containers::String& value, ShaderType& outType) {
    const Containers::String normalized = value.Trim().ToLower();
    if (normalized == "vertex" || normalized == "vs") {
        outType = ShaderType::Vertex;
        return true;
    }
    if (normalized == "fragment" || normalized == "pixel" || normalized == "ps") {
        outType = ShaderType::Fragment;
        return true;
    }
    if (normalized == "geometry" || normalized == "gs") {
        outType = ShaderType::Geometry;
        return true;
    }
    if (normalized == "compute" || normalized == "cs") {
        outType = ShaderType::Compute;
        return true;
    }
    if (normalized == "hull" || normalized == "hs") {
        outType = ShaderType::Hull;
        return true;
    }
    if (normalized == "domain" || normalized == "ds") {
        outType = ShaderType::Domain;
        return true;
    }

    return false;
}

bool TryParseShaderLanguage(const Containers::String& value, ShaderLanguage& outLanguage) {
    const Containers::String normalized = value.Trim().ToLower();
    if (normalized == "glsl") {
        outLanguage = ShaderLanguage::GLSL;
        return true;
    }
    if (normalized == "hlsl") {
        outLanguage = ShaderLanguage::HLSL;
        return true;
    }
    if (normalized == "spirv" || normalized == "spv") {
        outLanguage = ShaderLanguage::SPIRV;
        return true;
    }

    return false;
}

bool TryParseShaderBackend(const Containers::String& value, ShaderBackend& outBackend) {
    const Containers::String normalized = value.Trim().ToLower();
    if (normalized == "generic") {
        outBackend = ShaderBackend::Generic;
        return true;
    }
    if (normalized == "d3d12" || normalized == "dx12") {
        outBackend = ShaderBackend::D3D12;
        return true;
    }
    if (normalized == "opengl" || normalized == "gl") {
        outBackend = ShaderBackend::OpenGL;
        return true;
    }
    if (normalized == "vulkan" || normalized == "vk") {
        outBackend = ShaderBackend::Vulkan;
        return true;
    }

    return false;
}

Containers::String GetDefaultEntryPoint(ShaderLanguage language, ShaderType stage) {
    if (language != ShaderLanguage::HLSL) {
        return Containers::String("main");
    }

    switch (stage) {
    case ShaderType::Vertex: return "MainVS";
    case ShaderType::Fragment: return "MainPS";
    case ShaderType::Geometry: return "MainGS";
    case ShaderType::Compute: return "MainCS";
    case ShaderType::Hull: return "MainHS";
    case ShaderType::Domain: return "MainDS";
    default: return Containers::String();
    }
}

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

Containers::String ResolveShaderDependencyPath(
    const Containers::String& dependencyPath,
    const Containers::String& sourcePath) {
    if (dependencyPath.Empty()) {
        return dependencyPath;
    }

    const std::filesystem::path dependencyFsPath(dependencyPath.CStr());
    if (dependencyFsPath.is_absolute()) {
        return NormalizePathString(dependencyFsPath);
    }

    const std::filesystem::path sourceFsPath(sourcePath.CStr());
    if (sourceFsPath.is_absolute()) {
        return NormalizePathString(sourceFsPath.parent_path() / dependencyFsPath);
    }

    const Containers::String& resourceRoot = ResourceManager::Get().GetResourceRoot();
    if (!resourceRoot.Empty()) {
        return NormalizePathString(
            std::filesystem::path(resourceRoot.CStr()) /
            sourceFsPath.parent_path() /
            dependencyFsPath);
    }

    return NormalizePathString(sourceFsPath.parent_path() / dependencyFsPath);
}

bool ReadTextFile(const Containers::String& path, Containers::String& outText) {
    const Containers::Array<Core::uint8> data = ReadShaderFileData(path);
    if (data.Empty()) {
        return false;
    }

    outText = ToStdString(data).c_str();
    return true;
}

size_t CalculateShaderMemorySize(const Shader& shader) {
    size_t memorySize = sizeof(Shader) + shader.GetName().Length() + shader.GetPath().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 ShaderStageVariant& variant : pass.variants) {
            memorySize += variant.entryPoint.Length();
            memorySize += variant.profile.Length();
            memorySize += variant.sourceCode.Length();
            memorySize += variant.compiledBinary.Size();
        }
    }

    return memorySize;
}

ShaderType DetectShaderTypeFromPath(const Containers::String& path) {
    const Containers::String ext = GetPathExtension(path).ToLower();
    if (ext == "vert") return ShaderType::Vertex;
    if (ext == "frag") return ShaderType::Fragment;
    if (ext == "geom") return ShaderType::Geometry;
    if (ext == "comp") return ShaderType::Compute;

    return ShaderType::Fragment;
}

bool LooksLikeShaderManifest(const std::string& sourceText) {
    const size_t firstContentPos = SkipWhitespace(sourceText, 0);
    return firstContentPos < sourceText.size() &&
           sourceText[firstContentPos] == '{' &&
           sourceText.find("\"passes\"") != std::string::npos;
}

LoadResult LoadShaderManifest(const Containers::String& path, const std::string& jsonText) {
    std::string passesArray;
    if (!TryExtractArray(jsonText, "passes", passesArray)) {
        return LoadResult("Shader manifest is missing a valid passes array: " + path);
    }

    std::vector<std::string> passObjects;
    if (!SplitTopLevelArrayElements(passesArray, passObjects) || passObjects.empty()) {
        return LoadResult("Shader manifest does not contain any pass objects: " + path);
    }

    auto shader = std::make_unique<Shader>();
    IResource::ConstructParams params;
    params.path = path;
    params.guid = ResourceGUID::Generate(path);

    Containers::String manifestName;
    if (TryParseStringValue(jsonText, "name", manifestName) && !manifestName.Empty()) {
        params.name = manifestName;
    } else {
        const std::filesystem::path shaderPath(path.CStr());
        const std::string stem = shaderPath.stem().generic_string();
        params.name = stem.empty() ? path : Containers::String(stem.c_str());
    }

    shader->Initialize(params);

    for (const std::string& passObject : passObjects) {
        Containers::String passName;
        if (!TryParseStringValue(passObject, "name", passName) || passName.Empty()) {
            return LoadResult("Shader manifest pass is missing a valid name: " + path);
        }

        std::string tagsObject;
        if (TryExtractObject(passObject, "tags", tagsObject)) {
            if (!TryParseStringMapObject(
                    tagsObject,
                    [shaderPtr = shader.get(), &passName](const Containers::String& key, const Containers::String& value) {
                        shaderPtr->SetPassTag(passName, key, value);
                    })) {
                return LoadResult("Shader manifest pass tags could not be parsed: " + path);
            }
        }

        std::string variantsArray;
        if (!TryExtractArray(passObject, "variants", variantsArray)) {
            return LoadResult("Shader manifest pass is missing variants: " + path);
        }

        std::vector<std::string> variantObjects;
        if (!SplitTopLevelArrayElements(variantsArray, variantObjects) || variantObjects.empty()) {
            return LoadResult("Shader manifest pass does not contain any variants: " + path);
        }

        for (const std::string& variantObject : variantObjects) {
            ShaderStageVariant variant = {};

            Containers::String stageName;
            if (!TryParseStringValue(variantObject, "stage", stageName) ||
                !TryParseShaderType(stageName, variant.stage)) {
                return LoadResult("Shader manifest variant has an invalid stage: " + path);
            }

            Containers::String backendName;
            if (!TryParseStringValue(variantObject, "backend", backendName) ||
                !TryParseShaderBackend(backendName, variant.backend)) {
                return LoadResult("Shader manifest variant has an invalid backend: " + path);
            }

            Containers::String languageName;
            if (!TryParseStringValue(variantObject, "language", languageName) ||
                !TryParseShaderLanguage(languageName, variant.language)) {
                return LoadResult("Shader manifest variant has an invalid language: " + path);
            }

            Containers::String sourceCode;
            if (TryParseStringValue(variantObject, "sourceCode", sourceCode)) {
                variant.sourceCode = sourceCode;
            } else {
                Containers::String sourcePath;
                if (!TryParseStringValue(variantObject, "source", sourcePath) &&
                    !TryParseStringValue(variantObject, "sourcePath", sourcePath)) {
                    return LoadResult("Shader manifest variant is missing source/sourceCode: " + path);
                }

                const Containers::String resolvedSourcePath = ResolveShaderDependencyPath(sourcePath, path);
                if (!ReadTextFile(resolvedSourcePath, variant.sourceCode)) {
                    return LoadResult("Failed to read shader variant source: " + resolvedSourcePath);
                }
            }

            if (!TryParseStringValue(variantObject, "entryPoint", variant.entryPoint)) {
                variant.entryPoint = GetDefaultEntryPoint(variant.language, variant.stage);
            }

            if (!TryParseStringValue(variantObject, "profile", variant.profile)) {
                variant.profile = GetDefaultProfile(variant.language, variant.backend, variant.stage);
            }

            shader->AddPassVariant(passName, variant);
        }
    }

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

LoadResult LoadLegacySingleStageShader(const Containers::String& path, const std::string& sourceText) {
    auto shader = std::make_unique<Shader>();
    shader->m_path = path;
    shader->m_name = path;
    shader->m_guid = ResourceGUID::Generate(path);

    const Containers::String ext = GetPathExtension(path).ToLower();
    if (ext == "hlsl") {
        shader->SetShaderLanguage(ShaderLanguage::HLSL);
    } else {
        shader->SetShaderLanguage(ShaderLanguage::GLSL);
    }

    shader->SetShaderType(DetectShaderTypeFromPath(path));
    shader->SetSourceCode(sourceText.c_str());
    shader->m_isValid = true;
    shader->m_memorySize =
        sizeof(Shader) +
        shader->m_name.Length() +
        shader->m_path.Length() +
        shader->GetSourceCode().Length();

    return LoadResult(shader.release());
}

} // namespace

ShaderLoader::ShaderLoader() = default;

ShaderLoader::~ShaderLoader() = default;

Containers::Array<Containers::String> ShaderLoader::GetSupportedExtensions() const {
    Containers::Array<Containers::String> extensions;
    extensions.PushBack("vert");
    extensions.PushBack("frag");
    extensions.PushBack("geom");
    extensions.PushBack("comp");
    extensions.PushBack("glsl");
    extensions.PushBack("hlsl");
    extensions.PushBack("shader");
    return extensions;
}

bool ShaderLoader::CanLoad(const Containers::String& path) const {
    if (IsBuiltinShaderPath(path)) {
        return true;
    }

    const Containers::String ext = GetExtension(path).ToLower();
    return ext == "vert" || ext == "frag" || ext == "geom" ||
           ext == "comp" || ext == "glsl" || ext == "hlsl" || ext == "shader";
}

LoadResult ShaderLoader::Load(const Containers::String& path, const ImportSettings* settings) {
    (void)settings;

    if (IsBuiltinShaderPath(path)) {
        return CreateBuiltinShaderResource(path);
    }

    const Containers::Array<Core::uint8> data = ReadShaderFileData(path);
    if (data.Empty()) {
        return LoadResult("Failed to read shader file: " + path);
    }

    const std::string sourceText = ToStdString(data);
    const Containers::String ext = GetPathExtension(path).ToLower();
    if (ext == "shader" && LooksLikeShaderManifest(sourceText)) {
        return LoadShaderManifest(path, sourceText);
    }

    return LoadLegacySingleStageShader(path, sourceText);
}

ImportSettings* ShaderLoader::GetDefaultSettings() const {
    return nullptr;
}

ShaderType ShaderLoader::DetectShaderType(const Containers::String& path, const Containers::String& source) {
    (void)source;
    return DetectShaderTypeFromPath(path);
}

bool ShaderLoader::ParseShaderSource(const Containers::String& source, Shader* shader) {
    (void)source;
    (void)shader;
    return true;
}

} // namespace Resources
} // namespace XCEngine