From dd4ac2219af12a02eefc8327210427f78aaa9630 Mon Sep 17 00:00:00 2001
From: Pierre Krieger <pierre.krieger1708@gmail.com>
Date: Thu, 23 Feb 2017 07:40:37 +0100
Subject: [PATCH] Basic push constants reporting in vulkano_shaders

---
 vulkano-shaders/src/descriptor_sets.rs | 42 ++++++++++++++++++++++++--
 vulkano-shaders/src/structs.rs         |  3 +-
 2 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/vulkano-shaders/src/descriptor_sets.rs b/vulkano-shaders/src/descriptor_sets.rs
index 1f9e365de..e88d251b6 100644
--- a/vulkano-shaders/src/descriptor_sets.rs
+++ b/vulkano-shaders/src/descriptor_sets.rs
@@ -61,6 +61,21 @@ pub fn write_descriptor_sets(doc: &parse::Spirv) -> String {
         });
     }
 
+    // Looping to find all the push constant structs.
+    let mut push_constants_size = 0;
+    for instruction in doc.instructions.iter() {
+        let type_id = match instruction {
+            &parse::Instruction::TypePointer { type_id, storage_class: enums::StorageClass::StorageClassPushConstant, .. } => {
+                type_id
+            },
+            _ => continue
+        };
+
+        let (_, size, _) = ::structs::type_from_id(doc, type_id);
+        let size = size.expect("Found runtime-sized push constants");
+        push_constants_size = cmp::max(push_constants_size, size);
+    }
+
     // Writing the body of the `descriptor` method.
     let descriptor_body = descriptors.iter().map(|d| {
         format!("({set}, {binding}) => Some(DescriptorDesc {{
@@ -90,6 +105,25 @@ pub fn write_descriptor_sets(doc: &parse::Spirv) -> String {
                 name = d.name, set = d.set, binding = d.binding)
     }).collect::<Vec<_>>().concat();
 
+    // Writing the body of the `num_push_constants_ranges` method.
+    let num_push_constants_ranges_body = {
+        if push_constants_size == 0 {
+            "0"
+        } else {
+            "1"
+        }
+    };
+
+    // Writing the body of the `push_constants_range` method.
+    let push_constants_range_body = format!(r#"
+        if num != 0 || {pc_size} == 0 {{ return None; }}
+        Some(PipelineLayoutDescPcRange {{
+            offset: 0,
+            size: {pc_size},
+            stages: ShaderStages::all(),     // FIXME: wrong
+        }})
+    "#, pc_size = push_constants_size);
+
     format!(r#"
         #[derive(Debug, Clone)]
         pub struct Layout(ShaderStages);
@@ -115,11 +149,11 @@ pub fn write_descriptor_sets(doc: &parse::Spirv) -> String {
             }}
 
             fn num_push_constants_ranges(&self) -> usize {{
-                0       // FIXME:
+                {num_push_constants_ranges_body}
             }}
 
             fn push_constants_range(&self, num: usize) -> Option<PipelineLayoutDescPcRange> {{
-                None
+                {push_constants_range_body}
             }}
         }}
 
@@ -133,7 +167,9 @@ pub fn write_descriptor_sets(doc: &parse::Spirv) -> String {
             }}
         }}
         "#, num_sets = num_sets, num_bindings_in_set_body = num_bindings_in_set_body,
-            descriptor_by_name_body = descriptor_by_name_body, descriptor_body = descriptor_body)
+            descriptor_by_name_body = descriptor_by_name_body, descriptor_body = descriptor_body,
+            num_push_constants_ranges_body = num_push_constants_ranges_body,
+            push_constants_range_body = push_constants_range_body)
 }
 
 /// Assumes that `variable` is a variable with a `TypePointer` and returns the id of the pointed
diff --git a/vulkano-shaders/src/structs.rs b/vulkano-shaders/src/structs.rs
index c123524f4..588f4e989 100644
--- a/vulkano-shaders/src/structs.rs
+++ b/vulkano-shaders/src/structs.rs
@@ -210,7 +210,7 @@ fn is_builtin_member(doc: &parse::Spirv, id: u32, member_id: u32) -> bool {
 /// Returns the type name to put in the Rust struct, and its size and alignment.
 ///
 /// The size can be `None` if it's only known at runtime.
-fn type_from_id(doc: &parse::Spirv, searched: u32) -> (String, Option<usize>, usize) {
+pub fn type_from_id(doc: &parse::Spirv, searched: u32) -> (String, Option<usize>, usize) {
     for instruction in doc.instructions.iter() {
         match instruction {
             &parse::Instruction::TypeBool { result_id } if result_id == searched => {
@@ -304,6 +304,7 @@ fn type_from_id(doc: &parse::Spirv, searched: u32) -> (String, Option<usize>, us
                 return (format!("[{}]", t), None, t_align);
             },
             &parse::Instruction::TypeStruct { result_id, ref member_types } if result_id == searched => {
+                // TODO: take the Offset member decorate into account?
                 let name = ::name_from_id(doc, result_id);
                 let size = member_types.iter().filter_map(|&t| type_from_id(doc, t).1).fold(0, |a,b|a+b);
                 let align = member_types.iter().map(|&t| type_from_id(doc, t).2).max().unwrap_or(1);