XCEngine/tests/Resources/GaussianSplat/test_gaussian_splat_loader.cpp

#include <gtest/gtest.h>

#include <XCEngine/Core/Asset/AssetDatabase.h>
#include <XCEngine/Core/Asset/AssetRef.h>
#include <XCEngine/Core/Asset/ResourceManager.h>
#include <XCEngine/Core/Math/Bounds.h>
#include <XCEngine/Resources/GaussianSplat/GaussianSplat.h>
#include <XCEngine/Resources/GaussianSplat/GaussianSplatArtifactIO.h>
#include <XCEngine/Resources/GaussianSplat/GaussianSplatLoader.h>

#include <chrono>
#include <cmath>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <thread>
#include <vector>

using namespace XCEngine::Resources;
using namespace XCEngine::Math;

namespace {

constexpr float kSHC0 = 0.2820948f;

struct SyntheticGaussianSplatVertex {
    Vector3 position = Vector3::Zero();
    Vector3 dc0 = Vector3::Zero();
    float opacity = 0.0f;
    Vector3 scaleLog = Vector3::Zero();
    float rotationWXYZ[4] = { 1.0f, 0.0f, 0.0f, 0.0f };
    float sh[kGaussianSplatSHCoefficientCount] = {};
};

struct SampleArtifactData {
    GaussianSplatMetadata metadata;
    XCEngine::Containers::Array<GaussianSplatSection> sections;
    XCEngine::Containers::Array<XCEngine::Core::uint8> payload;
};

std::filesystem::path GetRoomPlyPath() {
    return std::filesystem::path(XCENGINE_TEST_ROOM_PLY_PATH);
}

void ExpectVector3Near(const Vector3& actual, const Vector3& expected, float epsilon = 1e-5f) {
    EXPECT_NEAR(actual.x, expected.x, epsilon);
    EXPECT_NEAR(actual.y, expected.y, epsilon);
    EXPECT_NEAR(actual.z, expected.z, epsilon);
}

void ExpectVector4Near(const Vector4& actual, const Vector4& expected, float epsilon = 1e-5f) {
    EXPECT_NEAR(actual.x, expected.x, epsilon);
    EXPECT_NEAR(actual.y, expected.y, epsilon);
    EXPECT_NEAR(actual.z, expected.z, epsilon);
    EXPECT_NEAR(actual.w, expected.w, epsilon);
}

void ExpectQuaternionNear(const Quaternion& actual, const Quaternion& expected, float epsilon = 1e-5f) {
    EXPECT_NEAR(actual.x, expected.x, epsilon);
    EXPECT_NEAR(actual.y, expected.y, epsilon);
    EXPECT_NEAR(actual.z, expected.z, epsilon);
    EXPECT_NEAR(actual.w, expected.w, epsilon);
}

float Sigmoid(float value) {
    return 1.0f / (1.0f + std::exp(-value));
}

Vector3 LinearScale(const Vector3& value) {
    return Vector3(
        std::abs(std::exp(value.x)),
        std::abs(std::exp(value.y)),
        std::abs(std::exp(value.z)));
}

Quaternion NormalizeRotationWXYZ(float w, float x, float y, float z) {
    const float magnitude = std::sqrt(w * w + x * x + y * y + z * z);
    if (magnitude <= 1e-8f) {
        return Quaternion::Identity();
    }

    return Quaternion(x / magnitude, y / magnitude, z / magnitude, w / magnitude);
}

Vector4 SH0ToColorOpacity(const Vector3& dc0, float opacityRaw) {
    return Vector4(
        dc0.x * kSHC0 + 0.5f,
        dc0.y * kSHC0 + 0.5f,
        dc0.z * kSHC0 + 0.5f,
        Sigmoid(opacityRaw));
}

void WriteBinaryFloat(std::ofstream& output, float value) {
    output.write(reinterpret_cast<const char*>(&value), sizeof(value));
}

void WriteSyntheticGaussianSplatPly(
    const std::filesystem::path& path,
    const std::vector<SyntheticGaussianSplatVertex>& vertices) {
    std::ofstream output(path, std::ios::binary | std::ios::trunc);
    ASSERT_TRUE(output.is_open());

    output << "ply\n";
    output << "format binary_little_endian 1.0\n";
    output << "element vertex " << vertices.size() << "\n";
    output << "property float opacity\n";
    output << "property float y\n";
    output << "property float scale_2\n";
    output << "property float rot_3\n";
    output << "property float f_dc_1\n";
    output << "property float x\n";
    output << "property float scale_0\n";
    output << "property float rot_1\n";
    output << "property float f_dc_2\n";
    output << "property float z\n";
    output << "property float scale_1\n";
    output << "property float rot_0\n";
    output << "property float f_dc_0\n";
    output << "property float rot_2\n";
    for (XCEngine::Core::uint32 index = 0; index < kGaussianSplatSHCoefficientCount; ++index) {
        output << "property float f_rest_" << index << "\n";
    }
    output << "end_header\n";

    for (const SyntheticGaussianSplatVertex& vertex : vertices) {
        WriteBinaryFloat(output, vertex.opacity);
        WriteBinaryFloat(output, vertex.position.y);
        WriteBinaryFloat(output, vertex.scaleLog.z);
        WriteBinaryFloat(output, vertex.rotationWXYZ[3]);
        WriteBinaryFloat(output, vertex.dc0.y);
        WriteBinaryFloat(output, vertex.position.x);
        WriteBinaryFloat(output, vertex.scaleLog.x);
        WriteBinaryFloat(output, vertex.rotationWXYZ[1]);
        WriteBinaryFloat(output, vertex.dc0.z);
        WriteBinaryFloat(output, vertex.position.z);
        WriteBinaryFloat(output, vertex.scaleLog.y);
        WriteBinaryFloat(output, vertex.rotationWXYZ[0]);
        WriteBinaryFloat(output, vertex.dc0.x);
        WriteBinaryFloat(output, vertex.rotationWXYZ[2]);
        for (float coefficient : vertex.sh) {
            WriteBinaryFloat(output, coefficient);
        }
    }
}

SampleArtifactData BuildSampleArtifactData() {
    const GaussianSplatPositionRecord positions[2] = {
        { Vector3(0.0f, 1.0f, 2.0f) },
        { Vector3(3.0f, 4.0f, 5.0f) }
    };
    const GaussianSplatOtherRecord other[2] = {
        { Quaternion::Identity(), Vector3(1.0f, 1.0f, 1.0f), 0.0f },
        { Quaternion(0.0f, 0.5f, 0.0f, 0.8660254f), Vector3(2.0f, 2.0f, 2.0f), 0.0f }
    };
    const GaussianSplatColorRecord colors[2] = {
        { Vector4(1.0f, 0.0f, 0.0f, 0.25f) },
        { Vector4(0.0f, 1.0f, 0.0f, 0.75f) }
    };

    GaussianSplatSHRecord sh[2];
    for (XCEngine::Core::uint32 index = 0; index < kGaussianSplatSHCoefficientCount; ++index) {
        sh[0].coefficients[index] = 0.01f * static_cast<float>(index + 1u);
        sh[1].coefficients[index] = -0.02f * static_cast<float>(index + 1u);
    }

    SampleArtifactData sample;
    sample.metadata.contentVersion = 3u;
    sample.metadata.splatCount = 2u;
    sample.metadata.bounds.SetMinMax(Vector3(-2.0f, -1.0f, -3.0f), Vector3(5.0f, 4.0f, 6.0f));
    sample.metadata.positionFormat = GaussianSplatSectionFormat::VectorFloat32;
    sample.metadata.otherFormat = GaussianSplatSectionFormat::OtherFloat32;
    sample.metadata.colorFormat = GaussianSplatSectionFormat::ColorRGBA32F;
    sample.metadata.shFormat = GaussianSplatSectionFormat::SHFloat32;

    sample.sections.Reserve(4u);
    size_t payloadOffset = 0u;
    auto appendSection = [&](GaussianSplatSectionType type,
                             GaussianSplatSectionFormat format,
                             const void* data,
                             size_t dataSize,
                             XCEngine::Core::uint32 elementCount,
                             XCEngine::Core::uint32 elementStride) {
        GaussianSplatSection section;
        section.type = type;
        section.format = format;
        section.dataOffset = payloadOffset;
        section.dataSize = dataSize;
        section.elementCount = elementCount;
        section.elementStride = elementStride;
        sample.sections.PushBack(section);

        const size_t newPayloadSize = sample.payload.Size() + dataSize;
        sample.payload.Resize(newPayloadSize);
        std::memcpy(sample.payload.Data() + payloadOffset, data, dataSize);
        payloadOffset = newPayloadSize;
    };

    appendSection(
        GaussianSplatSectionType::Positions,
        GaussianSplatSectionFormat::VectorFloat32,
        positions,
        sizeof(positions),
        2u,
        sizeof(GaussianSplatPositionRecord));
    appendSection(
        GaussianSplatSectionType::Other,
        GaussianSplatSectionFormat::OtherFloat32,
        other,
        sizeof(other),
        2u,
        sizeof(GaussianSplatOtherRecord));
    appendSection(
        GaussianSplatSectionType::Color,
        GaussianSplatSectionFormat::ColorRGBA32F,
        colors,
        sizeof(colors),
        2u,
        sizeof(GaussianSplatColorRecord));
    appendSection(
        GaussianSplatSectionType::SH,
        GaussianSplatSectionFormat::SHFloat32,
        sh,
        sizeof(sh),
        2u,
        sizeof(GaussianSplatSHRecord));

    return sample;
}

GaussianSplat BuildSampleGaussianSplat(const char* artifactPath) {
    SampleArtifactData sample = BuildSampleArtifactData();

    GaussianSplat gaussianSplat;
    XCEngine::Resources::IResource::ConstructParams params;
    params.name = "sample.xcgsplat";
    params.path = artifactPath;
    params.guid = ResourceGUID::Generate(params.path);
    gaussianSplat.Initialize(params);
    const bool created = gaussianSplat.CreateOwned(
        sample.metadata,
        std::move(sample.sections),
        std::move(sample.payload));
    EXPECT_TRUE(created);
    return gaussianSplat;
}

std::filesystem::path CreateTestProjectRoot(const char* folderName) {
    return std::filesystem::current_path() / "__xc_gaussian_splat_test_runtime" / folderName;
}

void LinkOrCopyFixture(const std::filesystem::path& sourcePath, const std::filesystem::path& destinationPath) {
    std::error_code ec;
    std::filesystem::create_hard_link(sourcePath, destinationPath, ec);
    if (ec) {
        ec.clear();
        std::filesystem::copy_file(sourcePath, destinationPath, std::filesystem::copy_options::overwrite_existing, ec);
    }
    ASSERT_FALSE(ec);
}

XCEngine::Core::uint32 ReadPlyVertexCount(const std::filesystem::path& path) {
    std::ifstream input(path, std::ios::binary);
    EXPECT_TRUE(input.is_open());

    std::string line;
    while (std::getline(input, line)) {
        if (line == "end_header") {
            break;
        }

        if (line.rfind("element vertex ", 0) == 0) {
            return static_cast<XCEngine::Core::uint32>(std::stoul(line.substr(15)));
        }
    }

    return 0u;
}

TEST(GaussianSplatLoader, GetResourceType) {
    GaussianSplatLoader loader;
    EXPECT_EQ(loader.GetResourceType(), ResourceType::GaussianSplat);
}

TEST(GaussianSplatLoader, CanLoad) {
    GaussianSplatLoader loader;
    EXPECT_TRUE(loader.CanLoad("sample.xcgsplat"));
    EXPECT_TRUE(loader.CanLoad("sample.XCGSPLAT"));
    EXPECT_FALSE(loader.CanLoad("sample.ply"));
}

TEST(GaussianSplatLoader, LoadInvalidPath) {
    GaussianSplatLoader loader;
    const LoadResult result = loader.Load("invalid/path/sample.xcgsplat");
    EXPECT_FALSE(result);
}

TEST(GaussianSplatLoader, WritesAndLoadsArtifact) {
    namespace fs = std::filesystem;

    const fs::path tempDir = fs::temp_directory_path() / "xc_gaussian_splat_artifact_test";
    const fs::path artifactPath = tempDir / "sample.xcgsplat";

    fs::remove_all(tempDir);
    fs::create_directories(tempDir);

    const GaussianSplat source = BuildSampleGaussianSplat(artifactPath.string().c_str());
    XCEngine::Containers::String errorMessage;
    ASSERT_TRUE(WriteGaussianSplatArtifactFile(artifactPath.string().c_str(), source, &errorMessage))
        << errorMessage.CStr();

    GaussianSplatLoader loader;
    const LoadResult result = loader.Load(artifactPath.string().c_str());
    ASSERT_TRUE(result);
    ASSERT_NE(result.resource, nullptr);

    auto* gaussianSplat = static_cast<GaussianSplat*>(result.resource);
    ASSERT_NE(gaussianSplat, nullptr);
    EXPECT_EQ(gaussianSplat->GetContentVersion(), 3u);
    EXPECT_EQ(gaussianSplat->GetSplatCount(), 2u);
    EXPECT_EQ(gaussianSplat->GetBounds().GetMin(), Vector3(-2.0f, -1.0f, -3.0f));
    EXPECT_EQ(gaussianSplat->GetBounds().GetMax(), Vector3(5.0f, 4.0f, 6.0f));
    ASSERT_EQ(gaussianSplat->GetSections().Size(), 4u);
    const GaussianSplatSection* shSection = gaussianSplat->FindSection(GaussianSplatSectionType::SH);
    ASSERT_NE(shSection, nullptr);
    EXPECT_EQ(shSection->elementCount, 2u);
    EXPECT_EQ(shSection->elementStride, sizeof(GaussianSplatSHRecord));
    ASSERT_NE(gaussianSplat->GetColorRecords(), nullptr);
    EXPECT_EQ(gaussianSplat->GetColorRecords()[1].colorOpacity, Vector4(0.0f, 1.0f, 0.0f, 0.75f));

    delete gaussianSplat;
    fs::remove_all(tempDir);
}

TEST(GaussianSplatLoader, ResourceManagerRegistersGaussianSplatLoader) {
    ResourceManager& manager = ResourceManager::Get();
    manager.Initialize();

    EXPECT_NE(manager.GetLoader(ResourceType::GaussianSplat), nullptr);

    manager.Shutdown();
}

TEST(GaussianSplatLoader, ResourceManagerLoadsArtifactByPath) {
    namespace fs = std::filesystem;

    const fs::path tempDir = fs::temp_directory_path() / "xc_gaussian_splat_manager_load_test";
    const fs::path artifactPath = tempDir / "sample.xcgsplat";

    fs::remove_all(tempDir);
    fs::create_directories(tempDir);

    const GaussianSplat source = BuildSampleGaussianSplat(artifactPath.string().c_str());
    XCEngine::Containers::String errorMessage;
    ASSERT_TRUE(WriteGaussianSplatArtifactFile(artifactPath.string().c_str(), source, &errorMessage))
        << errorMessage.CStr();

    ResourceManager& manager = ResourceManager::Get();
    manager.Initialize();

    {
        const auto handle = manager.Load<GaussianSplat>(artifactPath.string().c_str());
        ASSERT_TRUE(handle.IsValid());
        EXPECT_EQ(handle->GetSplatCount(), 2u);
        EXPECT_EQ(handle->GetContentVersion(), 3u);
        ASSERT_NE(handle->FindSection(GaussianSplatSectionType::Color), nullptr);
    }

    manager.UnloadAll();
    manager.Shutdown();
    fs::remove_all(tempDir);
}

TEST(GaussianSplatLoader, AssetDatabaseImportsSyntheticPlyAndLinearizesData) {
    namespace fs = std::filesystem;

    const fs::path projectRoot = CreateTestProjectRoot("gaussian_splat_synthetic_import");
    const fs::path assetsDir = projectRoot / "Assets";
    const fs::path sourcePath = assetsDir / "sample.ply";

    fs::remove_all(projectRoot);
    fs::create_directories(assetsDir);

    std::vector<SyntheticGaussianSplatVertex> vertices(2);
    vertices[0].position = Vector3(1.0f, 2.0f, 3.0f);
    vertices[0].dc0 = Vector3(0.2f, -0.1f, 0.0f);
    vertices[0].opacity = 0.25f;
    vertices[0].scaleLog = Vector3(0.0f, std::log(2.0f), std::log(4.0f));
    vertices[0].rotationWXYZ[0] = 2.0f;
    vertices[0].rotationWXYZ[1] = 0.0f;
    vertices[0].rotationWXYZ[2] = 0.0f;
    vertices[0].rotationWXYZ[3] = 0.0f;
    for (XCEngine::Core::uint32 index = 0; index < kGaussianSplatSHCoefficientCount; ++index) {
        vertices[0].sh[index] = 0.01f * static_cast<float>(index + 1u);
    }

    vertices[1].position = Vector3(-4.0f, -5.0f, -6.0f);
    vertices[1].dc0 = Vector3(1.0f, 0.5f, -0.5f);
    vertices[1].opacity = -1.0f;
    vertices[1].scaleLog = Vector3(std::log(0.5f), 0.0f, std::log(3.0f));
    vertices[1].rotationWXYZ[0] = 0.0f;
    vertices[1].rotationWXYZ[1] = 0.0f;
    vertices[1].rotationWXYZ[2] = 3.0f;
    vertices[1].rotationWXYZ[3] = 4.0f;
    for (XCEngine::Core::uint32 index = 0; index < kGaussianSplatSHCoefficientCount; ++index) {
        vertices[1].sh[index] = -0.02f * static_cast<float>(index + 1u);
    }

    WriteSyntheticGaussianSplatPly(sourcePath, vertices);

    AssetDatabase database;
    database.Initialize(projectRoot.string().c_str());

    AssetDatabase::ResolvedAsset resolved;
    ASSERT_TRUE(database.EnsureArtifact("Assets/sample.ply", ResourceType::GaussianSplat, resolved));
    ASSERT_TRUE(resolved.artifactReady);
    EXPECT_TRUE(fs::exists(resolved.artifactMainPath.CStr()));
    EXPECT_EQ(fs::path(resolved.artifactMainPath.CStr()).extension().generic_string(), ".xcgsplat");

    GaussianSplatLoader loader;
    LoadResult result = loader.Load(resolved.artifactMainPath);
    ASSERT_TRUE(result);
    ASSERT_NE(result.resource, nullptr);

    auto* gaussianSplat = static_cast<GaussianSplat*>(result.resource);
    ASSERT_NE(gaussianSplat, nullptr);
    EXPECT_EQ(gaussianSplat->GetSplatCount(), 2u);
    EXPECT_EQ(gaussianSplat->GetChunkCount(), 0u);
    EXPECT_EQ(gaussianSplat->GetCameraCount(), 0u);
    ExpectVector3Near(gaussianSplat->GetBounds().GetMin(), Vector3(-4.0f, -5.0f, -6.0f));
    ExpectVector3Near(gaussianSplat->GetBounds().GetMax(), Vector3(1.0f, 2.0f, 3.0f));

    ASSERT_NE(gaussianSplat->GetPositionRecords(), nullptr);
    ASSERT_NE(gaussianSplat->GetOtherRecords(), nullptr);
    ASSERT_NE(gaussianSplat->GetColorRecords(), nullptr);
    ASSERT_NE(gaussianSplat->GetSHRecords(), nullptr);

    ExpectVector3Near(gaussianSplat->GetPositionRecords()[0].position, vertices[0].position);
    ExpectVector3Near(gaussianSplat->GetPositionRecords()[1].position, vertices[1].position);
    ExpectQuaternionNear(
        gaussianSplat->GetOtherRecords()[0].rotation,
        NormalizeRotationWXYZ(2.0f, 0.0f, 0.0f, 0.0f));
    ExpectQuaternionNear(
        gaussianSplat->GetOtherRecords()[1].rotation,
        NormalizeRotationWXYZ(0.0f, 0.0f, 3.0f, 4.0f));
    ExpectVector3Near(gaussianSplat->GetOtherRecords()[0].scale, LinearScale(vertices[0].scaleLog));
    ExpectVector3Near(gaussianSplat->GetOtherRecords()[1].scale, LinearScale(vertices[1].scaleLog));
    ExpectVector4Near(gaussianSplat->GetColorRecords()[0].colorOpacity, SH0ToColorOpacity(vertices[0].dc0, vertices[0].opacity));
    ExpectVector4Near(gaussianSplat->GetColorRecords()[1].colorOpacity, SH0ToColorOpacity(vertices[1].dc0, vertices[1].opacity));
    EXPECT_NEAR(gaussianSplat->GetSHRecords()[0].coefficients[0], vertices[0].sh[0], 1e-6f);
    EXPECT_NEAR(gaussianSplat->GetSHRecords()[0].coefficients[44], vertices[0].sh[44], 1e-6f);
    EXPECT_NEAR(gaussianSplat->GetSHRecords()[1].coefficients[0], vertices[1].sh[0], 1e-6f);
    EXPECT_NEAR(gaussianSplat->GetSHRecords()[1].coefficients[44], vertices[1].sh[44], 1e-6f);

    delete gaussianSplat;
    database.Shutdown();
    fs::remove_all(projectRoot);
}

TEST(GaussianSplatLoader, RoomPlyBuildsArtifactAndLoadsThroughResourceManager) {
    namespace fs = std::filesystem;
    using namespace std::chrono_literals;

    const fs::path fixturePath = GetRoomPlyPath();
    ASSERT_TRUE(fs::exists(fixturePath));

    const XCEngine::Core::uint32 expectedVertexCount = ReadPlyVertexCount(fixturePath);
    ASSERT_GT(expectedVertexCount, 0u);

    const fs::path projectRoot = CreateTestProjectRoot("gaussian_splat_room_import");
    const fs::path assetsDir = projectRoot / "Assets";
    const fs::path roomPath = assetsDir / "room.ply";

    fs::remove_all(projectRoot);
    fs::create_directories(assetsDir);
    LinkOrCopyFixture(fixturePath, roomPath);

    ResourceManager& manager = ResourceManager::Get();
    manager.Initialize();
    manager.SetResourceRoot(projectRoot.string().c_str());

    {
        const auto firstHandle = manager.Load<GaussianSplat>("Assets/room.ply");
        ASSERT_TRUE(firstHandle.IsValid());
        EXPECT_EQ(firstHandle->GetSplatCount(), expectedVertexCount);
        EXPECT_EQ(firstHandle->GetPositionFormat(), GaussianSplatSectionFormat::VectorFloat32);
        EXPECT_EQ(firstHandle->GetOtherFormat(), GaussianSplatSectionFormat::OtherFloat32);
        EXPECT_EQ(firstHandle->GetColorFormat(), GaussianSplatSectionFormat::ColorRGBA32F);
        EXPECT_EQ(firstHandle->GetSHFormat(), GaussianSplatSectionFormat::SHFloat32);

        AssetRef assetRef;
        ASSERT_TRUE(manager.TryGetAssetRef("Assets/room.ply", ResourceType::GaussianSplat, assetRef));
        EXPECT_TRUE(assetRef.IsValid());

        AssetDatabase database;
        database.Initialize(projectRoot.string().c_str());

        AssetDatabase::ResolvedAsset firstResolve;
        ASSERT_TRUE(database.EnsureArtifact("Assets/room.ply", ResourceType::GaussianSplat, firstResolve));
        ASSERT_TRUE(firstResolve.artifactReady);
        EXPECT_TRUE(fs::exists(firstResolve.artifactMainPath.CStr()));
        EXPECT_EQ(fs::path(firstResolve.artifactMainPath.CStr()).extension().generic_string(), ".xcgsplat");
        const auto originalArtifactWriteTime = fs::last_write_time(firstResolve.artifactMainPath.CStr());

        std::this_thread::sleep_for(50ms);

        AssetDatabase::ResolvedAsset secondResolve;
        ASSERT_TRUE(database.EnsureArtifact("Assets/room.ply", ResourceType::GaussianSplat, secondResolve));
        EXPECT_EQ(firstResolve.artifactMainPath, secondResolve.artifactMainPath);
        EXPECT_EQ(originalArtifactWriteTime, fs::last_write_time(secondResolve.artifactMainPath.CStr()));

        database.Shutdown();

        manager.UnloadAll();

        const auto secondHandle = manager.Load<GaussianSplat>(assetRef);
        ASSERT_TRUE(secondHandle.IsValid());
        EXPECT_EQ(secondHandle->GetSplatCount(), expectedVertexCount);
        ASSERT_NE(secondHandle->FindSection(GaussianSplatSectionType::SH), nullptr);
    }

    manager.SetResourceRoot("");
    manager.Shutdown();
    fs::remove_all(projectRoot);
}

} // namespace