layers: Fix invalidating descriptor mode

Author: ziga-lunarg

layers/state_tracker/cmd_buffer_state.cpp

@@ -276,6 +276,19 @@ void CommandBuffer::SetDescriptorMode(vvl::DescriptorMode new_mode, vvl::Func fu
     }
 }
 
+// Some functions like vkCmdPushConstant are valid in Classic/Buffer, but invalid in Heap
+// So calling it doesn't "set" a mode, but instead only "invalidates"
+void CommandBuffer::InvalidateDescriptorMode(vvl::DescriptorMode invalidate_mode, vvl::DescriptorMode new_mode,
+                                             vvl::Func function) {
+    for (uint32_t i = 0; i < vvl::BindPointCount; i++) {
+        const vvl::DescriptorMode current_mode = lastBound[i].GetDescriptorMode();
+        if (current_mode == invalidate_mode) {
+            SetDescriptorMode(new_mode, function);
+            break;
+        }
+    }
+}
+
 // Reset the command buffer state
 // Maintain the createInfo and set state to CB_NEW, but clear all other state
 void CommandBuffer::ResetCBState() {

layers/state_tracker/cmd_buffer_state.h

@@ -593,6 +593,7 @@ class CommandBuffer : public RefcountedStateObject, public SubStateManager<Comma
     } descriptor_heap;
 
     void SetDescriptorMode(vvl::DescriptorMode new_mode, vvl::Func function);
+    void InvalidateDescriptorMode(vvl::DescriptorMode invalidate_mode, vvl::DescriptorMode new_mode, vvl::Func function);
 
     mutable std::shared_mutex lock;
     ReadLockGuard ReadLock() const { return ReadLockGuard(lock); }

layers/state_tracker/state_tracker.cpp

@@ -3301,11 +3301,12 @@ void DeviceState::PostCallRecordCmdPushConstants(VkCommandBuffer commandBuffer,
     auto pipeline_layout_state = Get<PipelineLayout>(layout);
     ASSERT_AND_RETURN(cb_state && pipeline_layout_state);
 
-    const DescriptorMode descriptor_mode =
-        pipeline_layout_state->has_descriptor_buffer ? vvl::DescriptorModeBuffer : vvl::DescriptorModeClassic;
-    cb_state->SetDescriptorMode(descriptor_mode, record_obj.location.function);
     cb_state->RecordCommand(record_obj.location);
     cb_state->RecordPushConstants(*pipeline_layout_state, stageFlags, offset, size, pValues);
+
+    const DescriptorMode descriptor_mode =
+        pipeline_layout_state->has_descriptor_buffer ? vvl::DescriptorModeBuffer : vvl::DescriptorModeClassic;
+    cb_state->InvalidateDescriptorMode(vvl::DescriptorModeHeap, descriptor_mode, record_obj.location.function);
 }
 
 void DeviceState::PostCallRecordCmdPushConstants2(VkCommandBuffer commandBuffer, const VkPushConstantsInfo* pPushConstantsInfo,