DescriptorSets

layout(set = 0, binding = 0) uniform UboVS {  
    mat4 mvp;  
} uboVS;  

layout(location = 0) in vec3 inPosition;  
layout(location = 0) out vec3 fragColor;  

void main() {  
    gl_Position = uboVS.mvp * vec4(inPosition, 1.0);  
    fragColor = inPosition;  
}

layout(set = 1, binding = 0) uniform UboMaterial {  
    vec4 color;  
    float roughness;  
} material;  

layout(location = 0) in vec3 fragColor;  
layout(location = 0) out vec4 outColor;  

void main() {  
    outColor = material.color * vec4(fragColor, 1.0);  
}

#include <vulkan/vulkan.h>
#include <vector>
#include <stdexcept>
#include <array>

// 假设这些函数和变量已在其他地方定义
VkDevice device;
VkPhysicalDevice physicalDevice;
uint32_t queueFamilyIndex;

// 用于存储 uniform 缓冲的结构体
struct UniformBufferObject {
    glm::mat4 mvp;
};

struct MaterialProperties {
    glm::vec4 color;
    float roughness;
};

// 创建缓冲的辅助函数
VkBuffer createBuffer(VkDeviceSize size, VkBufferUsageFlags usage, VkMemoryPropertyFlags properties, VkDeviceMemory& bufferMemory) {
    VkBuffer buffer;
    // 创建缓冲的代码...
    return buffer;
}

class VulkanMultipleDescriptorSets {
public:
    void setup() {
        createDescriptorSetLayouts();
        createUniformBuffers();
        createDescriptorPool();
        createDescriptorSets();
        createPipelineLayout();
    }

    void cleanup() {
        vkDestroyBuffer(device, uniformBufferMVP, nullptr);
        vkFreeMemory(device, uniformBufferMemoryMVP, nullptr);
        vkDestroyBuffer(device, uniformBufferMaterial, nullptr);
        vkFreeMemory(device, uniformBufferMemoryMaterial, nullptr);
        vkDestroyDescriptorSetLayout(device, descriptorSetLayoutMVP, nullptr);
        vkDestroyDescriptorSetLayout(device, descriptorSetLayoutMaterial, nullptr);
        vkDestroyDescriptorPool(device, descriptorPool, nullptr);
        vkDestroyPipelineLayout(device, pipelineLayout, nullptr);
    }

    void updateUniformBuffer(uint32_t currentImage) {
        // 更新 MVP 矩阵
        UniformBufferObject ubo{};
        // ... 计算 MVP 矩阵 ...
        void* data;
        vkMapMemory(device, uniformBufferMemoryMVP, 0, sizeof(ubo), 0, &data);
        memcpy(data, &ubo, sizeof(ubo));
        vkUnmapMemory(device, uniformBufferMemoryMVP);

        // 更新材质属性（如果需要）
        MaterialProperties material{};
        material.color = glm::vec4(1.0f, 0.0f, 0.0f, 1.0f);
        material.roughness = 0.5f;
        vkMapMemory(device, uniformBufferMemoryMaterial, 0, sizeof(material), 0, &data);
        memcpy(data, &material, sizeof(material));
        vkUnmapMemory(device, uniformBufferMemoryMaterial);
    }

    void bindDescriptorSets(VkCommandBuffer commandBuffer) {
        vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pipelineLayout, 
                                0, 2, descriptorSets.data(), 0, nullptr);
    }

private:
    VkDescriptorSetLayout descriptorSetLayoutMVP;
    VkDescriptorSetLayout descriptorSetLayoutMaterial;
    VkPipelineLayout pipelineLayout;
    VkBuffer uniformBufferMVP;
    VkDeviceMemory uniformBufferMemoryMVP;
    VkBuffer uniformBufferMaterial;
    VkDeviceMemory uniformBufferMemoryMaterial;
    VkDescriptorPool descriptorPool;
    std::array<VkDescriptorSet, 2> descriptorSets;

    void createDescriptorSetLayouts() {
        // MVP 矩阵的描述符集布局
        VkDescriptorSetLayoutBinding mvpLayoutBinding{};
        mvpLayoutBinding.binding = 0;
        mvpLayoutBinding.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
        mvpLayoutBinding.descriptorCount = 1;
        mvpLayoutBinding.stageFlags = VK_SHADER_STAGE_VERTEX_BIT;

        VkDescriptorSetLayoutCreateInfo layoutInfoMVP{};
        layoutInfoMVP.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
        layoutInfoMVP.bindingCount = 1;
        layoutInfoMVP.pBindings = &mvpLayoutBinding;

        if (vkCreateDescriptorSetLayout(device, &layoutInfoMVP, nullptr, &descriptorSetLayoutMVP) != VK_SUCCESS) {
            throw std::runtime_error("Failed to create descriptor set layout for MVP!");
        }

        // 材质属性的描述符集布局
        VkDescriptorSetLayoutBinding materialLayoutBinding{};
        materialLayoutBinding.binding = 0;
        materialLayoutBinding.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
        materialLayoutBinding.descriptorCount = 1;
        materialLayoutBinding.stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT;

        VkDescriptorSetLayoutCreateInfo layoutInfoMaterial{};
        layoutInfoMaterial.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
        layoutInfoMaterial.bindingCount = 1;
        layoutInfoMaterial.pBindings = &materialLayoutBinding;

        if (vkCreateDescriptorSetLayout(device, &layoutInfoMaterial, nullptr, &descriptorSetLayoutMaterial) != VK_SUCCESS) {
            throw std::runtime_error("Failed to create descriptor set layout for Material!");
        }
    }

    void createUniformBuffers() {
        VkDeviceSize bufferSizeMVP = sizeof(UniformBufferObject);
        uniformBufferMVP = createBuffer(bufferSizeMVP, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT, 
                                        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, 
                                        uniformBufferMemoryMVP);

        VkDeviceSize bufferSizeMaterial = sizeof(MaterialProperties);
        uniformBufferMaterial = createBuffer(bufferSizeMaterial, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT, 
                                             VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, 
                                             uniformBufferMemoryMaterial);
    }

    void createDescriptorPool() {
        std::array<VkDescriptorPoolSize, 2> poolSizes{};
        poolSizes[0].type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
        poolSizes[0].descriptorCount = 1;
        poolSizes[1].type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
        poolSizes[1].descriptorCount = 1;

        VkDescriptorPoolCreateInfo poolInfo{};
        poolInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
        poolInfo.poolSizeCount = static_cast<uint32_t>(poolSizes.size());
        poolInfo.pPoolSizes = poolSizes.data();
        poolInfo.maxSets = 2;

        if (vkCreateDescriptorPool(device, &poolInfo, nullptr, &descriptorPool) != VK_SUCCESS) {
            throw std::runtime_error("Failed to create descriptor pool!");
        }
    }

    void createDescriptorSets() {
        std::array<VkDescriptorSetLayout, 2> layouts = {descriptorSetLayoutMVP, descriptorSetLayoutMaterial};
        VkDescriptorSetAllocateInfo allocInfo{};
        allocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
        allocInfo.descriptorPool = descriptorPool;
        allocInfo.descriptorSetCount = static_cast<uint32_t>(layouts.size());
        allocInfo.pSetLayouts = layouts.data();

        if (vkAllocateDescriptorSets(device, &allocInfo, descriptorSets.data()) != VK_SUCCESS) {
            throw std::runtime_error("Failed to allocate descriptor sets!");
        }

        // 更新 MVP 描述符集
        VkDescriptorBufferInfo bufferInfoMVP{};
        bufferInfoMVP.buffer = uniformBufferMVP;
        bufferInfoMVP.offset = 0;
        bufferInfoMVP.range = sizeof(UniformBufferObject);

        VkWriteDescriptorSet descriptorWriteMVP{};
        descriptorWriteMVP.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
        descriptorWriteMVP.dstSet = descriptorSets[0];
        descriptorWriteMVP.dstBinding = 0;
        descriptorWriteMVP.dstArrayElement = 0;
        descriptorWriteMVP.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
        descriptorWriteMVP.descriptorCount = 1;
        descriptorWriteMVP.pBufferInfo = &bufferInfoMVP;

        // 更新材质描述符集
        VkDescriptorBufferInfo bufferInfoMaterial{};
        bufferInfoMaterial.buffer = uniformBufferMaterial;
        bufferInfoMaterial.offset = 0;
        bufferInfoMaterial.range = sizeof(MaterialProperties);

        VkWriteDescriptorSet descriptorWriteMaterial{};
        descriptorWriteMaterial.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
        descriptorWriteMaterial.dstSet = descriptorSets[1];
        descriptorWriteMaterial.dstBinding = 0;
        descriptorWriteMaterial.dstArrayElement = 0;
        descriptorWriteMaterial.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
        descriptorWriteMaterial.descriptorCount = 1;
        descriptorWriteMaterial.pBufferInfo = &bufferInfoMaterial;

        std::array<VkWriteDescriptorSet, 2> descriptorWrites = {descriptorWriteMVP, descriptorWriteMaterial};
        vkUpdateDescriptorSets(device, static_cast<uint32_t>(descriptorWrites.size()), descriptorWrites.data(), 0, nullptr);
    }

    void createPipelineLayout() {
        std::array<VkDescriptorSetLayout, 2> setLayouts = {descriptorSetLayoutMVP, descriptorSetLayoutMaterial};

        VkPipelineLayoutCreateInfo pipelineLayoutInfo{};
        pipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
        pipelineLayoutInfo.setLayoutCount = static_cast<uint32_t>(setLayouts.size());
        pipelineLayoutInfo.pSetLayouts = setLayouts.data();

        if (vkCreatePipelineLayout(device, &pipelineLayoutInfo, nullptr, &pipelineLayout) != VK_SUCCESS) {
            throw std::runtime_error("Failed to create pipeline layout!");
        }
    }
};