FactoryCraft-Client/FactoryCraft/DX12ChunkData.cpp

#include "Dx12ChunkData.h"

#include <DX12CommandQueue.h>

#include "Constants.h"

DX12ChunkData::DX12ChunkData(Framework::Point center)
    : ReferenceCounter(),
      aabbs(0),
      chunkInfoBuffer(0),
      blasScratchBuffer(0),
      blasBuffer(0),
      blockIndexBuffer(0),
      textureIdBuffer(0),
      lastObjectIndex(-1),
      lastHitGroupIndex(-1),
      chunkCenter(center),
      blockIndices(new int[CHUNK_SIZE * CHUNK_SIZE * WORLD_HEIGHT]),
      textureIds(),
      changed(0),
      bufferChanged(0),
      minZ(WORLD_HEIGHT),
      maxZ(0),
      minMaxChanged(0)
{
    memset(
        blockIndices, -1, sizeof(int) * CHUNK_SIZE * CHUNK_SIZE * WORLD_HEIGHT);
}

DX12ChunkData::~DX12ChunkData()
{
    if (aabbs)
    {
        aabbs->release();
    }
    if (chunkInfoBuffer)
    {
        chunkInfoBuffer->release();
    }
    if (blasScratchBuffer)
    {
        blasScratchBuffer->release();
    }
    if (blasBuffer)
    {
        blasBuffer->release();
    }
    if (blockIndexBuffer)
    {
        blockIndexBuffer->release();
    }
    if (textureIdBuffer)
    {
        textureIdBuffer->release();
    }
    delete[] blockIndices;
}

Framework::DX12Buffer* DX12ChunkData::zChunkInfoBuffer() const
{
    return chunkInfoBuffer;
}

Framework::DX12Buffer* DX12ChunkData::zBlasBuffer() const
{
    return blasBuffer;
}

Framework::DX12Buffer* DX12ChunkData::zBlockIndexBuffer() const
{
    return blockIndexBuffer;
}

Framework::DX12Buffer* DX12ChunkData::zTextureIdBuffer() const
{
    return textureIdBuffer;
}

void DX12ChunkData::setBlock(int index, Block* zBlock)
{
    cs.lock();
    if (zBlock)
    {
        if (blockIndices[index] < 0)
        {
            blockIndices[index] = textureIds.getEntryCount();
            for (int i = 0; i < 6; ++i)
            {
                textureIds.add(zBlock->zTexture()->zPolygonTexture(i)->getId());
            }
        }
        else
        {
            ArrayIterator<int> textureIt = textureIds.begin();
            int tmp = blockIndices[index];
            while (tmp > 0)
            {
                ++textureIt;
                --tmp;
            }
            for (int i = 0; i < 6; ++i)
            {
                textureIt.set(zBlock->zTexture()->zPolygonTexture(i)->getId());
                ++textureIt;
            }
        }
        if (zBlock->getLocation().z + 1 > maxZ)
        {
            maxZ = zBlock->getLocation().z + 1;
            minMaxChanged = 1;
        }
        if (zBlock->getLocation().z < minZ)
        {
            minZ = zBlock->getLocation().z;
            minMaxChanged = 1;
        }
    }
    else
    {
        if (blockIndices[index] >= 0)
        {
            int tmp = blockIndices[index];
            ArrayIterator<int> textureIt = textureIds.begin();
            while (tmp > 0)
            {
                ++textureIt;
                --tmp;
            }
            for (int i = 0; i < 6; ++i)
            {
                textureIt.remove();
            }
            int oldMinZ = minZ;
            int oldMaxZ = maxZ;
            maxZ = 0;
            minZ = WORLD_HEIGHT;
            for (int i = 0; i < CHUNK_SIZE * CHUNK_SIZE * WORLD_HEIGHT; ++i)
            {
                if (blockIndices[i] > blockIndices[index])
                {
                    blockIndices[i] -= 6;
                }
                if (blockIndices[i] >= 0)
                {
                    if (i % WORLD_HEIGHT + 1 > maxZ)
                    {
                        maxZ = i % WORLD_HEIGHT + 1;
                    }
                    if (i % WORLD_HEIGHT < minZ)
                    {
                        minZ = i % WORLD_HEIGHT;
                    }
                }
            }
            if (minZ != oldMinZ || maxZ != oldMaxZ)
            {
                minMaxChanged = 1;
            }
        }
        blockIndices[index] = -1;
    }
    changed = 1;
    cs.unlock();
}

void DX12ChunkData::updateBuffers(
    ID3D12Device5* zDevice, Framework::DX12CommandQueue* zQueue)
{
    cs.lock();
    if (!aabbs)
    {
        aabbs = new DX12Buffer(sizeof(D3D12_RAYTRACING_AABB),
            zDevice,
            dynamic_cast<Framework::DX12CommandQueue*>(zQueue->getThis()),
            D3D12_RESOURCE_FLAG_NONE);
        aabbs->setLength(sizeof(D3D12_RAYTRACING_AABB));
        D3D12_RAYTRACING_AABB data = {chunkCenter.x - 8.f,
            chunkCenter.y - 8.f,
            minZ,
            chunkCenter.x + 8.f,
            chunkCenter.y + 8.f,
            maxZ};
        aabbs->setData(&data, 1);
        aabbs->copyToGPU();
        minMaxChanged = 0;
    }
    if (minMaxChanged)
    {
        D3D12_RAYTRACING_AABB data = {chunkCenter.x - 8.f,
            chunkCenter.y - 8.f,
            minZ,
            chunkCenter.x + 8.f,
            chunkCenter.y + 8.f,
            maxZ};
        aabbs->setData(&data, 1);
        aabbs->copyToGPU();
        minMaxChanged = 0;
    }
    if (!blasBuffer)
    {
        D3D12_RAYTRACING_GEOMETRY_DESC geometryDesc = {};
        geometryDesc.Type
            = D3D12_RAYTRACING_GEOMETRY_TYPE_PROCEDURAL_PRIMITIVE_AABBS;
        geometryDesc.Flags = D3D12_RAYTRACING_GEOMETRY_FLAG_NONE;
        geometryDesc.AABBs.AABBCount = 1;
        geometryDesc.AABBs.AABBs.StartAddress
            = aabbs->zBuffer()->GetGPUVirtualAddress();
        geometryDesc.AABBs.AABBs.StrideInBytes = sizeof(D3D12_RAYTRACING_AABB);

        D3D12_BUILD_RAYTRACING_ACCELERATION_STRUCTURE_INPUTS prebuildDesc;
        prebuildDesc.Type
            = D3D12_RAYTRACING_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL;
        prebuildDesc.DescsLayout = D3D12_ELEMENTS_LAYOUT_ARRAY;
        prebuildDesc.NumDescs = 1;
        prebuildDesc.pGeometryDescs = &geometryDesc;
        prebuildDesc.Flags
            = D3D12_RAYTRACING_ACCELERATION_STRUCTURE_BUILD_FLAG_ALLOW_UPDATE;

        D3D12_RAYTRACING_ACCELERATION_STRUCTURE_PREBUILD_INFO info = {};
        zDevice->GetRaytracingAccelerationStructurePrebuildInfo(
            &prebuildDesc, &info);
        ID3D12Resource* oldBuffer = 0;
        if (blasBuffer)
        {
            oldBuffer = blasBuffer->zBuffer();
            oldBuffer->AddRef();
            blasBuffer->release();
        }
        blasBuffer = new DX12Buffer(1,
            zDevice,
            dynamic_cast<DX12CommandQueue*>(zQueue->getThis()),
            D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS);
        blasBuffer->setLength(
            ROUND_UP_POWER_OF_2(info.ResultDataMaxSizeInBytes, 256));
        blasBuffer->createBufferWithoutData(
            D3D12_RESOURCE_STATE_RAYTRACING_ACCELERATION_STRUCTURE);

        blasScratchBuffer = new DX12Buffer(1,
            zDevice,
            dynamic_cast<DX12CommandQueue*>(zQueue->getThis()),
            D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS);
        blasScratchBuffer->setLength(
            ROUND_UP_POWER_OF_2(info.ScratchDataSizeInBytes, 256));
        blasScratchBuffer->createBufferWithoutData(D3D12_RESOURCE_STATE_COMMON);

        D3D12_BUILD_RAYTRACING_ACCELERATION_STRUCTURE_DESC buildDesc;
        buildDesc.Inputs.Type
            = D3D12_RAYTRACING_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL;
        buildDesc.Inputs.DescsLayout = D3D12_ELEMENTS_LAYOUT_ARRAY;
        buildDesc.Inputs.NumDescs = 1;
        buildDesc.Inputs.pGeometryDescs = &geometryDesc;
        buildDesc.DestAccelerationStructureData
            = {blasBuffer->zBuffer()->GetGPUVirtualAddress()};
        buildDesc.ScratchAccelerationStructureData
            = {blasScratchBuffer->zBuffer()->GetGPUVirtualAddress()};
        buildDesc.SourceAccelerationStructureData = oldBuffer ? oldBuffer->GetGPUVirtualAddress() : 0;
        buildDesc.Inputs.Flags
            = oldBuffer
                ? D3D12_RAYTRACING_ACCELERATION_STRUCTURE_BUILD_FLAG_PERFORM_UPDATE
                : D3D12_RAYTRACING_ACCELERATION_STRUCTURE_BUILD_FLAG_ALLOW_UPDATE;

        zQueue->zCommandList()->BuildRaytracingAccelerationStructure(
            &buildDesc, 0, nullptr);
        if (oldBuffer)
        {
            oldBuffer->Release();
        }
        bufferChanged = 1;
    }
    if (!blockIndexBuffer)
    {
        blockIndexBuffer = new DX12Buffer(sizeof(int),
            zDevice,
            dynamic_cast<DX12CommandQueue*>(zQueue->getThis()),
            D3D12_RESOURCE_FLAG_NONE);
        blockIndexBuffer->setLength(
            sizeof(int) * CHUNK_SIZE * CHUNK_SIZE * WORLD_HEIGHT);
        bufferChanged = 1;
    }
    if (textureIdBuffer
        && textureIdBuffer->getElementCount() < textureIds.getEntryCount())
    {
        textureIdBuffer->release();
        textureIdBuffer = 0;
    }
    if (!textureIdBuffer)
    {
        textureIdBuffer = new DX12Buffer(sizeof(int),
            zDevice,
            dynamic_cast<DX12CommandQueue*>(zQueue->getThis()),
            D3D12_RESOURCE_FLAG_NONE);
        textureIdBuffer->setLength(textureIds.getEntryCount() * sizeof(int));
        bufferChanged = 1;
    }
    if (!chunkInfoBuffer)
    {
        chunkInfoBuffer = new DX12Buffer(sizeof(ChunkShaderInfo),
            zDevice,
            dynamic_cast<DX12CommandQueue*>(zQueue->getThis()),
            D3D12_RESOURCE_FLAG_NONE);
        chunkInfoBuffer->setLength(sizeof(ChunkShaderInfo));
    }
    if (changed)
    {
        for (int i = 0; i < CHUNK_SIZE * CHUNK_SIZE * WORLD_HEIGHT; ++i)
        {
            assert(blockIndices[i] < 0
                   || blockIndices[i] <= textureIds.getEntryCount() - 6);
        }
        int* textureIdData = new int[textureIds.getEntryCount()];
        ArrayIterator<int> textureIt = textureIds.begin();
        int index = 0;
        while (textureIt)
        {
            textureIdData[index++] = textureIt.val();
            ++textureIt;
        }
        blockIndexBuffer->setData(blockIndices, 1);
        textureIdBuffer->setData(textureIdData, 1);
        blockIndexBuffer->copyToGPU();
        textureIdBuffer->copyToGPU(sizeof(int) * textureIds.getEntryCount());
        delete[] textureIdData;
        ChunkShaderInfo info
            = {chunkCenter, (unsigned int)textureIds.getEntryCount()};
        chunkInfoBuffer->setData(&info, 1);
        chunkInfoBuffer->copyToGPU();
        changed = 0;
    }
    cs.unlock();
}

int DX12ChunkData::getLastObjectIndex() const
{
    return lastObjectIndex;
}

int DX12ChunkData::getLastHitGroupIndex() const
{
    return lastHitGroupIndex;
}

void DX12ChunkData::setLastObjectIndex(int index)
{
    lastObjectIndex = index;
}

void DX12ChunkData::setLastHitGroupIndex(int index)
{
    lastHitGroupIndex = index;
}

bool DX12ChunkData::wasBufferChanged() const
{
    return bufferChanged;
}

void DX12ChunkData::setBufferChanged(bool changed)
{
    bufferChanged = changed;
}