Reflect specialization constants in vulkano-shaders (#774)

vulkano-shaders/examples/example.rs

@@ -14,6 +14,8 @@ fn main() {
     let shader = r#"
 #version 450
 
+layout(constant_id = 5) const int index = 2;
+
 struct S {
     vec3 val1;
     bool val2[5];
@@ -25,11 +27,11 @@ layout(set = 0, binding = 1) uniform Block {
     S u_data;
 } block;
 
-in vec2 v_texcoords;
-out vec4 f_color;
+layout(location = 0) in vec2 v_texcoords;
+layout(location = 0) out vec4 f_color;
 
 void main() {
-    if (block.u_data.val2[3]) {
+    if (block.u_data.val2[index]) {
         f_color = texture(u_texture, v_texcoords);
     } else {
         f_color = vec4(1.0);

vulkano-shaders/src/entry_point.rs

@@ -53,10 +53,16 @@ pub fn write_entry_point(doc: &parse::Spirv, instruction: &parse::Instruction) -
                                         true,
                                     _ => false,
                                 });
+    
+    let spec_consts_struct = if ::spec_consts::has_specialization_constants(doc) {
+        "SpecializationConstants"
+    } else {
+        "()"
+    };
 
     let (ty, f_call) = {
         if let enums::ExecutionModel::ExecutionModelGLCompute = *execution {
-            (format!("::vulkano::pipeline::shader::ComputeEntryPoint<(), Layout>"),
+            (format!("::vulkano::pipeline::shader::ComputeEntryPoint<{}, Layout>", spec_consts_struct),
             format!("compute_entry_point(::std::ffi::CStr::from_ptr(NAME.as_ptr() as \
                     *const _), Layout(ShaderStages {{ compute: true, .. ShaderStages::none() \
                     }}))"))
@@ -131,8 +137,9 @@ pub fn write_entry_point(doc: &parse::Spirv, instruction: &parse::Instruction) -
                 enums::ExecutionModel::ExecutionModelKernel => unreachable!(),
             };
 
-            let t = format!("::vulkano::pipeline::shader::GraphicsEntryPoint<(), {0}Input, \
-                                {0}Output, Layout>",
+            let t = format!("::vulkano::pipeline::shader::GraphicsEntryPoint<{0}, {1}Input, \
+                                {1}Output, Layout>",
+                            spec_consts_struct,
                             capitalized_ep_name);
             let f = format!("graphics_entry_point(::std::ffi::CStr::from_ptr(NAME.as_ptr() \
                                 as *const _), {0}Input, {0}Output, Layout({2}), {1})",

vulkano-shaders/src/lib.rs

@@ -24,6 +24,7 @@ mod descriptor_sets;
 mod entry_point;
 mod enums;
 mod parse;
+mod spec_consts;
 mod structs;
 
 pub fn build_glsl_shaders<'a, I>(shaders: I)
@@ -108,6 +109,10 @@ pub fn reflect<R>(name: &str, mut spirv: R) -> Result<String, Error>
         use vulkano::descriptor::pipeline_layout::PipelineLayoutDesc;
         #[allow(unused_imports)]
         use vulkano::descriptor::pipeline_layout::PipelineLayoutDescPcRange;
+        #[allow(unused_imports)]
+        use vulkano::pipeline::shader::SpecializationConstants as SpecConstsTrait;
+        #[allow(unused_imports)]
+        use vulkano::pipeline::shader::SpecializationMapEntry;
     "#,
     );
 
@@ -202,6 +207,9 @@ impl {name} {{
 
         // descriptor sets
         output.push_str(&descriptor_sets::write_descriptor_sets(&doc));
+
+        // specialization constants
+        output.push_str(&spec_consts::write_specialization_constants(&doc));
     }
 
     Ok(output)

vulkano-shaders/src/parse.rs

@@ -166,6 +166,24 @@ pub enum Instruction {
         result_id: u32,
         data: Vec<u32>,
     },
+    SpecConstantTrue {
+        result_type_id: u32,
+        result_id: u32,
+    },
+    SpecConstantFalse {
+        result_type_id: u32,
+        result_id: u32,
+    },
+    SpecConstant {
+        result_type_id: u32,
+        result_id: u32,
+        data: Vec<u32>,
+    },
+    SpecConstantComposite {
+        result_type_id: u32,
+        result_id: u32,
+        data: Vec<u32>,
+    },
     FunctionEnd,
     Variable {
         result_type_id: u32,
@@ -318,6 +336,24 @@ fn decode_instruction(opcode: u16, operands: &[u32]) -> Result<Instruction, Pars
                result_id: operands[1],
                data: operands[2 ..].to_owned(),
            },
+           48 => Instruction::SpecConstantTrue {
+               result_type_id: operands[0],
+               result_id: operands[1],
+           },
+           49 => Instruction::SpecConstantFalse {
+               result_type_id: operands[0],
+               result_id: operands[1],
+           },
+           50 => Instruction::SpecConstant {
+               result_type_id: operands[0],
+               result_id: operands[1],
+               data: operands[2 ..].to_owned(),
+           },
+           51 => Instruction::SpecConstantComposite {
+               result_type_id: operands[0],
+               result_id: operands[1],
+               data: operands[2 ..].to_owned(),
+           },
            56 => Instruction::FunctionEnd,
            59 => Instruction::Variable {
                result_type_id: operands[0],

vulkano-shaders/src/spec_consts.rs

@@ -0,0 +1,152 @@
+// Copyright (c) 2017 The vulkano developers
+// Licensed under the Apache License, Version 2.0
+// <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
+// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
+// at your option. All files in the project carrying such
+// notice may not be copied, modified, or distributed except
+// according to those terms.
+
+use std::mem;
+
+use enums;
+use parse;
+
+/// Returns true if the document has specialization constants.
+pub fn has_specialization_constants(doc: &parse::Spirv) -> bool {
+    for instruction in doc.instructions.iter() {
+        match instruction {
+            &parse::Instruction::SpecConstantTrue { .. } => return true,
+            &parse::Instruction::SpecConstantFalse { .. } => return true,
+            &parse::Instruction::SpecConstant { .. } => return true,
+            &parse::Instruction::SpecConstantComposite { .. } => return true,
+            _ => (),
+        }
+    }
+
+    false
+}
+
+/// Writes the `SpecializationConstants` struct that contains the specialization constants and
+/// implements the `Default` and the `vulkano::pipeline::shader::SpecializationConstants` traits.
+pub fn write_specialization_constants(doc: &parse::Spirv) -> String {
+    struct SpecConst {
+        name: String,
+        constant_id: u32,
+        rust_ty: String,
+        rust_size: usize,
+        rust_alignment: usize,
+        default_value: String,
+    }
+
+    let mut spec_consts = Vec::new();
+
+    for instruction in doc.instructions.iter() {
+        let (type_id, result_id, default_value) = match instruction {
+            &parse::Instruction::SpecConstantTrue { result_type_id, result_id } => {
+                (result_type_id, result_id, "1u32".to_string())
+            },
+            &parse::Instruction::SpecConstantFalse { result_type_id, result_id } => {
+                (result_type_id, result_id, "0u32".to_string())
+            },
+            &parse::Instruction::SpecConstant { result_type_id, result_id, ref data } => {
+                let data = data.iter().map(|d| d.to_string() + "u32").collect::<Vec<_>>().join(", ");
+                let def_val = format!("unsafe {{ ::std::mem::transmute([{}]) }}", data);
+                (result_type_id, result_id, def_val)
+            },
+            &parse::Instruction::SpecConstantComposite { result_type_id, result_id, ref data } => {
+                let data = data.iter().map(|d| d.to_string() + "u32").collect::<Vec<_>>().join(", ");
+                let def_val = format!("unsafe {{ ::std::mem::transmute([{}]) }}", data);
+                (result_type_id, result_id, def_val)
+            },
+            _ => continue,
+        };
+
+        let (rust_ty, rust_size, rust_alignment) = spec_const_type_from_id(doc, type_id);
+        let rust_size = rust_size.expect("Found runtime-sized specialization constant");
+
+        let constant_id = doc.instructions.iter().filter_map(|i| {
+            match i {
+                &parse::Instruction::Decorate
+                    { target_id, decoration: enums::Decoration::DecorationSpecId, ref params }
+                    if target_id == result_id =>
+                {
+                    Some(params[0])
+                },
+                _ => None,
+            }
+        }).next().expect("Found a specialization constant with no SpecId decoration");
+
+        spec_consts.push(SpecConst {
+            name: ::name_from_id(doc, result_id),
+            constant_id,
+            rust_ty,
+            rust_size,
+            rust_alignment,
+            default_value,
+        });
+    }
+
+    let map_entries = {
+        let mut map_entries = Vec::new();
+        let mut curr_offset = 0;
+        for c in &spec_consts {
+            map_entries.push(format!("SpecializationMapEntry {{
+                constant_id: {},
+                offset: {},
+                size: {},
+            }}", c.constant_id, curr_offset, c.rust_size));
+
+            assert_ne!(c.rust_size, 0);
+            curr_offset += c.rust_size;
+            curr_offset = c.rust_alignment * (1 + (curr_offset - 1) / c.rust_alignment);
+        }
+        map_entries
+    };
+
+    format!(r#"
+
+#[derive(Debug, Copy, Clone)]
+#[repr(C)]
+pub struct SpecializationConstants {{
+    {struct_def}
+}}
+
+impl Default for SpecializationConstants {{
+    fn default() -> SpecializationConstants {{
+        SpecializationConstants {{
+            {def_vals}
+        }}
+    }}
+}}
+
+unsafe impl SpecConstsTrait for SpecializationConstants {{
+    fn descriptors() -> &'static [SpecializationMapEntry] {{
+        &[
+            {map_entries}
+        ]
+    }}
+}}
+
+    "#,
+        struct_def = spec_consts.iter().map(|c| format!("{}: {}", c.name, c.rust_ty))
+                                .collect::<Vec<_>>().join(", "),
+        def_vals = spec_consts.iter().map(|c| format!("{}: {}", c.name, c.default_value))
+                                .collect::<Vec<_>>().join(", "),
+        map_entries = map_entries.join(", ")
+    )
+}
+
+// Wrapper around `type_from_id` that also handles booleans.
+fn spec_const_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 => {
+                return ("u32".to_owned(), Some(mem::size_of::<u32>()), mem::align_of::<u32>());
+            },
+            _ => ()
+        }
+    }
+
+    ::structs::type_from_id(doc, searched)
+}

vulkano-shaders/src/structs.rs

@@ -252,13 +252,7 @@ pub fn type_from_id(doc: &parse::Spirv, searched: u32) -> (String, Option<usize>
     for instruction in doc.instructions.iter() {
         match instruction {
             &parse::Instruction::TypeBool { result_id } if result_id == searched => {
-                #[repr(C)]
-                struct Foo {
-                    data: bool,
-                    after: u8,
-                }
-                let size = unsafe { (&(&*(0 as *const Foo)).after) as *const u8 as usize };
-                return ("bool".to_owned(), Some(size), mem::align_of::<Foo>());
+                panic!("Can't put booleans in structs")
             },
             &parse::Instruction::TypeInt {
                 result_id,