vulkan-project/texture_atlas.cpp

#include "texture_atlas.hpp"

#include "vulkan_instance.hpp"
#include "exceptions.hpp"

#include "stb_image.h"

#include <cstdint>
#include <cstdio>
#include <cstring>
#include <gz-util/exceptions.hpp>
#include <gz-util/util/string_concepts.hpp>
#include <gz-util/util/string_conversion.hpp>
#include <ratio>
#include <vulkan/vulkan_core.h>
#include <glm/glm.hpp>

namespace gz::vk {
std::string TextureImageArea::toString() const {
    return "<(" + gz::toString(x) + "," + gz::toString(y) + "), (" + gz::toString(width) + "x" + gz::toString(height) + ")>";
}


//
// INIT & CLEANUP
//
TextureAtlas::TextureAtlas(VulkanInstance& instance, uint16_t slotWidth, uint16_t slotHeight, uint16_t slotCountX, uint16_t slotCountY)
    : vk(instance), slotWidth(slotWidth), slotHeight(slotHeight), slotCountX(slotCountX), slotCountY(slotCountY)
{
    // whole image
    freeAreas.insert({ 0, 0, slotCountX, slotCountY});
    createImageResources();
    vk.registerCleanupCallback(std::bind(&TextureAtlas::cleanup, this));
    VulkanInstance::registerObjectUsingVulkan(ObjectUsingVulkan(std::string("TextureAtlas-") + gz::toString(slotWidth) + "x" + gz::toString(slotHeight),
                { &textureImage, &textureImageMemory, &textureImageView, &textureSampler }, 
                {}));
}


void TextureAtlas::cleanup() {
    VulkanInstance::vLog("TextureAtlas::cleanup, textureSampler", reinterpret_cast<uint64_t>(textureSampler), "textureImageView", reinterpret_cast<uint64_t>(textureImageView));
    vkDestroySampler(vk.device, textureSampler, nullptr);
    vkDestroyImageView(vk.device, textureImageView, nullptr);
    vkDestroyImage(vk.device, textureImage, nullptr);
    vkFreeMemory(vk.device, textureImageMemory, nullptr);
}


void TextureAtlas::createImageResources() {
    vk.createImage(slotWidth * slotCountX, slotHeight * slotCountY, VK_FORMAT_R8G8B8A8_SRGB, VK_IMAGE_TILING_LINEAR, 
            VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT, 
            VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, 
            textureImage, textureImageMemory);
    VkCommandBuffer cmdBuffer = vk.beginSingleTimeCommands(vk.commandPoolGraphics);
    vk.transitionImageLayout(textureImage, VK_FORMAT_R8G8B8A8_SRGB, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, &cmdBuffer);
    VkImageSubresourceRange subresourceRange{};
    subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    subresourceRange.baseMipLevel = 0;
    subresourceRange.levelCount = 1;
    subresourceRange.baseArrayLayer = 0;
    subresourceRange.layerCount = 1;
    vkCmdClearColorImage(cmdBuffer, textureImage, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, &missingTextureColor, 1, &subresourceRange);
    vk.transitionImageLayout(textureImage, VK_FORMAT_R8G8B8A8_SRGB, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, &cmdBuffer);
    vk.endSingleTimeCommands(cmdBuffer, vk.commandPoolGraphics, vk.graphicsQ);

    vk.createImageView(VK_FORMAT_R8G8B8A8_SRGB, textureImage, textureImageView, VK_IMAGE_ASPECT_COLOR_BIT);
    vk.createTextureSampler(textureSampler);
}


//
// ADD TEXTURE
//
std::pair<glm::vec2, glm::vec2> TextureAtlas::addTexture(uint8_t* pixels, uint16_t width, uint16_t height) {
    if (freeAreas.empty()) {
        throw Exception("No texture slots available", "TextureAtlas::addTexture");
    }
    // needed slot count in x/y direction
    uint16_t slotsX, slotsY;
    // width / slotHeight (+ 1 if another slot is not fully filled)
    slotsX = width / slotWidth + ((width % slotWidth != 0) ? 1 : 0);
    slotsY = height / slotHeight + ((height % slotHeight != 0) ? 1 : 0);

    // find smallest area that fits the texture
    auto freeArea = freeAreas.end();
    for (auto it = freeAreas.begin(); it != freeAreas.end(); it++) {
        // if fits the texture
        if (it->width >= slotsX and it->height >= slotsY) {
            freeArea = it;
            break;
        }
    }
    // check if match
    if (freeArea == freeAreas.end()) {
        throw Exception("Not enough texture slots available to fit the texture", "TextureAtlas::addTexture");
    }

    // consume free area
    auto node = freeAreas.extract(freeArea);
    TextureImageArea& selectedArea = node.value();

    // topleft slot for texture
    uint16_t x = selectedArea.x;
    uint16_t y = selectedArea.y;
    // update selectedArea 
    uint16_t selectedAreaNewY = selectedArea.y + slotsY;
    uint16_t selectedAreaNewWidth = slotsX;
    uint16_t selectedAreaNewHeight = selectedArea.height - slotsY;
    // if split needed
    if (slotsX < selectedArea.width and slotsY < selectedArea.height) {
        // split so that the size difference between the new free area and the remaining free area is maximal
        TextureImageArea newFreeArea;
        newFreeArea.x = selectedArea.x + slotsX;  // right of new texture
        newFreeArea.y = selectedArea.y;
        newFreeArea.width = selectedArea.width - slotsX;
        // try full height first first
        newFreeArea.height = selectedArea.height;

        // size difference with (newArea.height == full height) < size difference with (newArea.height == textureHeight)
        int sizeDiffNewAreaFullHeight = static_cast<int>(newFreeArea.width * newFreeArea.height) - selectedAreaNewWidth * selectedAreaNewHeight;
        int sizeDiffNewAreaTextureHeight = static_cast<int>(newFreeArea.width * (selectedArea.height - slotsY)) - selectedArea.width * selectedAreaNewHeight;
        if (sizeDiffNewAreaFullHeight * sizeDiffNewAreaFullHeight < sizeDiffNewAreaTextureHeight * sizeDiffNewAreaTextureHeight) {
            newFreeArea.height = slotsY;
            selectedAreaNewWidth = selectedArea.width;
        }
        // insert new area
        freeAreas.insert(std::move(newFreeArea));
    }

    if (selectedAreaNewHeight > 0 and selectedAreaNewWidth > 0) {
        // insert selected area with updated width, height and y
        selectedArea.width = selectedAreaNewWidth;
        selectedArea.height = selectedAreaNewHeight;
        selectedArea.y = selectedAreaNewY;
        freeAreas.insert(std::move(node));
    }
    mergeFreeAreas();
    VulkanInstance::vLog("TextureAtlas::addTexture: Adding texture at position x,y=(", x, y, "), with slotCountX,Y=(", slotsX, slotsY, "), with dimensions w,h=(", width, height, ")");
    blitTextureOnImage(pixels, x, y, width, height);
    // topleft normalized: slot / slotCount
    // bottomright normalized: (slot + (textureSize/slotCountize)) / slotCount
    return { glm::vec2(static_cast<float>(x) / slotCountX, static_cast<float>(y) / slotCountY), 
             glm::vec2((static_cast<float>(x) + static_cast<float>(width) / slotWidth) / slotCountX,
                (static_cast<float>(y) + static_cast<float>(height) / slotHeight) / slotCountY) };
}


void TextureAtlas::blitTextureOnImage(uint8_t* pixels, uint16_t slotX, uint16_t slotY, uint16_t textureWidth, uint16_t textureHeight) {
    constexpr size_t BYTES_PER_PIXEL = 4;
    VkDeviceSize imageSize = textureWidth * textureHeight * BYTES_PER_PIXEL;
    VkBuffer stagingBuffer;
    VkDeviceMemory stagingBufferMemory;
    vk.createBuffer(imageSize, VK_BUFFER_USAGE_TRANSFER_SRC_BIT, 
            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, 
            stagingBuffer, stagingBufferMemory);
    void* data;
    vkMapMemory(vk.device, stagingBufferMemory, NO_OFFSET, imageSize, NO_FLAGS, &data);
    memcpy(data, pixels, static_cast<size_t>(imageSize));
    vkUnmapMemory(vk.device, stagingBufferMemory);

    vk.transitionImageLayout(textureImage, VK_FORMAT_R8G8B8A8_SRGB, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
    vk.copyBufferToImage(stagingBuffer, textureImage, 
            static_cast<int32_t>(slotWidth * slotX), static_cast<int32_t>(slotHeight * slotY),  // offset
            static_cast<uint32_t>(textureWidth), static_cast<uint32_t>(textureHeight));  // dimensions
    vk.transitionImageLayout(textureImage, VK_FORMAT_R8G8B8A8_SRGB, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
    
    vkDestroyBuffer(vk.device, stagingBuffer, nullptr);
    vkFreeMemory(vk.device, stagingBufferMemory, nullptr); 
}


void TextureAtlas::mergeFreeAreas() {
begin:
    VulkanInstance::vLog("TextureAtlas::mergeFreeAreas before merge:", freeAreas);
    // surely not the most efficient way, but it should only be run at load time and thus not be too problematic
    // iterate over the areas for each area
    for (auto it = freeAreas.begin(); it != freeAreas.end(); it++) {
        for (auto it2 = freeAreas.begin(); it2 != freeAreas.end(); it2++) {
            if (it == it2) { continue; }
            // if same y and same height and next to each other
            if (it->y == it2->y and it->height == it2->height and it2->x == it->x + it->width) {
                auto area1 = freeAreas.extract(it);  // node
                area1.value().width += it2->width;
                freeAreas.insert(std::move(area1));
                freeAreas.erase(it2);
                goto begin;  // start over since the iterators are invalid and new merges might be possible
            }
            // if same x and same with and below/above each other
            if (it->x == it2->x and it->width == it2->width and it2->y == it->y + it->height) {
                auto area1 = freeAreas.extract(it);  // node
                area1.value().height += it2->height;
                freeAreas.insert(std::move(area1));
                freeAreas.erase(it2);
                goto begin;  // start over since the iterators are invalid and new merges might be possible
            }
        }
    }
    VulkanInstance::vLog("TextureAtlas::mergeFreeAreas after merge:", freeAreas);
    
}

std::string TextureAtlas::toString() const {
    return "TextureAtlas(SlotSize: " + gz::toString(slotWidth) + "x" + gz::toString(slotHeight)
        + ", SlotCount: " + gz::toString(slotCountX) + "x" + gz::toString(slotCountY) + ", FreeAreas: " + gz::toString(freeAreas) + ")";
}


}  // namespace gz