Crafter.Graphics/Crafter.Graphics-VulkanPipeline.cppm

module;

#include <cstdint>
#include <vulkan/vulkan.h>
#include <array>
#include "VulkanInitializers.hpp"
#define GLM_FORCE_RADIANS
#define GLM_FORCE_DEPTH_ZERO_TO_ONE
#define GLM_ENABLE_EXPERIMENTAL
#include <glm/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include <glm/gtc/matrix_inverse.hpp>
#include <glm/gtc/type_ptr.hpp>
#include "VulkanBuffer.h"
#include "camera.hpp"
#include <unordered_map>

export module Crafter.Graphics:VulkanPipeline;
import :VulkanDevice;
import :VulkanShader;
import :WindowWaylandVulkan;

namespace Crafter {
	struct DescriptorEntry {
		VkDescriptorType type;
		std::uint32_t occurrences = 0;
	};

    export template <typename MeshShader, typename FragmentShader>
    class VulkanPipeline {
		private:
		constexpr static std::uint32_t totalDescriptorCount = MeshShader::descriptorCount+FragmentShader::descriptorCount;
		consteval static std::uint32_t GetUniqueDiscriptorCount() {
			DescriptorEntry types[] = {{VK_DESCRIPTOR_TYPE_SAMPLER, 0},{VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 0},{VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE, 0},{VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 0},{VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER, 0},{VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER, 0},{VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 0},{VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 0},{VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC, 0},{VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC, 0},{VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 0},{VK_DESCRIPTOR_TYPE_INLINE_UNIFORM_BLOCK, 0},{VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR, 0},{VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_NV, 0},{VK_DESCRIPTOR_TYPE_SAMPLE_WEIGHT_IMAGE_QCOM, 0},{VK_DESCRIPTOR_TYPE_BLOCK_MATCH_IMAGE_QCOM, 0},{VK_DESCRIPTOR_TYPE_MUTABLE_EXT, 0},{VK_DESCRIPTOR_TYPE_PARTITIONED_ACCELERATION_STRUCTURE_NV, 0},{VK_DESCRIPTOR_TYPE_INLINE_UNIFORM_BLOCK_EXT, 0},{VK_DESCRIPTOR_TYPE_MUTABLE_VALVE, 0}};

			for(const DescriptorBinding& binding : MeshShader::descriptors) {
				for(DescriptorEntry& type : types) {
					if(type.type == binding.type) {
						type.occurrences++;
					}
				}
			}

			for(const DescriptorBinding& binding : FragmentShader::descriptors) {
				for(DescriptorEntry& type : types) {
					if(type.type == binding.type) {
						type.occurrences++;
					}
				}
			}

			std::uint32_t size = 0;
			for(DescriptorEntry& type : types) {
				size+=type.occurrences;
			}

			return size;
		}
		constexpr static std::uint32_t uniqueDescriptorCount = GetUniqueDiscriptorCount();
		consteval static std::array<VkDescriptorPoolSize, uniqueDescriptorCount> GetPoolSizes() {
			std::array<VkDescriptorPoolSize, uniqueDescriptorCount> types = {};
			std::uint32_t i = 0;

			for(const DescriptorBinding& binding : MeshShader::descriptors) {
				for(VkDescriptorPoolSize& type : types) {
					if(type.type == binding.type) {
						type.descriptorCount++;
						goto next;
					}
				}
				types[i].type = binding.type;
				types[i].descriptorCount = 1;
				next:;
			}

			for(const DescriptorBinding& binding : FragmentShader::descriptors) {
				for(VkDescriptorPoolSize& type : types) {
					if(type.type == binding.type) {
						type.descriptorCount++;
						goto next2;
					}
				}
				types[i].type = binding.type;
				types[i].descriptorCount = 1;
				next2:;
			}

			return types;
		}
		consteval static std::array<VkDescriptorSetLayoutBinding, totalDescriptorCount> GetDescriptorSet() {
			std::array<VkDescriptorSetLayoutBinding, totalDescriptorCount> set;

			std::uint32_t i = 0;
			for(const DescriptorBinding& binding : MeshShader::descriptors) {
				set[i] = {binding.slot, binding.type, 1, VK_SHADER_STAGE_MESH_BIT_EXT, nullptr};
				i++;
			}

			for(const DescriptorBinding& binding : FragmentShader::descriptors) {
				set[i] = {binding.slot, binding.type, 1, VK_SHADER_STAGE_FRAGMENT_BIT, nullptr};
				i++;
			}

			return set;
		}
        public:
        inline static VkPipeline pipeline;
        inline static VkPipelineLayout pipelineLayout;
		inline static VkDescriptorSet descriptorSet = VK_NULL_HANDLE;
		inline static VkDescriptorSetLayout descriptorSetLayout = VK_NULL_HANDLE;
		inline static VkDescriptorPool descriptorPool = VK_NULL_HANDLE;

		inline static struct UniformData {
			glm::mat4 projection;
			glm::mat4 model;
			glm::mat4 view;
		} uniformData;
		inline static vks::Buffer uniformBuffer;

        static void CreatePipeline() {
			Camera camera;
			camera.type = Camera::CameraType::lookat;
			camera.setPerspective(60.0f, 128 / 128, 0.1f, 512.0f);
			camera.setRotation(glm::vec3(0.0f, 15.0f, 0.0f));
			camera.setTranslation(glm::vec3(0.0f, 0.0f, -5.0f));
			
			VulkanDevice::CreateBuffer(VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, &uniformBuffer, sizeof(UniformData));
			VulkanDevice::CHECK_VK_RESULT(uniformBuffer.map());
			uniformData.projection = camera.matrices.perspective;
			uniformData.view = camera.matrices.view;
			uniformData.model = glm::mat4(1.0f);
			memcpy(uniformBuffer.mapped, &uniformData, sizeof(UniformData));

		// Pool
		std::array<VkDescriptorPoolSize, uniqueDescriptorCount> poolSizes = GetPoolSizes();
		VkDescriptorPoolCreateInfo descriptorPoolInfo = vks::initializers::descriptorPoolCreateInfo(uniqueDescriptorCount, poolSizes.data(), 1);
		VulkanDevice::CHECK_VK_RESULT(vkCreateDescriptorPool(VulkanDevice::device, &descriptorPoolInfo, nullptr, &descriptorPool));

		// Layout
		constexpr std::array<VkDescriptorSetLayoutBinding, totalDescriptorCount> setLayoutBindings = GetDescriptorSet();
		VkDescriptorSetLayoutCreateInfo descriptorLayoutInfo = vks::initializers::descriptorSetLayoutCreateInfo(setLayoutBindings.data(), totalDescriptorCount);
		VulkanDevice::CHECK_VK_RESULT(vkCreateDescriptorSetLayout(VulkanDevice::device, &descriptorLayoutInfo, nullptr, &descriptorSetLayout));

		// Set
		VkDescriptorSetAllocateInfo allocInfo = vks::initializers::descriptorSetAllocateInfo(descriptorPool, &descriptorSetLayout, 1);
		VulkanDevice::CHECK_VK_RESULT(vkAllocateDescriptorSets(VulkanDevice::device, &allocInfo, &descriptorSet));

		std::vector<VkWriteDescriptorSet> modelWriteDescriptorSets = {
			vks::initializers::writeDescriptorSet(descriptorSet, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 0, &uniformBuffer.descriptor),
		};
		vkUpdateDescriptorSets(VulkanDevice::device, static_cast<uint32_t>(modelWriteDescriptorSets.size()), modelWriteDescriptorSets.data(), 0, nullptr);

        // Layout
		VkPipelineLayoutCreateInfo pipelineLayoutInfo = vks::initializers::pipelineLayoutCreateInfo(&descriptorSetLayout, 1);
		VulkanDevice::CHECK_VK_RESULT(vkCreatePipelineLayout(VulkanDevice::device, &pipelineLayoutInfo, nullptr, &pipelineLayout));

		// Pipeline
		VkPipelineInputAssemblyStateCreateInfo inputAssemblyState = vks::initializers::pipelineInputAssemblyStateCreateInfo(VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST, 0, VK_FALSE);
		VkPipelineRasterizationStateCreateInfo rasterizationState = vks::initializers::pipelineRasterizationStateCreateInfo(VK_POLYGON_MODE_FILL, VK_CULL_MODE_BACK_BIT, VK_FRONT_FACE_COUNTER_CLOCKWISE, 0);
		VkPipelineColorBlendAttachmentState blendAttachmentState = vks::initializers::pipelineColorBlendAttachmentState(0xf, VK_FALSE);
		VkPipelineColorBlendStateCreateInfo colorBlendState = vks::initializers::pipelineColorBlendStateCreateInfo(1, &blendAttachmentState);
		VkPipelineDepthStencilStateCreateInfo depthStencilState = vks::initializers::pipelineDepthStencilStateCreateInfo(VK_TRUE, VK_TRUE, VK_COMPARE_OP_LESS_OR_EQUAL);
		VkPipelineViewportStateCreateInfo viewportState = vks::initializers::pipelineViewportStateCreateInfo(1, 1, 0);
		VkPipelineMultisampleStateCreateInfo multisampleState = vks::initializers::pipelineMultisampleStateCreateInfo(VK_SAMPLE_COUNT_1_BIT, 0);
		std::vector<VkDynamicState> dynamicStateEnables = { VK_DYNAMIC_STATE_VIEWPORT, VK_DYNAMIC_STATE_SCISSOR };
		VkPipelineDynamicStateCreateInfo dynamicState = vks::initializers::pipelineDynamicStateCreateInfo(dynamicStateEnables);
		std::array<VkPipelineShaderStageCreateInfo, 2> shaderStages;

		VkFormat format = VK_FORMAT_R8G8B8A8_UNORM;

		VkPipelineRenderingCreateInfoKHR pipeline_create{VK_STRUCTURE_TYPE_PIPELINE_RENDERING_CREATE_INFO_KHR};
		pipeline_create.pNext                   = VK_NULL_HANDLE;
		pipeline_create.colorAttachmentCount    = 1;
		pipeline_create.pColorAttachmentFormats = &format;
		pipeline_create.depthAttachmentFormat   = VulkanDevice::depthFormat;

		VkGraphicsPipelineCreateInfo pipelineCI = vks::initializers::pipelineCreateInfo(pipelineLayout, VK_NULL_HANDLE, 0);

		pipelineCI.pNext = &pipeline_create;
		pipelineCI.pRasterizationState = &rasterizationState;
		pipelineCI.pColorBlendState = &colorBlendState;
		pipelineCI.pMultisampleState = &multisampleState;
		pipelineCI.pViewportState = &viewportState;
		pipelineCI.pDepthStencilState = &depthStencilState;
		pipelineCI.pDynamicState = &dynamicState;
		pipelineCI.stageCount = static_cast<uint32_t>(shaderStages.size());
		pipelineCI.pStages = shaderStages.data();

		// Not using a vertex shader, mesh shading doesn't require vertex input state
		pipelineCI.pInputAssemblyState = nullptr;
		pipelineCI.pVertexInputState = nullptr;

        //Mesh stage of the pipeline
        shaderStages[0].sType  = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
        shaderStages[0].stage  = MeshShader::_stage;
        shaderStages[0].module = MeshShader::shader;
        shaderStages[0].pName  = MeshShader::_entrypoint.value;
		shaderStages[0].flags = 0;
   
        // Fragment stage of the pipeline
        shaderStages[1].sType  = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
        shaderStages[1].stage  = FragmentShader::_stage;
        shaderStages[1].module = FragmentShader::shader;
        shaderStages[1].pName  = FragmentShader::_entrypoint.value;
		shaderStages[1].flags = 0;
               
		VulkanDevice::CHECK_VK_RESULT(vkCreateGraphicsPipelines(VulkanDevice::device, VK_NULL_HANDLE, 1, &pipelineCI, nullptr, &pipeline));
        }
    };
}