refactor shader module

WinteryFox · Oct 26, 2023 · 336113f · 336113f
1 parent 237bf39
commit 336113f
Show file tree

Hide file tree

Showing 7 changed files with 31 additions and 32 deletions.
diff --git a/engine/ShaderProgram.h b/engine/ShaderProgram.h
@@ -7,14 +7,20 @@ namespace Vixen {
     template<class T, class = std::enable_if<std::is_base_of<ShaderModule, T>::value>>
     class ShaderProgram {
     protected:
-        std::map<ShaderModule::Stage, std::shared_ptr<T>> modules;
+        std::shared_ptr<T> vertex;
+
+        std::shared_ptr<T> fragment;
 
     public:
-        explicit ShaderProgram(const std::map<ShaderModule::Stage, std::shared_ptr<T>> &modules)
-                : modules(modules) {}
+        ShaderProgram(const std::shared_ptr<T> &vertex, const std::shared_ptr<T> &fragment)
+                : vertex(vertex), fragment(fragment) {}
+
+        const std::shared_ptr<T> &getVertex() const {
+            return vertex;
+        }
 
-        const std::map<ShaderModule::Stage, std::shared_ptr<T>> &getModules() const {
-            return modules;
+        const std::shared_ptr<T> &getFragment() const {
+            return fragment;
         }
     };
 }
diff --git a/engine/vk/VkPipeline.cpp b/engine/vk/VkPipeline.cpp
@@ -11,21 +11,9 @@ namespace Vixen::Vk {
         config(config),
         pipelineLayout(device, program),
         renderPass(device, program, swapchain) {
-        auto modules = program.getModules();
-
         std::vector<VkPipelineShaderStageCreateInfo> stages;
-        stages.reserve(modules.size());
-        for (const auto &[stage, module]: modules) {
-            stages.emplace_back(VkPipelineShaderStageCreateInfo{
-                    .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
-                    .pNext = VK_NULL_HANDLE,
-                    .flags = 0,
-                    .stage = getVulkanShaderStage(stage),
-                    .module = module->getModule(),
-                    .pName = module->getEntrypoint().c_str(),
-                    .pSpecializationInfo = VK_NULL_HANDLE
-            });
-        }
+        stages.emplace_back(program.getVertex()->createInfo());
+        stages.emplace_back(program.getFragment()->createInfo());
 
         std::vector<VkDynamicState> dynamicStates{
                 VK_DYNAMIC_STATE_VIEWPORT,
@@ -38,7 +26,7 @@ namespace Vixen::Vk {
                 .pDynamicStates = dynamicStates.data()
         };
 
-        const auto &vertexModule = modules[ShaderModule::Stage::VERTEX];
+        const auto &vertexModule = program.getVertex();
 
         std::vector<VkVertexInputBindingDescription> vertexBindings{};
         for (const auto &binding: vertexModule->getBindings()) {

diff --git a/engine/vk/VkShaderModule.cpp b/engine/vk/VkShaderModule.cpp
@@ -27,7 +27,15 @@ namespace Vixen::Vk {
         vkDestroyShaderModule(device->getDevice(), module, nullptr);
     }
 
-    ::VkShaderModule VkShaderModule::getModule() const {
-        return module;
+    VkPipelineShaderStageCreateInfo VkShaderModule::createInfo() {
+        return {
+                .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
+                .pNext = VK_NULL_HANDLE,
+                .flags = 0,
+                .stage = getVulkanShaderStage(stage),
+                .module = module,
+                .pName = entrypoint.c_str(),
+                .pSpecializationInfo = VK_NULL_HANDLE
+        };
     }
 }
diff --git a/engine/vk/VkShaderModule.h b/engine/vk/VkShaderModule.h
@@ -41,7 +41,7 @@ namespace Vixen::Vk {
 
         ~VkShaderModule();
 
-        [[nodiscard]] ::VkShaderModule getModule() const;
+        VkPipelineShaderStageCreateInfo createInfo();
 
         class Builder {
             // TODO: This builder has slightly confusing API, you can compile before setting the stage meaning its possible to mistakenly have the wrong stage set at compile time

diff --git a/engine/vk/VkShaderProgram.cpp b/engine/vk/VkShaderProgram.cpp
@@ -1,6 +1,8 @@
 #include "VkShaderProgram.h"
 
 namespace Vixen::Vk {
-    VkShaderProgram::VkShaderProgram(const std::map<ShaderModule::Stage, std::shared_ptr<VkShaderModule>> &modules)
-            : ShaderProgram(modules) {}
+    VkShaderProgram::VkShaderProgram(
+            const std::shared_ptr<VkShaderModule> &vertex,
+            const std::shared_ptr<VkShaderModule> &fragment
+    ) : ShaderProgram(vertex, fragment) {}
 }
diff --git a/engine/vk/VkShaderProgram.h b/engine/vk/VkShaderProgram.h
@@ -14,6 +14,6 @@
 namespace Vixen::Vk {
     class VkShaderProgram : public ShaderProgram<VkShaderModule> {
     public:
-        explicit VkShaderProgram(const std::map<ShaderModule::Stage, std::shared_ptr<VkShaderModule>> &modules);
+        VkShaderProgram(const std::shared_ptr<VkShaderModule> &vertex, const std::shared_ptr<VkShaderModule> &fragment);
     };
 }
diff --git a/engine/vk/test/main.cpp b/engine/vk/test/main.cpp
@@ -44,12 +44,7 @@ int main() {
     auto fragment = Vixen::Vk::VkShaderModule::Builder()
             .setStage(Vixen::ShaderModule::Stage::FRAGMENT)
             .compileFromFile(vixen.device, "../../editor/shaders/triangle.frag");
-    auto program = Vixen::Vk::VkShaderProgram(
-            {
-                    {Vixen::ShaderModule::Stage::VERTEX,   vertex},
-                    {Vixen::ShaderModule::Stage::FRAGMENT, fragment}
-            }
-    );
+    auto program = Vixen::Vk::VkShaderProgram(vertex, fragment);
 
     int width;
     int height;