From b9498952a13497f6825e1a6015787fd8a71b4e60 Mon Sep 17 00:00:00 2001
From: Connor Fitzgerald <connorwadefitzgerald@gmail.com>
Date: Fri, 24 Jul 2020 14:41:46 -0400
Subject: [PATCH] Fix push constant pipeline invalidation

---
 wgpu-core/src/command/bind.rs    | 24 +++++++++++++++++++++---
 wgpu-core/src/command/compute.rs | 10 ++++++----
 wgpu-core/src/command/render.rs  | 10 ++++++----
 3 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/wgpu-core/src/command/bind.rs b/wgpu-core/src/command/bind.rs
index c38f670459..e53bae67c2 100644
--- a/wgpu-core/src/command/bind.rs
+++ b/wgpu-core/src/command/bind.rs
@@ -3,9 +3,9 @@
  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
 
 use crate::{
-    binding_model::BindGroup,
+    binding_model::{BindGroup, PipelineLayout},
     device::SHADER_STAGE_COUNT,
-    hub::GfxBackend,
+    hub::{GfxBackend, Storage},
     id::{BindGroupId, BindGroupLayoutId, PipelineLayoutId},
     Stored, MAX_BIND_GROUPS,
 };
@@ -151,7 +151,25 @@ impl Binder {
         self.entries.clear();
     }
 
-    pub(crate) fn reset_expectations(&mut self, length: usize) {
+    pub(crate) fn change_pipeline_layout<B: GfxBackend>(
+        &mut self,
+        guard: &Storage<PipelineLayout<B>, PipelineLayoutId>,
+        new_id: PipelineLayoutId,
+    ) {
+        let old_id_opt = self.pipeline_layout_id.replace(new_id);
+        let new = &guard[new_id];
+
+        let length = if let Some(old_id) = old_id_opt {
+            let old = &guard[old_id];
+            if old.push_constant_ranges == new.push_constant_ranges {
+                new.bind_group_layout_ids.len()
+            } else {
+                0
+            }
+        } else {
+            0
+        };
+
         for entry in self.entries[length..].iter_mut() {
             entry.expected_layout_id = None;
         }
diff --git a/wgpu-core/src/command/compute.rs b/wgpu-core/src/command/compute.rs
index 9ae644dbf0..9c56c70686 100644
--- a/wgpu-core/src/command/compute.rs
+++ b/wgpu-core/src/command/compute.rs
@@ -251,10 +251,12 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
                     // Rebind resources
                     if state.binder.pipeline_layout_id != Some(pipeline.layout_id.value) {
                         let pipeline_layout = &pipeline_layout_guard[pipeline.layout_id.value];
-                        state.binder.pipeline_layout_id = Some(pipeline.layout_id.value);
-                        state
-                            .binder
-                            .reset_expectations(pipeline_layout.bind_group_layout_ids.len());
+
+                        state.binder.change_pipeline_layout(
+                            &*pipeline_layout_guard,
+                            pipeline.layout_id.value,
+                        );
+
                         let mut is_compatible = true;
 
                         for (index, (entry, &bgl_id)) in state
diff --git a/wgpu-core/src/command/render.rs b/wgpu-core/src/command/render.rs
index b63b6002fc..e184a06247 100644
--- a/wgpu-core/src/command/render.rs
+++ b/wgpu-core/src/command/render.rs
@@ -1103,10 +1103,12 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
                     // Rebind resource
                     if state.binder.pipeline_layout_id != Some(pipeline.layout_id.value) {
                         let pipeline_layout = &pipeline_layout_guard[pipeline.layout_id.value];
-                        state.binder.pipeline_layout_id = Some(pipeline.layout_id.value);
-                        state
-                            .binder
-                            .reset_expectations(pipeline_layout.bind_group_layout_ids.len());
+
+                        state.binder.change_pipeline_layout(
+                            &*pipeline_layout_guard,
+                            pipeline.layout_id.value,
+                        );
+
                         let mut is_compatible = true;
 
                         for (index, (entry, &bgl_id)) in state