Edited push_constants command to take stage instead of offset

Passing `offset` is mostly used to select which stage we want to push
the data to, and thus, its safer for the user to pass which stage they
want, and we can take care of computing the offset and sizes.

Also, passing the `stage` allows us to look through the
`overlapping_push_constant_ranges` from the `pipeline_layout` to check
which ranges we need to push to.

examples/src/bin/deferred/frame/ambient_lighting_system.rs

@@ -23,6 +23,7 @@ use vulkano::pipeline::graphics::vertex_input::BuffersDefinition;
 use vulkano::pipeline::graphics::viewport::{Viewport, ViewportState};
 use vulkano::pipeline::{GraphicsPipeline, Pipeline, PipelineBindPoint};
 use vulkano::render_pass::Subpass;
+use vulkano::shader::ShaderStage;
 
 /// Allows applying an ambient lighting to a scene.
 pub struct AmbientLightingSystem {
@@ -146,7 +147,11 @@ impl AmbientLightingSystem {
                 0,
                 descriptor_set,
             )
-            .push_constants(self.pipeline.layout().clone(), 0, push_constants)
+            .push_constants(
+                self.pipeline.layout().clone(),
+                ShaderStage::Fragment,
+                push_constants,
+            )
             .bind_vertex_buffers(0, self.vertex_buffer.clone())
             .draw(self.vertex_buffer.len() as u32, 1, 0, 0)
             .unwrap();

examples/src/bin/deferred/frame/directional_lighting_system.rs

@@ -24,6 +24,7 @@ use vulkano::pipeline::graphics::vertex_input::BuffersDefinition;
 use vulkano::pipeline::graphics::viewport::{Viewport, ViewportState};
 use vulkano::pipeline::{GraphicsPipeline, Pipeline, PipelineBindPoint};
 use vulkano::render_pass::Subpass;
+use vulkano::shader::ShaderStage;
 
 /// Allows applying a directional light source to a scene.
 pub struct DirectionalLightingSystem {
@@ -160,7 +161,11 @@ impl DirectionalLightingSystem {
                 0,
                 descriptor_set,
             )
-            .push_constants(self.pipeline.layout().clone(), 0, push_constants)
+            .push_constants(
+                self.pipeline.layout().clone(),
+                ShaderStage::Fragment,
+                push_constants,
+            )
             .bind_vertex_buffers(0, self.vertex_buffer.clone())
             .draw(self.vertex_buffer.len() as u32, 1, 0, 0)
             .unwrap();

examples/src/bin/deferred/frame/point_lighting_system.rs

@@ -24,6 +24,7 @@ use vulkano::pipeline::graphics::vertex_input::BuffersDefinition;
 use vulkano::pipeline::graphics::viewport::{Viewport, ViewportState};
 use vulkano::pipeline::{GraphicsPipeline, Pipeline, PipelineBindPoint};
 use vulkano::render_pass::Subpass;
+use vulkano::shader::ShaderStage;
 
 pub struct PointLightingSystem {
     gfx_queue: Arc<Queue>,
@@ -172,7 +173,11 @@ impl PointLightingSystem {
                 0,
                 descriptor_set,
             )
-            .push_constants(self.pipeline.layout().clone(), 0, push_constants)
+            .push_constants(
+                self.pipeline.layout().clone(),
+                ShaderStage::Fragment,
+                push_constants,
+            )
             .bind_vertex_buffers(0, self.vertex_buffer.clone())
             .draw(self.vertex_buffer.len() as u32, 1, 0, 0)
             .unwrap();

examples/src/bin/interactive_fractal/fractal_compute_pipeline.rs

@@ -18,6 +18,7 @@ use vulkano::descriptor_set::{PersistentDescriptorSet, WriteDescriptorSet};
 use vulkano::device::Queue;
 use vulkano::image::ImageAccess;
 use vulkano::pipeline::{ComputePipeline, Pipeline, PipelineBindPoint};
+use vulkano::shader::ShaderStage;
 use vulkano::sync::GpuFuture;
 
 pub struct FractalComputePipeline {
@@ -126,7 +127,11 @@ impl FractalComputePipeline {
         builder
             .bind_pipeline_compute(self.pipeline.clone())
             .bind_descriptor_sets(PipelineBindPoint::Compute, pipeline_layout.clone(), 0, set)
-            .push_constants(pipeline_layout.clone(), 0, push_constants)
+            .push_constants(
+                pipeline_layout.clone(),
+                ShaderStage::Fragment,
+                push_constants,
+            )
             .dispatch([img_dims[0] / 8, img_dims[1] / 8, 1])
             .unwrap();
         let command_buffer = builder.build().unwrap();

examples/src/bin/multi_window_game_of_life/game_of_life.rs

@@ -20,6 +20,7 @@ use vulkano::device::Queue;
 use vulkano::format::Format;
 use vulkano::image::ImageAccess;
 use vulkano::pipeline::{ComputePipeline, Pipeline, PipelineBindPoint};
+use vulkano::shader::ShaderStage;
 use vulkano::sync::GpuFuture;
 
 /// Pipeline holding double buffered grid & color image.
@@ -153,7 +154,11 @@ impl GameOfLifeComputePipeline {
         builder
             .bind_pipeline_compute(self.compute_life_pipeline.clone())
             .bind_descriptor_sets(PipelineBindPoint::Compute, pipeline_layout.clone(), 0, set)
-            .push_constants(pipeline_layout.clone(), 0, push_constants)
+            .push_constants(
+                pipeline_layout.clone(),
+                ShaderStage::Compute,
+                push_constants,
+            )
             .dispatch([img_dims[0] / 8, img_dims[1] / 8, 1])
             .unwrap();
     }

examples/src/bin/push-constants.rs

@@ -16,6 +16,7 @@ use vulkano::device::physical::{PhysicalDevice, PhysicalDeviceType};
 use vulkano::device::{Device, DeviceExtensions, Features};
 use vulkano::instance::{Instance, InstanceExtensions};
 use vulkano::pipeline::{ComputePipeline, Pipeline, PipelineBindPoint};
+use vulkano::shader::ShaderStage;
 use vulkano::sync;
 use vulkano::sync::GpuFuture;
 use vulkano::Version;
@@ -138,7 +139,11 @@ fn main() {
             0,
             set.clone(),
         )
-        .push_constants(pipeline.layout().clone(), 0, push_constants)
+        .push_constants(
+            pipeline.layout().clone(),
+            ShaderStage::Compute,
+            push_constants,
+        )
         .dispatch([1024, 1, 1])
         .unwrap();
     let command_buffer = builder.build().unwrap();

examples/src/bin/shader-types-sharing.rs

@@ -35,6 +35,7 @@ use vulkano::device::physical::{PhysicalDevice, PhysicalDeviceType};
 use vulkano::device::{Device, DeviceExtensions, Features, Queue};
 use vulkano::instance::{Instance, InstanceExtensions};
 use vulkano::pipeline::{ComputePipeline, Pipeline, PipelineBindPoint};
+use vulkano::shader::ShaderStage;
 use vulkano::sync;
 use vulkano::sync::GpuFuture;
 use vulkano::Version;
@@ -192,7 +193,7 @@ fn main() {
                 0,
                 set.clone(),
             )
-            .push_constants(pipeline.layout().clone(), 0, parameters)
+            .push_constants(pipeline.layout().clone(), ShaderStage::Compute, parameters)
             .dispatch([1024, 1, 1])
             .unwrap();
         let command_buffer = builder.build().unwrap();

vulkano/src/command_buffer/auto.rs

@@ -79,7 +79,7 @@ use crate::render_pass::Framebuffer;
 use crate::render_pass::LoadOp;
 use crate::render_pass::Subpass;
 use crate::sampler::Filter;
-use crate::shader::ShaderStages;
+use crate::shader::ShaderStage;
 use crate::sync::AccessCheckError;
 use crate::sync::AccessFlags;
 use crate::sync::GpuFuture;
@@ -1818,71 +1818,68 @@ impl<L, P> AutoCommandBufferBuilder<L, P> {
     ///
     /// # Panics
     ///
-    /// - Panics if `offset` is not a multiple of 4.
+    /// - Panics if it could not find `stage` in `pipeline_layout`.
     /// - Panics if the size of `push_constants` is not a multiple of 4.
+    /// - Panics if the size of `push_constants` does not match the push_constant size of `stage`.
     /// - Panics if any of the bytes in `push_constants` do not fall within any of the pipeline
     ///   layout's push constant ranges.
     pub fn push_constants<Pc>(
         &mut self,
         pipeline_layout: Arc<PipelineLayout>,
-        offset: u32,
+        stage: ShaderStage,
         push_constants: Pc,
     ) -> &mut Self {
         let size = mem::size_of::<Pc>() as u32;
-
         if size == 0 {
             return self;
         }
 
-        assert!(offset % 4 == 0, "the offset must be a multiple of 4");
+        let stages = stage.into();
+        // find the offset from the pipeline layout
+        let (stage_offset, stage_size) = pipeline_layout
+            .push_constant_ranges()
+            .iter()
+            .find(|range| range.stages == stages)
+            .map(|range| (range.offset, range.size))
+            .unwrap_or_else(|| {
+                panic!(
+                    "No push constant range for stage {:?} in pipeline layout {:?}",
+                    stage, pipeline_layout
+                )
+            });
+
         assert!(
             size % 4 == 0,
             "the size of push_constants must be a multiple of 4"
         );
+        assert_eq!(
+            stage_size, size,
+            "the size of push_constants must match the size of the range"
+        );
 
-        // Figure out which shader stages in the pipeline layout overlap this byte range.
-        // Also check that none of the bytes being set are outside all push constant ranges.
-        let shader_stages = pipeline_layout
-            .push_constant_ranges()
-            .iter()
-            .filter(|range| range.offset < offset + size && offset < range.offset + range.size)
-            .try_fold(
-                (ShaderStages::none(), offset),
-                |(shader_stages, last_bound), range| {
-                    if range.offset > last_bound {
-                        Err(())
-                    } else {
-                        Ok((
-                            shader_stages.union(&range.stages),
-                            last_bound.max(range.offset + range.size),
-                        ))
-                    }
-                },
-            )
-            .and_then(|(shader_stages, last_bound)| {
-                if shader_stages == ShaderStages::none() || last_bound < offset + size {
-                    Err(())
-                } else {
-                    Ok(shader_stages)
-                }
-            })
-            .expect(
-                "not all bytes in push_constants fall within the pipeline layout's push constant ranges",
-            );
-
-        unsafe {
-            let data = slice::from_raw_parts(
-                (&push_constants as *const Pc as *const u8).offset(offset as isize),
-                size as usize,
-            );
+        // SAFETY: `&push_constants` is a valid pointer, and the size of the struct is `size`,
+        //         thus, getting a slice of the whole struct is safe if its not modified.
+        let whole_data = unsafe {
+            slice::from_raw_parts(&push_constants as *const Pc as *const u8, size as usize)
+        };
 
-            self.inner.push_constants::<[u8]>(
-                pipeline_layout.clone(),
-                shader_stages,
-                offset,
-                size,
-                data,
-            );
+        // get all overlapping ranges of the push constants, and submit
+        // `push_constants` command for each of them
+        for range in pipeline_layout
+            .overlapping_push_constant_ranges()
+            .iter()
+            .filter(|range| range.stages.intersects(&stages))
+        {
+            let relative_offset = (range.offset - stage_offset) as usize;
+            unsafe {
+                self.inner.push_constants::<[u8]>(
+                    pipeline_layout.clone(),
+                    range.stages,
+                    range.offset,
+                    range.size,
+                    &whole_data[relative_offset..(relative_offset + range.size as usize)],
+                );
+            }
         }
 
         self