image.cpp

#include <cassert>
#include <cstdio>
#include <vk_mem_alloc.h>
#include <volk.h>

#include "backend/vulkan/debug.h"
#include "backend/vulkan/utils.h"
#include "backend/vulkan/vk_api.h"
#include "backend/vulkan/vk_conversion.h"

namespace rhi {
using namespace vk;

bool create_image(Device &device, ImageDesc &desc, Image &out) {
    VkImageCreateInfo ci{};
    ci.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
    if (desc.is_cubemap) {
        ci.flags = VK_IMAGE_CREATE_CUBE_COMPATIBLE_BIT;
    }
    ci.imageType = desc.depth > 1 ? VK_IMAGE_TYPE_3D : VK_IMAGE_TYPE_2D;
    ci.extent.width = desc.width;
    ci.extent.height = desc.height;
    ci.extent.depth = desc.depth;
    ci.mipLevels = desc.mip_levels;
    ci.arrayLayers = desc.layers;
    ci.format = image_format_to_vk(desc.format);
    ci.tiling = VK_IMAGE_TILING_OPTIMAL;
    ci.usage = image_usage_to_vk(device, desc);
    ci.samples = sample_count_to_vk(desc.sample_count);
    ci.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
    VmaAllocationCreateInfo alloc_ci = {};
    alloc_ci.usage = VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE;
    ci.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
    VkImage vk_image;
    VmaAllocation vk_alloc;
    VK_CHECK(vmaCreateImage(device.memory.allocator, &ci, &alloc_ci, &vk_image, &vk_alloc, nullptr));
    out._handle = vk_image;
    out._alloc = vk_alloc;
    out.desc = desc;
    if (!desc.name.empty()) {
        VK_NAME_EX(device, vk_image, desc.name.c_str());
    }
    return true;
}

static rhi::ImageAspect aspect_from_format(rhi::ImageFormat fmt) {
    return (fmt == rhi::ImageFormat::D32_FLOAT) ? rhi::ImageAspect::Depth : rhi::ImageAspect::Color;
}

static bool
create_vk_image_view(Device &device, const ImageViewDesc &desc, VkImageView &out, rhi::ImageAspect &out_aspect) {
    if (desc.image == nullptr || desc.image->_handle == 0) {
        return false;
    }
    rhi::ImageAspect aspect = desc.aspect;
    if (aspect == rhi::ImageAspect::None) {
        aspect = aspect_from_format(desc.image->desc.format);
    }
    VkImageViewCreateInfo ci{};
    ci.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
    ci.image = desc.image->_handle;
    ci.viewType = image_view_type_to_vk(desc.dimension);
    ci.format = image_format_to_vk(desc.image->desc.format);
    ci.subresourceRange.aspectMask = image_aspect_to_vk(aspect);
    ci.subresourceRange.baseMipLevel = desc.mip_start;
    ci.subresourceRange.levelCount = desc.mip_count;
    ci.subresourceRange.baseArrayLayer = desc.layer_start;
    ci.subresourceRange.layerCount = desc.layer_count;
    ci.components.r = component_swizzle_to_vk(desc.swizzle_r);
    ci.components.g = component_swizzle_to_vk(desc.swizzle_g);
    ci.components.b = component_swizzle_to_vk(desc.swizzle_b);
    ci.components.a = component_swizzle_to_vk(desc.swizzle_a);
    VK_CHECK(vkCreateImageView(device.logical, &ci, nullptr, &out));
    out_aspect = aspect;
    return true;
}

void destroy_image(Device &device, Image &img) {
    vmaDestroyImage(device.memory.allocator, img._handle, img._alloc);
    img = {};
}

bool create_image_view(Device &device, const ImageViewDesc &desc, ImageView &out) {
    if (desc.image != nullptr) {
        if (desc.type == ImageViewType::Storage) {
            assert(
                Any(desc.image->desc.usage, ImageUsage::Storage) &&
                "Storage view requires ImageUsage::Storage at creation"
            );
        } else if (desc.type == ImageViewType::Sampled) {
            assert(
                Any(desc.image->desc.usage, ImageUsage::Sampled) &&
                "Sampled view requires ImageUsage::Sampled at creation"
            );
        }
    }
    VkImageView vk_view = VK_NULL_HANDLE;
    rhi::ImageAspect aspect = rhi::ImageAspect::None;
    if (!create_vk_image_view(device, desc, vk_view, aspect)) {
        return false;
    }
    const bool storage = (desc.type == ImageViewType::Storage);
    u64 slot = write_image_view(device, vk_view, storage);
    if (slot == UINT64_MAX) {
        vkDestroyImageView(device.logical, vk_view, nullptr);
        return false;
    }
    out._handle = vk_view;
    out.slot = slot;
    out.desc = desc;
    out.desc.aspect = aspect;
    if (desc.image != nullptr && !desc.image->desc.name.empty()) {
        const char *kind = (desc.type == ImageViewType::Storage) ? "storage" : "sampled";
        char view_name[256];
        (void)snprintf(view_name, sizeof(view_name), "%s.view:%s", desc.image->desc.name.c_str(), kind);
        VK_NAME_EX(device, vk_view, view_name);
    }
    return true;
}

u64 handle_id(Device &device, ImageView &view) {
    (void)device;
    return view.slot;
}

ImageView get_cached_image_view(Device &device, Image &image, const ImageViewDesc &desc) {
    ImageViewDesc key = desc;
    key.image = &image;
    if (key.aspect == rhi::ImageAspect::None) {
        key.aspect = aspect_from_format(image.desc.format);
    }
    if (key.dimension == rhi::TextureViewDimension::TEXTURE_2D && image.desc.is_cubemap) {
        key.dimension = rhi::TextureViewDimension::TEXTURE_CUBE;
    }
    for (auto &[k, v] : image.view_cache) {
        if (k.image == key.image && k.aspect == key.aspect && k.type == key.type && k.dimension == key.dimension &&
            k.mip_start == key.mip_start && k.mip_count == key.mip_count && k.layer_start == key.layer_start &&
            k.layer_count == key.layer_count) {
            ImageView out{};
            out.desc = k;
            out._handle = v;
            out.slot = UINT64_MAX;
            return out;
        }
    }
    VkImageView vk_view = VK_NULL_HANDLE;
    rhi::ImageAspect aspect = rhi::ImageAspect::None;
    create_vk_image_view(device, key, vk_view, aspect);
    key.aspect = aspect;
    image.view_cache.emplace_back(key, vk_view);
    ImageView out{};
    out.desc = key;
    out._handle = vk_view;
    out.slot = UINT64_MAX;
    return out;
}

void destroy_image_views(Device &device, Image &image) {
    for (auto &[k, v] : image.view_cache) {
        vkDestroyImageView(device.logical, v, nullptr);
    }
    image.view_cache.clear();
}

void destroy_image_view(Device &device, ImageView &view) {
    if (view._handle != nullptr) {
        vkDestroyImageView(device.logical, view._handle, nullptr);
        view._handle = nullptr;
    }
    if (view.slot != UINT64_MAX) {
        RESOURCE_SLOTS.free(view.slot);
        view.slot = UINT64_MAX;
    }
}

} // namespace rhi