XCEngine/MVS/3DGS-D3D12/include/XC3DGSD3D12/App.h

#pragma once

#include <windows.h>

#include <memory>
#include <string>
#include <vector>
#include <wrl/client.h>

#include "XC3DGSD3D12/GaussianPlyLoader.h"
#include "XCEngine/RHI/RHIEnums.h"
#include "XCEngine/RHI/RHITypes.h"
#include "XCEngine/RHI/D3D12/D3D12CommandAllocator.h"
#include "XCEngine/RHI/D3D12/D3D12Buffer.h"
#include "XCEngine/RHI/D3D12/D3D12CommandList.h"
#include "XCEngine/RHI/D3D12/D3D12CommandQueue.h"
#include "XCEngine/RHI/D3D12/D3D12DescriptorHeap.h"
#include "XCEngine/RHI/D3D12/D3D12Device.h"
#include "XCEngine/RHI/D3D12/D3D12ResourceView.h"
#include "XCEngine/RHI/D3D12/D3D12SwapChain.h"
#include "XCEngine/RHI/D3D12/D3D12Texture.h"

namespace XCEngine {
namespace RHI {
class RHIDescriptorPool;
class RHIDescriptorSet;
class RHIPipelineLayout;
class RHIPipelineState;
} // namespace RHI
} // namespace XCEngine

namespace XC3DGSD3D12 {

struct PreparedSplatView {
    float clipPosition[4] = {};
    float axis1[2] = {};
    float axis2[2] = {};
    uint32_t packedColor[2] = {};
};

class App {
public:
    App();
    ~App();

    bool Initialize(HINSTANCE instance, int showCommand);
    int Run();
    void SetFrameLimit(unsigned int frameLimit);
    void SetGaussianScenePath(std::wstring scenePath);
    void SetSummaryPath(std::wstring summaryPath);
    void SetScreenshotPath(std::wstring screenshotPath);
    const std::wstring& GetLastErrorMessage() const;

private:
    static constexpr int kBackBufferCount = 2;
    static constexpr int kDefaultWidth = 1280;
    static constexpr int kDefaultHeight = 720;

    static LRESULT CALLBACK StaticWindowProc(HWND hwnd, UINT message, WPARAM wParam, LPARAM lParam);
    LRESULT WindowProc(HWND hwnd, UINT message, WPARAM wParam, LPARAM lParam);

    bool RegisterWindowClass(HINSTANCE instance);
    bool CreateMainWindow(HINSTANCE instance, int showCommand);
    bool LoadGaussianScene();
    bool InitializeRhi();
    bool InitializeGaussianGpuResources();
    bool InitializePreparePassResources();
    bool InitializeSortResources();
    bool InitializeDebugDrawResources();
    bool InitializeCompositeResources();
    void ShutdownGaussianGpuResources();
    void ShutdownPreparePassResources();
    void ShutdownSortResources();
    void ShutdownDebugDrawResources();
    void ShutdownCompositeResources();
    void Shutdown();
    bool CaptureSortSnapshot();
    bool CapturePass3HistogramDebug();
    void RenderFrame(bool captureScreenshot);

    HWND m_hwnd = nullptr;
    HINSTANCE m_instance = nullptr;
    int m_width = kDefaultWidth;
    int m_height = kDefaultHeight;
    bool m_running = false;
    bool m_isInitialized = false;
    bool m_hasRenderedAtLeastOneFrame = false;
    unsigned int m_frameLimit = 0;
    unsigned int m_renderedFrameCount = 0;
    std::wstring m_gaussianScenePath = L"room.ply";
    std::wstring m_summaryPath;
    std::wstring m_screenshotPath = L"phase3_debug_points.ppm";
    std::wstring m_sortKeySnapshotPath = L"phase3_sortkeys.txt";
    std::wstring m_lastErrorMessage;
    GaussianSplatRuntimeData m_gaussianSceneData;
    XCEngine::RHI::D3D12Buffer m_gaussianPositionBuffer;
    XCEngine::RHI::D3D12Buffer m_gaussianOtherBuffer;
    XCEngine::RHI::D3D12Buffer m_gaussianShBuffer;
    XCEngine::RHI::D3D12Texture m_gaussianColorTexture;
    std::unique_ptr<XCEngine::RHI::D3D12ResourceView> m_gaussianPositionView;
    std::unique_ptr<XCEngine::RHI::D3D12ResourceView> m_gaussianOtherView;
    std::unique_ptr<XCEngine::RHI::D3D12ResourceView> m_gaussianShView;
    std::unique_ptr<XCEngine::RHI::D3D12ResourceView> m_gaussianColorView;
    std::vector<Microsoft::WRL::ComPtr<ID3D12Resource>> m_gaussianUploadBuffers;
    XCEngine::RHI::D3D12Buffer* m_preparedViewBuffer = nullptr;
    std::unique_ptr<XCEngine::RHI::D3D12ResourceView> m_preparedViewSrv;
    std::unique_ptr<XCEngine::RHI::D3D12ResourceView> m_preparedViewUav;
    XCEngine::RHI::RHIPipelineLayout* m_preparePipelineLayout = nullptr;
    XCEngine::RHI::RHIPipelineState* m_preparePipelineState = nullptr;
    XCEngine::RHI::RHIDescriptorPool* m_prepareDescriptorPool = nullptr;
    XCEngine::RHI::RHIDescriptorSet* m_prepareDescriptorSet = nullptr;
    XCEngine::RHI::D3D12Buffer* m_sortKeyBuffer = nullptr;
    XCEngine::RHI::D3D12Buffer* m_sortKeyScratchBuffer = nullptr;
    XCEngine::RHI::D3D12Buffer* m_orderBuffer = nullptr;
    XCEngine::RHI::D3D12Buffer* m_orderScratchBuffer = nullptr;
    XCEngine::RHI::D3D12Buffer* m_passHistogramBuffer = nullptr;
    XCEngine::RHI::D3D12Buffer* m_globalHistogramBuffer = nullptr;
    std::unique_ptr<XCEngine::RHI::D3D12ResourceView> m_sortKeyUav;
    std::unique_ptr<XCEngine::RHI::D3D12ResourceView> m_sortKeyScratchUav;
    std::unique_ptr<XCEngine::RHI::D3D12ResourceView> m_orderBufferSrv;
    std::unique_ptr<XCEngine::RHI::D3D12ResourceView> m_orderBufferUav;
    std::unique_ptr<XCEngine::RHI::D3D12ResourceView> m_orderScratchUav;
    std::unique_ptr<XCEngine::RHI::D3D12ResourceView> m_passHistogramUav;
    std::unique_ptr<XCEngine::RHI::D3D12ResourceView> m_globalHistogramUav;
    XCEngine::RHI::RHIPipelineLayout* m_buildSortKeyPipelineLayout = nullptr;
    XCEngine::RHI::RHIPipelineState* m_buildSortKeyPipelineState = nullptr;
    XCEngine::RHI::RHIDescriptorPool* m_buildSortKeyDescriptorPool = nullptr;
    XCEngine::RHI::RHIDescriptorSet* m_buildSortKeyDescriptorSet = nullptr;
    XCEngine::RHI::RHIPipelineLayout* m_radixSortPipelineLayout = nullptr;
    XCEngine::RHI::RHIPipelineState* m_radixSortInitPipelineState = nullptr;
    XCEngine::RHI::RHIPipelineState* m_radixSortUpsweepPipelineState = nullptr;
    XCEngine::RHI::RHIPipelineState* m_radixSortGlobalHistogramPipelineState = nullptr;
    XCEngine::RHI::RHIPipelineState* m_radixSortScanPipelineState = nullptr;
    XCEngine::RHI::RHIPipelineState* m_radixSortDownsweepPipelineState = nullptr;
    XCEngine::RHI::RHIDescriptorPool* m_radixSortDescriptorPool = nullptr;
    XCEngine::RHI::RHIDescriptorSet* m_radixSortDescriptorSetPrimaryToScratch = nullptr;
    XCEngine::RHI::RHIDescriptorSet* m_radixSortDescriptorSetScratchToPrimary = nullptr;
    XCEngine::RHI::RHIPipelineLayout* m_debugPipelineLayout = nullptr;
    XCEngine::RHI::RHIPipelineState* m_debugPipelineState = nullptr;
    XCEngine::RHI::RHIDescriptorPool* m_debugDescriptorPool = nullptr;
    XCEngine::RHI::RHIDescriptorSet* m_debugDescriptorSet = nullptr;
    XCEngine::RHI::D3D12Texture m_splatRenderTargetTexture;
    std::unique_ptr<XCEngine::RHI::D3D12ResourceView> m_splatRenderTargetRtv;
    std::unique_ptr<XCEngine::RHI::D3D12ResourceView> m_splatRenderTargetSrv;
    XCEngine::RHI::RHIPipelineLayout* m_compositePipelineLayout = nullptr;
    XCEngine::RHI::RHIPipelineState* m_compositePipelineState = nullptr;
    XCEngine::RHI::RHIDescriptorPool* m_compositeDescriptorPool = nullptr;
    XCEngine::RHI::RHIDescriptorSet* m_compositeDescriptorSet = nullptr;

    XCEngine::RHI::D3D12Device m_device;
    XCEngine::RHI::D3D12CommandQueue m_commandQueue;
    XCEngine::RHI::D3D12SwapChain m_swapChain;
    XCEngine::RHI::D3D12CommandAllocator m_commandAllocator;
    XCEngine::RHI::D3D12CommandList m_commandList;
    XCEngine::RHI::D3D12Texture m_depthStencil;
    XCEngine::RHI::D3D12DescriptorHeap m_rtvHeap;
    XCEngine::RHI::D3D12DescriptorHeap m_dsvHeap;
    XCEngine::RHI::D3D12ResourceView m_rtvs[kBackBufferCount];
    XCEngine::RHI::D3D12ResourceView m_dsv;
};

} // namespace XC3DGSD3D12