From 4cca177c9da813338ad0f11384aa2b58a95d072e Mon Sep 17 00:00:00 2001
From: Lucas Alber <lucasd.alber@gmail.com>
Date: Tue, 12 Nov 2024 18:51:05 +0100
Subject: [PATCH] merian: Pipeline: Add syntactic sugar to push descriptor
 methods

---
 .../descriptor_set_layout_builder.hpp         | 10 +++-
 .../merian/vk/memory/resource_allocations.hpp |  4 +-
 include/merian/vk/pipeline/pipeline.hpp       | 60 +++++++++++++++++++
 .../merian/vk/pipeline/pipeline_layout.hpp    |  5 ++
 .../pipeline/specialization_info_builder.hpp  |  3 +-
 5 files changed, 78 insertions(+), 4 deletions(-)

diff --git a/include/merian/vk/descriptors/descriptor_set_layout_builder.hpp b/include/merian/vk/descriptors/descriptor_set_layout_builder.hpp
index 16eda4a..a4b01ab 100644
--- a/include/merian/vk/descriptors/descriptor_set_layout_builder.hpp
+++ b/include/merian/vk/descriptors/descriptor_set_layout_builder.hpp
@@ -153,7 +153,8 @@ class DescriptorSetLayoutBuilder {
         return *this;
     }
 
-    DescriptorSetLayoutBuilder& add_binding(const std::vector<vk::DescriptorSetLayoutBinding>& bindings) {
+    DescriptorSetLayoutBuilder&
+    add_binding(const std::vector<vk::DescriptorSetLayoutBinding>& bindings) {
         for (const auto& binding : bindings) {
             add_binding(binding);
         }
@@ -188,6 +189,13 @@ class DescriptorSetLayoutBuilder {
         return make_shared<DescriptorSetLayout>(context, std::move(sorted_bindings), flags);
     }
 
+    DescriptorSetLayoutHandle
+    build_push_descriptor_layout(const ContextHandle& context,
+                                 const vk::DescriptorSetLayoutCreateFlags flags = {}) {
+        return build_layout(context,
+                            flags | vk::DescriptorSetLayoutCreateFlagBits::ePushDescriptorKHR);
+    }
+
     // --------------------------------------------------------------------------------------------------------------------
 
   private:
diff --git a/include/merian/vk/memory/resource_allocations.hpp b/include/merian/vk/memory/resource_allocations.hpp
index f2011cf..348671e 100644
--- a/include/merian/vk/memory/resource_allocations.hpp
+++ b/include/merian/vk/memory/resource_allocations.hpp
@@ -309,7 +309,7 @@ class Texture : public std::enable_shared_from_this<Texture> {
         return image->get_memory();
     }
 
-    const SamplerHandle get_sampler() const {
+    const SamplerHandle& get_sampler() const {
         return sampler;
     }
 
@@ -322,7 +322,7 @@ class Texture : public std::enable_shared_from_this<Texture> {
         return view;
     }
 
-    const vk::DescriptorImageInfo get_descriptor_info() {
+    vk::DescriptorImageInfo get_descriptor_info() const {
         return vk::DescriptorImageInfo{*get_sampler(), view, image->get_current_layout()};
     }
 
diff --git a/include/merian/vk/pipeline/pipeline.hpp b/include/merian/vk/pipeline/pipeline.hpp
index ebe4c8d..5e3fa6b 100644
--- a/include/merian/vk/pipeline/pipeline.hpp
+++ b/include/merian/vk/pipeline/pipeline.hpp
@@ -2,6 +2,7 @@
 
 #include "merian/vk/context.hpp"
 #include "merian/vk/descriptors/descriptor_set.hpp"
+#include "merian/vk/memory/resource_allocations.hpp"
 #include "merian/vk/pipeline/pipeline_layout.hpp"
 
 #include <spdlog/spdlog.h>
@@ -47,6 +48,65 @@ class Pipeline : public std::enable_shared_from_this<Pipeline> {
         cmd.pushDescriptorSetKHR(get_pipeline_bind_point(), *pipeline_layout, set, writes);
     }
 
+    // shortcut to push the descriptor set 0.
+    void push_descriptor_set(const vk::CommandBuffer& cmd,
+                             const std::vector<vk::WriteDescriptorSet>& writes) {
+        push_descriptor_set(cmd, 0, writes);
+    }
+
+  private:
+    vk::WriteDescriptorSet make_descriptor_write(const vk::DescriptorBufferInfo& buffer_info,
+                                                 const uint32_t set,
+                                                 const uint32_t binding) {
+
+        return vk::WriteDescriptorSet{
+            {},
+            binding,
+            0,
+            1,
+            pipeline_layout->get_descriptor_set_layout(set)->get_type_for_binding(binding),
+            nullptr,
+            &buffer_info,
+        };
+    }
+
+    vk::WriteDescriptorSet make_descriptor_write(const vk::DescriptorImageInfo& image_info,
+                                                 const uint32_t set,
+                                                 const uint32_t binding) {
+        return vk::WriteDescriptorSet{
+            {},
+            binding,
+            0,
+            1,
+            pipeline_layout->get_descriptor_set_layout(set)->get_type_for_binding(binding),
+            &image_info,
+            nullptr,
+        };
+    }
+
+    template <typename... T, std::size_t... Is>
+    void push_descriptor_set(const vk::CommandBuffer& cmd,
+                             const uint32_t set,
+                             const std::index_sequence<Is...> /*unused*/,
+                             const T&... resources) {
+        const std::vector<vk::WriteDescriptorSet> writes = {
+            make_descriptor_write(resources, set, Is)...};
+        push_descriptor_set(cmd, set, writes);
+    }
+
+  public:
+    template <typename... T>
+    void
+    push_descriptor_set(const vk::CommandBuffer& cmd, const uint32_t set, const T&... resources) {
+        push_descriptor_set(cmd, set, std::index_sequence_for<T...>{}, resources...);
+    }
+
+    template <typename... T>
+    std::enable_if_t<(std::negation_v<std::is_same<int32_t, T>> && ...)>
+    push_descriptor_set(const vk::CommandBuffer& cmd, const T&... resources) {
+        push_descriptor_set(cmd, 0, std::index_sequence_for<T...>{}, resources...);
+    }
+
     // ---------------------------------------------------------------------------
 
     virtual vk::PipelineBindPoint get_pipeline_bind_point() const = 0;
diff --git a/include/merian/vk/pipeline/pipeline_layout.hpp b/include/merian/vk/pipeline/pipeline_layout.hpp
index 1880342..e211f17 100644
--- a/include/merian/vk/pipeline/pipeline_layout.hpp
+++ b/include/merian/vk/pipeline/pipeline_layout.hpp
@@ -51,6 +51,11 @@ class PipelineLayout : public std::enable_shared_from_this<PipelineLayout> {
         return ranges[id];
     }
 
+    const std::shared_ptr<DescriptorSetLayout>& get_descriptor_set_layout(const uint32_t set = 0) {
+        assert(set < shared_descriptor_set_layouts.size());
+        return shared_descriptor_set_layouts[set];
+    }
+
   private:
     const ContextHandle context;
     const std::vector<vk::PushConstantRange> ranges;
diff --git a/include/merian/vk/pipeline/specialization_info_builder.hpp b/include/merian/vk/pipeline/specialization_info_builder.hpp
index 157c9ec..44bd9b6 100644
--- a/include/merian/vk/pipeline/specialization_info_builder.hpp
+++ b/include/merian/vk/pipeline/specialization_info_builder.hpp
@@ -106,7 +106,8 @@ class SpecializationInfoBuilder {
 
     SpecializationInfoHandle build() const {
         std::vector<vk::SpecializationMapEntry> vec_entries;
-        for (auto& map_entry : entries) {
+        vec_entries.reserve(entries.size());
+        for (const auto& map_entry : entries) {
             vec_entries.push_back(map_entry.second);
         }
         SpecializationInfoHandle result =