From f5587619979b791c8d37110393c4a9b4c50e0a34 Mon Sep 17 00:00:00 2001
From: Marijn Suijten <marijn@traverseresearch.nl>
Date: Sat, 29 Jul 2023 10:51:03 +0200
Subject: [PATCH] extensions/amd: Add VK_AMD_shader_info (#773)

---
 Changelog.md                          |  1 +
 ash/src/extensions/amd/mod.rs         |  2 +
 ash/src/extensions/amd/shader_info.rs | 86 +++++++++++++++++++++++++++
 3 files changed, 89 insertions(+)
 create mode 100644 ash/src/extensions/amd/shader_info.rs

diff --git a/Changelog.md b/Changelog.md
index d404c28..d8ac096 100644
--- a/Changelog.md
+++ b/Changelog.md
@@ -16,6 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added `VK_GOOGLE_display_timing` device extension (#765)
 - Added `VK_ANDROID_external_memory_android_hardware_buffer` device extension (#769)
 - Added `VK_AMD_buffer_marker` device extension (#772)
+- Added `VK_AMD_shader_info` device extension (#773)
 
 ### Changed
 
diff --git a/ash/src/extensions/amd/mod.rs b/ash/src/extensions/amd/mod.rs
index 8ba8f50..2399991 100644
--- a/ash/src/extensions/amd/mod.rs
+++ b/ash/src/extensions/amd/mod.rs
@@ -1,3 +1,5 @@
 pub use self::buffer_marker::BufferMarker;
+pub use self::shader_info::{ShaderInfo, ShaderInfoResult};
 
 mod buffer_marker;
+mod shader_info;
diff --git a/ash/src/extensions/amd/shader_info.rs b/ash/src/extensions/amd/shader_info.rs
new file mode 100644
index 0000000..9c01aec
--- /dev/null
+++ b/ash/src/extensions/amd/shader_info.rs
@@ -0,0 +1,86 @@
+use crate::prelude::*;
+use crate::vk;
+use crate::{Device, Instance};
+use std::ffi::CStr;
+use std::mem;
+
+/// <https://registry.khronos.org/vulkan/specs/1.3-extensions/man/html/VK_AMD_shader_info.html>
+#[derive(Clone)]
+pub struct ShaderInfo {
+    handle: vk::Device,
+    fp: vk::AmdShaderInfoFn,
+}
+
+impl ShaderInfo {
+    pub fn new(instance: &Instance, device: &Device) -> Self {
+        let handle = device.handle();
+        let fp = vk::AmdShaderInfoFn::load(|name| unsafe {
+            mem::transmute(instance.get_device_proc_addr(handle, name.as_ptr()))
+        });
+        Self { handle, fp }
+    }
+
+    /// <https://registry.khronos.org/vulkan/specs/1.3-extensions/man/html/vkGetShaderInfoAMD.html>
+    #[inline]
+    pub unsafe fn get_shader_info(
+        &self,
+        pipeline: vk::Pipeline,
+        shader_stage: vk::ShaderStageFlags,
+        info_type: vk::ShaderInfoTypeAMD,
+    ) -> VkResult<ShaderInfoResult> {
+        let load_data = |count: &mut usize, data: *mut u8| {
+            (self.fp.get_shader_info_amd)(
+                self.handle,
+                pipeline,
+                shader_stage,
+                info_type,
+                count,
+                data.cast(),
+            )
+        };
+
+        match info_type {
+            vk::ShaderInfoTypeAMD::STATISTICS => {
+                let mut statistics_info = mem::MaybeUninit::<vk::ShaderStatisticsInfoAMD>::uninit();
+                load_data(
+                    &mut std::mem::size_of_val(&statistics_info),
+                    statistics_info.as_mut_ptr().cast(),
+                )
+                .result()?;
+                Ok(ShaderInfoResult::StatisticsInfo(
+                    statistics_info.assume_init(),
+                ))
+            }
+            vk::ShaderInfoTypeAMD::BINARY => {
+                read_into_uninitialized_vector(load_data).map(ShaderInfoResult::Binary)
+            }
+            vk::ShaderInfoTypeAMD::DISASSEMBLY => {
+                read_into_uninitialized_vector(load_data).map(ShaderInfoResult::Disassembly)
+            }
+            #[cfg(feature = "debug")]
+            x => unimplemented!("ShaderInfoTypeAMD {:?}", x),
+            #[cfg(not(feature = "debug"))]
+            x => unimplemented!("ShaderInfoTypeAMD {}", x.0),
+        }
+    }
+
+    pub const NAME: &'static CStr = vk::AmdShaderInfoFn::NAME;
+
+    #[inline]
+    pub fn fp(&self) -> &vk::AmdShaderInfoFn {
+        &self.fp
+    }
+
+    #[inline]
+    pub fn device(&self) -> vk::Device {
+        self.handle
+    }
+}
+
+#[derive(Clone)]
+#[cfg_attr(feature = "debug", derive(Debug))]
+pub enum ShaderInfoResult {
+    StatisticsInfo(vk::ShaderStatisticsInfoAMD),
+    Binary(Vec<u8>),
+    Disassembly(Vec<u8>),
+}